popertots/ragmcp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
e1f7cffdef493cbb82e85155d7fc4b21ec701e73
ragmcp/mcp_codebase.py
T
popertots e1f7cffdef path
2025-11-15 01:54:46 +00:00

3713 lines
144 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
import os
import signal
import json
from collections import defaultdict, deque
import re
import subprocess
import logging
from pathlib import Path
from typing import List, Dict, Tuple, Any, Optional
from threading import RLock
from contextlib import contextmanager
from functools import lru_cache
import time
import hashlib
import pathspec
import requests
import numpy as np
from mcp.server.fastmcp import FastMCP
from langchain_ollama import OllamaEmbeddings
from langchain_community.vectorstores import Chroma
from rank_bm25 import BM25Okapi
import javalang
import sqlparse
import sqlglot
from enhanced_toon import EnhancedToon
from graph.graph import LocalGraph
os.environ["ANONYMIZED_TELEMETRY"] = "False"
# -----------------------------
# Configuration
# -----------------------------
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(levelname)s - %(message)s",
handlers=[
logging.FileHandler("mcp_codebase.log"), # Log to file
logging.StreamHandler() # Also log to console
]
)
logger = logging.getLogger("rag-mcp")
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("./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"
# Models / behavior
EMBEDDING_MODEL = "bge-m3" # Ollama local embedding model
RERANKER_MODEL_OLLAMA = "dengcao/Qwen3-Reranker-8B:Q4_K_M" # prompt-based reranker via Ollama
ENABLE_RERANK = True # toggle reranking
RERANK_TOP_N = 10
# Hybrid weights (vector first: 80%, BM25: 20%)
VECTOR_WEIGHT = 0.8
BM25_WEIGHT = 0.2
# Embedding batch configuration
EMBEDDING_BATCH_SIZE = 200
# Contextual weighting configuration
CONTEXTUAL_WEIGHTS = {
# SQL weights
"sql_schema": 1.3, # CREATE TABLE, ALTER TABLE, schema definitions
"sql_function": 1.2, # Function/procedure definitions
"sql_view": 1.15, # VIEW definitions
"sql_index": 1.1, # INDEX definitions
# Python weights
"python_class": 1.25, # Class definitions (enhanced)
"python_function": 1.15, # Function definitions
"python_method": 1.2, # Method definitions (enhanced)
"python_async": 1.1, # Async functions
# Go weights
"go_type": 1.25, # Type definitions (struct, interface)
"go_function": 1.15, # Function definitions
"go_method": 1.2, # Method definitions (enhanced)
"go_struct": 1.3, # Struct definitions (enhanced)
# Java weights
"java_class": 1.3, # Class definitions (enhanced)
"java_interface": 1.25, # Interface definitions (enhanced)
"java_method": 1.2, # Method definitions (enhanced)
"java_constructor": 1.15, # Constructor definitions (enhanced)
"java_enum": 1.2, # Enum definitions (enhanced)
"java_field": 1.1, # Field definitions (enhanced)
# Rust weights (new)
"rust_struct": 1.3, # Struct definitions
"rust_enum": 1.25, # Enum definitions
"rust_trait": 1.35, # Trait definitions (very important in Rust)
"rust_impl": 1.2, # Implementation blocks
"rust_function": 1.15, # Function definitions
"rust_method": 1.2, # Method definitions
"rust_async": 1.1, # Async functions
"rust_unsafe": 1.15, # Unsafe blocks/functions
# TypeScript/JavaScript weights (for Svelte scripts)
"typescript_interface": 1.25, # Interface definitions
"typescript_class": 1.2, # Class definitions
"typescript_function": 1.15, # Function definitions
"typescript_type": 1.2, # Type definitions
# Svelte weights (new)
"svelte_component": 1.4, # Component definitions (very important)
"svelte_script": 1.1, # Script blocks
"svelte_style": 1.05, # Style blocks
"svelte_markup": 1.15, # Markup/template
"svelte_prop": 1.25, # Component props (enhanced)
"svelte_reactive": 1.2, # Reactive declarations
# Default weight for unspecified types
"default": 1.0
}
SUPPORTED_EXTENSIONS = {
".py", ".go", ".rs", ".ts", ".tsx", ".js", ".jsx",
".svelte", ".sql", ".pgsql", ".sh", ".bash", ".zsh",
".rb", ".java", ".cpp", ".c", ".h", ".hpp"
}
# MCP
mcp = FastMCP("codebase-rag")
# Global state
vectorstore: Optional[Chroma] = None
bm25: Optional[BM25Okapi] = None
bm25_corpus: List[str] = []
chunks_metadata: List[Dict[str, Any]] = []
embeddings: Optional[OllamaEmbeddings] = None
_startup_lock = RLock()
index_build_time: float = 0.0
# -----------------------------
# Thread-Safe Context Manager
# -----------------------------
@contextmanager
def safe_search():
"""Thread-safe context manager for search operations."""
with _startup_lock:
if vectorstore is None or bm25 is None:
raise RuntimeError("Indexes not ready. Please wait for initialization or rebuild.")
yield
# -----------------------------
# File Hashing for Cache Invalidation
# -----------------------------
def get_file_hash(filepath: Path) -> str:
"""Generate hash of file content + mtime for cache key."""
try:
stat = filepath.stat()
content_sample = filepath.read_bytes()[:1024] # First 1KB for speed
hash_input = f"{stat.st_mtime}:{stat.st_size}:{content_sample}".encode()
return hashlib.md5(hash_input).hexdigest()
except Exception:
return ""
# -----------------------------
# Contextual Weight Calculator
# -----------------------------
def calculate_contextual_weight(metadata: Dict[str, Any]) -> float:
"""
Calculate contextual weight multiplier based on chunk metadata.
Higher weights for schema-defining statements, function definitions, etc.
"""
language = metadata.get("language", "").lower()
chunk_type = metadata.get("type", "").lower()
block_type = metadata.get("block_type", "").lower()
# SQL-specific weights
if language == "sql":
if chunk_type in ("create", "alter", "create table", "alter table"):
return CONTEXTUAL_WEIGHTS.get("sql_schema", 1.0)
elif chunk_type in ("function", "procedure", "create function", "create procedure"):
return CONTEXTUAL_WEIGHTS.get("sql_function", 1.0)
elif chunk_type in ("view", "create view"):
return CONTEXTUAL_WEIGHTS.get("sql_view", 1.0)
elif chunk_type in ("index", "create index"):
return CONTEXTUAL_WEIGHTS.get("sql_index", 1.0)
# Python-specific weights (enhanced)
elif language == "python":
if chunk_type == "classdef":
return CONTEXTUAL_WEIGHTS.get("python_class", 1.0)
elif chunk_type in ("functiondef", "asyncfunctiondef"):
if metadata.get("is_method"):
return CONTEXTUAL_WEIGHTS.get("python_method", 1.0)
elif chunk_type == "asyncfunctiondef" or metadata.get("is_async"):
return CONTEXTUAL_WEIGHTS.get("python_async", 1.0)
else:
return CONTEXTUAL_WEIGHTS.get("python_function", 1.0)
# Go-specific weights (enhanced)
elif language == "go":
if chunk_type in ("struct", "interface"):
return CONTEXTUAL_WEIGHTS.get("go_struct", 1.0)
elif chunk_type == "type":
return CONTEXTUAL_WEIGHTS.get("go_type", 1.0)
elif chunk_type == "method":
return CONTEXTUAL_WEIGHTS.get("go_method", 1.0)
elif chunk_type == "func":
return CONTEXTUAL_WEIGHTS.get("go_function", 1.0)
# Java-specific weights (enhanced)
elif language == "java":
if chunk_type == "class":
return CONTEXTUAL_WEIGHTS.get("java_class", 1.0)
elif chunk_type == "interface":
return CONTEXTUAL_WEIGHTS.get("java_interface", 1.0)
elif chunk_type == "method":
return CONTEXTUAL_WEIGHTS.get("java_method", 1.0)
elif chunk_type == "constructor":
return CONTEXTUAL_WEIGHTS.get("java_constructor", 1.0)
elif chunk_type == "enum":
return CONTEXTUAL_WEIGHTS.get("java_enum", 1.0)
elif chunk_type == "field":
return CONTEXTUAL_WEIGHTS.get("java_field", 1.0)
# Rust-specific weights (new)
elif language == "rust":
if chunk_type == "struct":
return CONTEXTUAL_WEIGHTS.get("rust_struct", 1.0)
elif chunk_type == "enum":
return CONTEXTUAL_WEIGHTS.get("rust_enum", 1.0)
elif chunk_type == "trait":
return CONTEXTUAL_WEIGHTS.get("rust_trait", 1.0)
elif chunk_type == "impl":
return CONTEXTUAL_WEIGHTS.get("rust_impl", 1.0)
elif chunk_type == "function":
if metadata.get("is_async"):
return CONTEXTUAL_WEIGHTS.get("rust_async", 1.0)
elif metadata.get("is_unsafe"):
return CONTEXTUAL_WEIGHTS.get("rust_unsafe", 1.0)
elif metadata.get("is_method", False):
return CONTEXTUAL_WEIGHTS.get("rust_method", 1.0)
else:
return CONTEXTUAL_WEIGHTS.get("rust_function", 1.0)
# TypeScript/JavaScript weights (for Svelte scripts)
elif language == "typescript":
if chunk_type == "interface":
return CONTEXTUAL_WEIGHTS.get("typescript_interface", 1.0)
elif chunk_type == "class":
return CONTEXTUAL_WEIGHTS.get("typescript_class", 1.0)
elif chunk_type == "function":
return CONTEXTUAL_WEIGHTS.get("typescript_function", 1.0)
elif chunk_type == "type":
return CONTEXTUAL_WEIGHTS.get("typescript_type", 1.0)
# Svelte-specific weights (new)
elif language == "svelte":
if chunk_type == "component":
return CONTEXTUAL_WEIGHTS.get("svelte_component", 1.0)
elif chunk_type == "prop":
return CONTEXTUAL_WEIGHTS.get("svelte_prop", 1.0)
elif chunk_type == "reactive":
return CONTEXTUAL_WEIGHTS.get("svelte_reactive", 1.0)
elif block_type == "markup":
return CONTEXTUAL_WEIGHTS.get("svelte_markup", 1.0)
elif block_type == "script":
return CONTEXTUAL_WEIGHTS.get("svelte_script", 1.0)
elif block_type == "style":
return CONTEXTUAL_WEIGHTS.get("svelte_style", 1.0)
# CSS weights (for Svelte styles)
elif language == "css":
if block_type == "style":
return CONTEXTUAL_WEIGHTS.get("svelte_style", 1.0)
return CONTEXTUAL_WEIGHTS.get("default", 1.0)
def apply_contextual_weights_to_embeddings(embeddings_list: List[List[float]],
metadata_list: List[Dict[str, Any]]) -> List[List[float]]:
"""
Apply contextual weights to embedding vectors by scaling them.
This biases the vector space without changing the embedding model.
"""
weighted_embeddings = []
if len(embeddings_list) != len(metadata_list):
raise ValueError(
f"Embedding/metadata mismatch: {len(embeddings_list)} embeddings vs {len(metadata_list)} metadata. "
"This will break downstream systems (Chroma)."
)
for embedding, metadata in zip(embeddings_list, metadata_list):
weight = calculate_contextual_weight(metadata)
# Convert to numpy for easier manipulation
emb_array = np.array(embedding)
# Scale the embedding vector by the weight
# This effectively increases the magnitude, making it more "important"
weighted_emb = emb_array * weight
# Optionally normalize to maintain consistent vector magnitudes
# Comment out if you want the raw weighted vectors
norm = np.linalg.norm(weighted_emb)
if norm > 0:
weighted_emb = weighted_emb / norm * np.linalg.norm(emb_array)
weighted_embeddings.append(weighted_emb.tolist())
if weight != 1.0:
logger.debug(f"Applied weight {weight:.2f} to {metadata.get('file')}:{metadata.get('type')}")
return weighted_embeddings
# -----------------------------
# Utility: .ragignore
# -----------------------------
def load_ragignore(base_path: Path) -> pathspec.PathSpec:
if RAGIGNORE_PATH.exists():
with open(RAGIGNORE_PATH, "r", encoding="utf-8") as f:
return pathspec.PathSpec.from_lines("gitwildmatch", f)
default_patterns = [
".git/", "node_modules/", "__pycache__/", "*.pyc",
"*.so", "*.dll", "*.exe", "venv/", ".venv/", "dist/",
"build/", "*.log", "*.tmp", ".DS_Store"
]
return pathspec.PathSpec.from_lines("gitwildmatch", default_patterns)
# -----------------------------
# Python AST chunking with caching
# -----------------------------
@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."""
filepath = Path(filepath_str)
chunks = []
try:
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),
"name": getattr(node, "name", ""),
"type": type(node).__name__,
"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', '<anon>')}\nDocstring: {docstring}\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"
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)
def _extract_import_info(node):
"""Extract import information for module dependencies."""
if isinstance(node, ast.Import):
return {
"type": "import",
"modules": [alias.name for alias in node.names],
"level": 0
}
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
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]
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 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
# Try to find the containing class
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:
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))
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"):
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
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):
"""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)
return None
def parse_python_file(filepath: Path) -> List[Dict]:
"""Parse Python file with caching."""
file_hash = get_file_hash(filepath)
cached_result = parse_python_file_cached(str(filepath), file_hash)
return [{"text": text, "metadata": dict(meta)} for text, meta in cached_result]
@lru_cache(maxsize=1000)
def parse_sql_file_cached(filepath_str: str, file_hash: str) -> tuple:
"""Cached SQL parsing. Returns tuple for hashability."""
filepath = Path(filepath_str)
try:
sql_code = filepath.read_text(encoding="utf-8")
except Exception as e:
logger.warning(f"parse_sql_file: cannot read {filepath}: {e}")
return tuple()
chunks: List[Dict[str, Any]] = []
statements = [s for s in sqlparse.split(sql_code) if s and s.strip()]
lines = sql_code.splitlines()
search_pos = 0
for stmt in statements:
start_idx = sql_code.find(stmt, search_pos)
if start_idx == -1:
start_idx = sql_code.find(stmt)
if start_idx == -1:
start_line = len(lines)
else:
start_line = sql_code[:start_idx].count("\n") + 1
search_pos = start_idx + len(stmt)
leading_comments = include_leading_comments(lines, start_line)
stmt_clean = stmt.strip()
meta: Dict[str, Any] = {
"file": str(filepath),
"name": None,
"type": None,
"line": start_line,
"language": "sql"
}
try:
# Special handling for CREATE TYPE ENUM (PostgreSQL)
enum_match = re.match(
r'CREATE\s+TYPE\s+(\w+)\s+AS\s+ENUM\s*\((.*?)\)',
stmt_clean,
re.IGNORECASE | re.DOTALL
)
if enum_match:
enum_name = enum_match.group(1)
enum_values = enum_match.group(2)
meta["type"] = "create type enum"
meta["name"] = enum_name
meta["enum_values"] = ",".join([v.strip().strip("'\"") for v in enum_values.split(",") if v.strip()])
chunk_text = f"File: {filepath}\nType: {meta['type']}\nName: {meta['name']}\n"
if leading_comments:
chunk_text += f"\nComments:\n{leading_comments}\n\n"
chunk_text += f"SQL:\n{stmt_clean}"
chunks.append({"text": chunk_text, "metadata": meta})
continue
# Special handling for CREATE EXTENSION
extension_match = re.match(
r'CREATE\s+EXTENSION\s+(?:IF\s+NOT\s+EXISTS\s+)?(\w+)',
stmt_clean,
re.IGNORECASE
)
if extension_match:
extension_name = extension_match.group(1)
meta["type"] = "create extension"
meta["name"] = extension_name
chunk_text = f"File: {filepath}\nType: {meta['type']}\nName: {meta['name']}\n"
if leading_comments:
chunk_text += f"\nComments:\n{leading_comments}\n\n"
chunk_text += f"SQL:\n{stmt_clean}"
chunks.append({"text": chunk_text, "metadata": meta})
continue
# Special handling for CREATE TRIGGER
trigger_match = re.match(
r'CREATE\s+(?:OR\s+REPLACE\s+)?TRIGGER\s+(\w+)',
stmt_clean,
re.IGNORECASE
)
if trigger_match:
trigger_name = trigger_match.group(1)
meta["type"] = "create trigger"
meta["name"] = trigger_name
chunk_text = f"File: {filepath}\nType: {meta['type']}\nName: {meta['name']}\n"
if leading_comments:
chunk_text += f"\nComments:\n{leading_comments}\n\n"
chunk_text += f"SQL:\n{stmt_clean}"
chunks.append({"text": chunk_text, "metadata": meta})
continue
# Try sqlglot for standard SQL
parsed = sqlglot.parse_one(stmt_clean, read="postgres")
stmt_type = getattr(parsed, "key", None) or parsed.token_type if hasattr(parsed, "token_type") else None
meta["type"] = str(stmt_type).lower() if stmt_type else "statement"
tables = [t.this for t in parsed.find_all(sqlglot.exp.Table)]
meta["tables"] = ",".join([str(t) for t in tables if isinstance(t, str)])
ctes = []
for cte in parsed.find_all(sqlglot.exp.CTE):
try:
alias = cte.alias_or_name
if alias:
ctes.append(alias)
except Exception:
pass
meta["ctes"] = ",".join(ctes)
funcs = []
for f in parsed.find_all(sqlglot.exp.Func):
try:
name = f.name
if name:
funcs.append(name)
except Exception:
pass
meta["functions"] = ",".join(funcs)
if meta.get("tables") and meta["tables"]:
meta["name"] = meta["tables"].split(",")[0]
elif ctes:
meta["name"] = ctes[0]
elif funcs:
meta["name"] = funcs[0]
chunk_text = f"File: {filepath}\nType: {meta.get('type')}\nName: {meta.get('name')}\n"
if leading_comments:
chunk_text += f"\nComments:\n{leading_comments}\n\n"
chunk_text += f"SQL:\n{stmt_clean}"
chunks.append({"text": chunk_text, "metadata": meta})
continue
except Exception:
pass
# Fallback: simple heuristics if sqlglot fails
try:
parsed_tok = sqlparse.parse(stmt_clean)[0]
first_token = parsed_tok.token_first(skip_cm=True)
stmt_type = first_token.value.upper() if first_token else "UNKNOWN"
except Exception:
stmt_type = "UNKNOWN"
meta["type"] = stmt_type.lower()
chunk_text = f"File: {filepath}\nType: {meta['type']}\nName: {meta.get('name')}\n"
if leading_comments:
chunk_text += f"\nComments:\n{leading_comments}\n\n"
chunk_text += f"SQL:\n{stmt_clean}"
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)
cached_result = parse_sql_file_cached(str(filepath), file_hash)
return [{"text": text, "metadata": dict(meta)} for text, meta in cached_result]
def include_leading_comments(lines: List[str], stmt_start_line: int, max_context_lines: int = 8) -> str:
"""
Collect contiguous comment lines immediately above stmt_start_line (1-based).
Supports single-line comments (--), and block comments (/* ... */).
Returns a string containing the comment block (no leading comment markers).
"""
comments: List[str] = []
idx = stmt_start_line - 2 # convert to 0-based index of the line above the statement
collected = 0
# Collect single-line comments (-- ...) and blank lines (include blanks between comments)
while idx >= 0 and collected < max_context_lines:
raw = lines[idx]
stripped = raw.strip()
if stripped.startswith("--"):
# remove leading -- and optional space
comments.insert(0, stripped[2:].lstrip())
idx -= 1
collected += 1
continue
# block comment end
if stripped.endswith("*/"):
# gather entire block
block_lines = []
while idx >= 0:
block_line = lines[idx].rstrip()
block_lines.insert(0, block_line)
if block_line.strip().startswith("/*"):
break
idx -= 1
# remove /* and */ markers and join
cleaned = []
for bl in block_lines:
s = bl.strip()
if s.startswith("/*"):
s = s[2:].lstrip()
if s.endswith("*/"):
s = s[:-2].rstrip()
cleaned.append(s)
comments = cleaned + comments
break
# blank line - allow if we already have comments (so comments may be separated by one blank)
if stripped == "":
if comments:
comments.insert(0, "")
idx -= 1
collected += 1
continue
else:
break
# otherwise stop when non-comment found
break
# join into paragraph
if comments:
# remove possible leading/trailing empty lines
while comments and comments[0] == "":
comments.pop(0)
while comments and comments[-1] == "":
comments.pop()
return "\n".join(comments).strip()
return ""
# -----------------------------
# Go AST via helper binary with caching
# -----------------------------
# Replace the parse_go_file_cached function in mcp_codebase.py (around line 685)
# The issue is that the Go binary outputs lowercase JSON keys, but Python expects uppercase
@lru_cache(maxsize=1000)
def parse_go_file_cached(filepath_str: str, file_hash: str) -> tuple:
"""Cached Go AST parsing with enhanced metadata for graph relationships."""
filepath = Path(filepath_str)
chunks = []
helper = Path("./tools/parse_go_ast")
if helper.exists():
try:
proc = subprocess.run(
[str(helper), str(filepath)],
capture_output=True,
text=True,
check=True,
timeout=20
)
decls = json.loads(proc.stdout)
lines = filepath.read_text(encoding="utf-8").splitlines(keepends=True)
for d in decls:
# Use lowercase keys (matching the Go JSON output)
start = max(0, d.get("start_line", 1) - 1)
end = d.get("end_line", start + 1)
name = d.get("name", "")
typ = d.get("type", "")
doc_comment = d.get("doc_comment", "")
receiver = d.get("receiver", "")
fields = d.get("fields", [])
methods = d.get("methods", [])
# Skip package declarations
if typ == "package":
continue
# Build display type
display_type = typ
if receiver:
display_type = f"method ({receiver})"
chunk_code = "".join(lines[start:end])
chunk_text = f"File: {filepath}\nType: {display_type}\nName: {name}\n"
if doc_comment:
chunk_text += f"Doc:\n{doc_comment}\n\n"
if fields:
chunk_text += f"Fields: {', '.join(fields)}\n"
if methods:
chunk_text += f"Methods: {', '.join(methods)}\n"
chunk_text += f"Code:\n{chunk_code}"
# Build comprehensive metadata
metadata = {
"file": str(filepath),
"name": name,
"type": typ,
"receiver": receiver,
"line": start + 1,
"language": "go",
}
# Enhanced metadata for graph relationships
if fields:
metadata["fields"] = fields
if methods:
metadata["interface_methods"] = methods
if receiver:
metadata["receiver_type"] = receiver
# Extract receiver type without pointer for type matching
receiver_base = receiver.replace('*', '')
metadata["receiver_base_type"] = receiver_base
chunks.append({
"text": chunk_text,
"metadata": metadata
})
if chunks: # If helper worked, return its results
return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)
except subprocess.CalledProcessError as e:
logger.warning(f"parse_go_ast failed for {filepath}: {e.stderr}; falling back to regex")
except json.JSONDecodeError as e:
logger.warning(f"parse_go_ast output invalid JSON for {filepath}: {e}; falling back to regex")
except Exception as e:
logger.warning(f"parse_go_file helper failed for {filepath}: {e}; falling back to regex")
# Fallback: Use regex-based parsing
try:
content = filepath.read_text(encoding="utf-8")
lines = content.splitlines(keepends=True)
# Pattern to match Go function declarations
func_pattern = re.compile(
r'^func\s+(?:\([^)]+\)\s+)?(\w+)\s*\([^)]*\)(?:\s*\([^)]*\)|\s+[\w\[\].*]+)?\s*\{',
re.MULTILINE
)
for match in func_pattern.finditer(content):
func_name = match.group(1)
start_pos = match.start()
start_line = content[:start_pos].count('\n')
end_line = find_brace_block_end_go(lines, start_line)
func_code = "".join(lines[start_line:end_line + 1])
comments = extract_go_comments(lines, start_line)
# Try to extract receiver type from function signature
receiver_type = ""
func_sig = content[start_pos:match.end()]
receiver_match = re.search(r'func\s*\(([^)]+)\)', func_sig)
if receiver_match:
receiver_part = receiver_match.group(1)
# Extract type from receiver (e.g., "u *User" -> "*User")
receiver_type_match = re.search(r'\*?(\w+)', receiver_part.split()[-1] if ' ' in receiver_part else receiver_part)
if receiver_type_match:
receiver_type = receiver_type_match.group(0)
chunk_text = f"File: {filepath}\nType: function\nName: {func_name}\n"
if comments:
chunk_text += f"Comments:\n{comments}\n\n"
chunk_text += f"Code:\n{func_code}"
metadata = {
"file": str(filepath),
"name": func_name,
"type": "function",
"line": start_line + 1,
"language": "go",
}
# Receiver type for methods
if receiver_type:
metadata["receiver_type"] = receiver_type
metadata["receiver_base_type"] = receiver_type.replace('*', '')
metadata["type"] = "method"
chunks.append({
"text": chunk_text,
"metadata": metadata
})
# Structs with field extraction
struct_pattern = re.compile(r'^type\s+(\w+)\s+struct\s*\{', re.MULTILINE)
for match in struct_pattern.finditer(content):
struct_name = match.group(1)
start_pos = match.start()
start_line = content[:start_pos].count('\n')
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)
if struct_body_match:
field_lines = struct_body_match.group(1).split('\n')
for line in field_lines:
field_match = re.match(r'\s*(\w+)\s+', line.strip())
if field_match:
field_names.append(field_match.group(1))
chunk_text = f"File: {filepath}\nType: struct\nName: {struct_name}\n"
if comments:
chunk_text += f"Comments:\n{comments}\n\n"
if field_names:
chunk_text += f"Fields: {', '.join(field_names)}\n"
chunk_text += f"Code:\n{struct_code}"
metadata = {
"file": str(filepath),
"name": struct_name,
"type": "struct",
"line": start_line + 1,
"language": "go",
}
# Struct fields for graph relationships
if field_names:
metadata["fields"] = field_names
chunks.append({
"text": chunk_text,
"metadata": metadata
})
# Interfaces with method extraction
interface_pattern = re.compile(r'^type\s+(\w+)\s+interface\s*\{', re.MULTILINE)
for match in interface_pattern.finditer(content):
interface_name = match.group(1)
start_pos = match.start()
start_line = content[:start_pos].count('\n')
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)
if interface_body_match:
method_lines = interface_body_match.group(1).split('\n')
for line in method_lines:
method_match = re.match(r'\s*(\w+)\s*\([^)]*\)', line.strip())
if method_match:
method_names.append(method_match.group(1))
chunk_text = f"File: {filepath}\nType: interface\nName: {interface_name}\n"
if comments:
chunk_text += f"Comments:\n{comments}\n\n"
if method_names:
chunk_text += f"Methods: {', '.join(method_names)}\n"
chunk_text += f"Code:\n{interface_code}"
metadata = {
"file": str(filepath),
"name": interface_name,
"type": "interface",
"line": start_line + 1,
"language": "go",
}
# Interface methods for graph relationships
if method_names:
metadata["interface_methods"] = method_names
chunks.append({
"text": chunk_text,
"metadata": metadata
})
except Exception as e:
logger.warning(f"Regex-based Go parsing failed for {filepath}: {e}")
return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)
def find_brace_block_end_go(lines: List[str], start_line: int) -> int:
"""Find the end of a Go brace block starting at start_line."""
brace_count = 0
in_block = False
for i in range(start_line, len(lines)):
line = lines[i]
for char in line:
if char == '{':
brace_count += 1
in_block = True
elif char == '}':
brace_count -= 1
if in_block and brace_count == 0:
return i
return len(lines) - 1
def extract_go_comments(lines: List[str], func_line_index: int, max_lines: int = 10) -> str:
"""Extract leading comments above a Go function/struct/interface."""
comments = []
idx = func_line_index - 1
while idx >= 0 and len(comments) < max_lines:
line = lines[idx].strip()
if line.startswith('//'):
# Remove // and trim
comments.insert(0, line[2:].strip())
idx -= 1
elif line.startswith('/*') or '*/' in line:
# Block comment - collect entire block
if '*/' in line and not line.startswith('/*'):
# End of block, go backwards to find start
block_lines = [line]
idx -= 1
while idx >= 0:
block_line = lines[idx].strip()
block_lines.insert(0, block_line)
if block_line.startswith('/*'):
break
idx -= 1
# Clean up block comment markers
block_text = ' '.join(block_lines)
block_text = block_text.replace('/*', '').replace('*/', '').replace('*', '').strip()
comments.insert(0, block_text)
idx -= 1
else:
# Single line block comment
clean = line.replace('/*', '').replace('*/', '').strip()
comments.insert(0, clean)
idx -= 1
elif line == '':
# Allow one blank line
if comments:
idx -= 1
else:
break
else:
break
# Clean up
while comments and comments[0] == '':
comments.pop(0)
while comments and comments[-1] == '':
comments.pop(-1)
return '\n'.join(comments)
def parse_go_file(filepath: Path) -> List[Dict]:
"""Parse Go file with caching."""
file_hash = get_file_hash(filepath)
cached_result = parse_go_file_cached(str(filepath), file_hash)
if not cached_result:
return chunk_text_file(filepath, "go")
return [{"text": text, "metadata": dict(meta)} for text, meta in cached_result]
# -----------------------------
# Java parsing (javalang + manual body extraction) with caching
# -----------------------------
@lru_cache(maxsize=1000)
def parse_java_file_cached(filepath_str: str) -> tuple:
"""Cached Java parsing. Returns tuple for hashability."""
filepath = Path(filepath_str)
chunks = []
try:
code = filepath.read_text(encoding="utf-8")
lines = code.splitlines(keepends=True)
tree = javalang.parse.parse(code)
def find_brace_block_end(start_idx: int) -> int:
brace_count = 0
for i in range(start_idx, len(lines)):
for char in lines[i]:
if char == '{':
brace_count += 1
elif char == '}':
brace_count -= 1
if brace_count == 0:
return i
return len(lines) - 1
for path, node in tree:
if isinstance(node, javalang.tree.ClassDeclaration):
start_line = node.position.line - 1 if node.position else 0
end_line = find_brace_block_end(start_line)
chunk_text = "".join(lines[start_line:end_line + 1])
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": {
"file": str(filepath),
"name": node.name,
"type": "class",
"line": start_line + 1,
"language": "java",
"extends": extends_class,
"implements": implements_interfaces
}
})
elif isinstance(node, javalang.tree.MethodDeclaration):
if not hasattr(node, 'position') or not node.position:
continue
start_line = node.position.line - 1
end_line = find_brace_block_end(start_line)
chunk_text = "".join(lines[start_line:end_line + 1])
doc = get_context_with_comments(filepath, start_line + 1)
# Find enclosing class name
class_name = None
for p in reversed(path):
if isinstance(p, javalang.tree.ClassDeclaration):
class_name = p.name
break
return_type = node.return_type.name if node.return_type else "void"
parameters = []
if node.parameters:
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": {
"file": str(filepath),
"name": node.name,
"type": "method",
"class": class_name,
"line": start_line + 1,
"language": "java",
"return_type": return_type,
"parameters": parameters
}
})
elif isinstance(node, javalang.tree.FieldDeclaration):
# Fields may have multiple variable declarators
if not hasattr(node, 'position') or not node.position:
continue
start_line = node.position.line - 1
# Field is usually one line, but include next if annotation
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"
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}",
"metadata": {
"file": str(filepath),
"name": declarator.name,
"type": "field",
"line": start_line + 1,
"language": "java",
"field_type": field_type
}
})
elif isinstance(node, javalang.tree.ConstructorDeclaration):
if not hasattr(node, 'position') or not node.position:
continue
start_line = node.position.line - 1
end_line = find_brace_block_end(start_line)
chunk_text = "".join(lines[start_line:end_line + 1])
doc = get_context_with_comments(filepath, start_line + 1)
class_name = None
for p in reversed(path):
if isinstance(p, javalang.tree.ClassDeclaration):
class_name = p.name
break
chunks.append({
"text": f"File: {filepath}\nType: Constructor\nClass: {class_name or 'unknown'}\nDoc:\n{doc}\nCode:\n{chunk_text}",
"metadata": {
"file": str(filepath),
"name": f"{class_name} constructor",
"type": "constructor",
"class": class_name,
"line": start_line + 1,
"language": "java"
}
})
elif isinstance(node, javalang.tree.EnumDeclaration):
start_line = node.position.line - 1 if node.position else 0
end_line = find_brace_block_end(start_line)
chunk_text = "".join(lines[start_line:end_line + 1])
doc = get_context_with_comments(filepath, start_line + 1)
chunks.append({
"text": f"File: {filepath}\nType: Enum\nName: {node.name}\nDoc:\n{doc}\nCode:\n{chunk_text}",
"metadata": {
"file": str(filepath),
"name": node.name,
"type": "enum",
"line": start_line + 1,
"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)
chunk_text = "".join(lines[start_line:end_line + 1])
doc = get_context_with_comments(filepath, start_line + 1)
chunks.append({
"text": f"File: {filepath}\nType: Interface\nName: {node.name}\nDoc:\n{doc}\nCode:\n{chunk_text}",
"metadata": {
"file": str(filepath),
"name": node.name,
"type": "interface",
"line": start_line + 1,
"language": "java"
}
})
elif isinstance(node, javalang.tree.AnnotationDeclaration):
start_line = node.position.line - 1 if node.position else 0
end_line = find_brace_block_end(start_line)
chunk_text = "".join(lines[start_line:end_line + 1])
doc = get_context_with_comments(filepath, start_line + 1)
chunks.append({
"text": f"File: {filepath}\nType: Annotation\nName: {node.name}\nDoc:\n{doc}\nCode:\n{chunk_text}",
"metadata": {
"file": str(filepath),
"name": node.name,
"type": "annotation",
"line": start_line + 1,
"language": "java"
}
})
except javalang.parser.JavaSyntaxError as e:
logger.warning(f"Failed to parse Java file {filepath}: {e}")
except Exception as e:
logger.warning(f"Unexpected error parsing Java file {filepath}: {e}")
return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)
def parse_java_file(filepath: Path) -> List[Dict]:
"""Parse Java source file with caching."""
cached_result = parse_java_file_cached(str(filepath))
return [{"text": text, "metadata": dict(meta)} for text, meta in cached_result]
# -----------------------------
# Svelte support with caching
# -----------------------------
@lru_cache(maxsize=500)
def parse_svelte_file_cached(filepath_str: str, file_hash: str) -> tuple:
"""Cached Svelte parsing with enhanced multi-language graph relationships."""
filepath = Path(filepath_str)
chunks = []
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"<script(?:\s+[^>]*)?>(.*?)</script>", 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 <script ...> attributes
# Extract script context (module, context="module")
is_module = 'context="module"' in script_attrs
script_type = "script_module" if is_module else "script"
ts_helper = Path("./tools/parse_ts.js")
if ts_helper.exists():
tmp = filepath.parent / f".tmp_{filepath.name}_{idx}.ts"
try:
tmp.write_text(script_content, encoding="utf-8")
proc = subprocess.run(
["node", str(ts_helper), str(tmp)],
capture_output=True, text=True, check=True, timeout=15
)
decls = json.loads(proc.stdout)
script_lines = script_content.splitlines(keepends=True)
for d in decls:
start_rel = max(0, d.get("start_line", 1) - 1)
end_rel = d.get("end_line", start_rel + 1)
chunk_code = "".join(script_lines[start_rel:end_rel])
doc = d.get("doc_comment", "").strip()
abs_start = script_start_line + start_rel
# Extract component relationships
decl_name = d.get("name")
decl_type = d.get("type")
# Track props for component interface
if decl_type == "variable" and decl_name and not is_module:
# Look for prop patterns: export let prop = ...
if re.search(r'export\s+let\s+' + re.escape(decl_name), chunk_code):
component_props.append(decl_name)
decl_type = "prop"
# Track reactive declarations: $: reactiveVar = ...
elif chunk_code.strip().startswith('$:'):
component_context.append(decl_name)
decl_type = "reactive"
chunk_text = f"File: {filepath}\nType: {decl_type}\nName: {decl_name}\nDoc: {doc}\nCode:\n{chunk_code}"
metadata = {
"file": str(filepath),
"name": decl_name,
"type": decl_type,
"line": abs_start,
"language": "typescript",
"block_type": script_type,
"component_name": component_name
}
# Add specific metadata for different declaration types
if decl_type == "prop":
metadata["is_exported"] = True
metadata["component_prop"] = True
elif decl_type == "reactive":
metadata["is_reactive"] = True
elif decl_type == "function":
metadata["is_function"] = True
# Try to detect event dispatchers
if "createEventDispatcher" in chunk_code:
metadata["is_event_dispatcher"] = True
chunks.append({
"text": chunk_text,
"metadata": metadata
})
except Exception as e:
logger.warning(f"TS parser failed for {filepath}: {e}; falling back to heuristic")
chunks.append(_make_script_chunk(filepath, script_content, script_start_line, script_type, component_name))
finally:
tmp.unlink(missing_ok=True)
else:
chunks.append(_make_script_chunk(filepath, script_content, script_start_line, script_type, component_name))
idx += 1
# Style blocks with CSS relationship extraction
for m in re.finditer(r"<style(?:\s+[^>]*)?>(.*?)</style>", text, flags=re.DOTALL):
style_content = m.group(1)
style_start_line = text[:m.start(1)].count("\n") + 1
# Extract CSS classes and relationships
css_classes = re.findall(r'\.([a-zA-Z][\w-]*)\s*{', style_content)
css_selectors = re.findall(r'([a-zA-Z][\w-]*)\s*{', style_content)
chunk_text = f"File: {filepath}\nType: style\nComponent: {component_name}\n"
if css_classes:
chunk_text += f"CSS Classes: {', '.join(css_classes)}\n"
chunk_text += f"Style Content:\n{style_content}"
chunks.append({
"text": chunk_text,
"metadata": {
"file": str(filepath),
"name": f"style_block_{idx}",
"type": "style",
"line": style_start_line,
"language": "css",
"block_type": "style",
"component_name": component_name,
"css_classes": css_classes,
"css_selectors": css_selectors
}
})
idx += 1
# Markup with Svelte-specific element analysis
markup = re.sub(r"<script[^>]*>.*?</script>", "", text, flags=re.DOTALL)
markup = re.sub(r"<style[^>]*>.*?</style>", "", 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"
if prop_bindings:
chunk_text += f"Prop Bindings: {', '.join([f'{prop}={val}' for prop, val in prop_bindings[:5]])}\n"
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": {
"file": str(filepath),
"name": "markup",
"type": "markup",
"line": 1,
"language": "svelte",
"block_type": "markup",
"component_name": component_name,
"used_components": list(set(used_components)),
"prop_bindings": [prop for prop, _ in prop_bindings],
"event_handlers": list(set(event_handlers)),
"component_props": component_props, # Props discovered in scripts
"reactive_vars": component_context # Reactive context discovered
}
})
except Exception as e:
logger.warning(f"parse_svelte_file failed {filepath}: {e}")
return tuple()
return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)
def _make_script_chunk(filepath, script_content, start_line, block_type, component_name=None):
"""Enhanced script chunk creation with basic analysis."""
# Basic heuristic analysis for when TypeScript parser fails
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"
if exports:
chunk_text += f"Exports: {', '.join(exports)}\n"
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",
"type": block_type,
"line": start_line,
"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]:
"""Parse Svelte file with caching."""
file_hash = get_file_hash(filepath)
cached_result = parse_svelte_file_cached(str(filepath), file_hash)
if not cached_result:
return chunk_text_file(filepath, "svelte")
return [{"text": text, "metadata": dict(meta)} for text, meta in cached_result]
# -----------------------------
# Heuristic chunking for other languages
# -----------------------------
def chunk_text_file(filepath: Path, language: str) -> List[Dict]:
try:
lines = filepath.read_text(encoding="utf-8").splitlines()
except Exception as e:
logger.warning(f"chunk_text_file read failed {filepath}: {e}")
return []
chunks = []
current_chunk = []
start_line = 0
for i, ln in enumerate(lines):
stripped = ln.strip()
if stripped.startswith(("def ", "func ", "fn ", "class ", "struct ", "impl ", "export ", "const ", "let ")) or stripped.startswith(("//", "#", "--")) or stripped.endswith("{"):
if current_chunk:
text = "\n".join(current_chunk)
chunks.append({
"text": f"File: {filepath}\nLanguage: {language}\n\n{text}",
"metadata": {"file": str(filepath), "name": f"chunk_{len(chunks)}", "type": "code_block", "line": start_line + 1, "language": language}
})
current_chunk = []
start_line = i
current_chunk.append(ln)
if current_chunk:
text = "\n".join(current_chunk)
chunks.append({
"text": f"File: {filepath}\nLanguage: {language}\n\n{text}",
"metadata": {"file": str(filepath), "name": f"chunk_{len(chunks)}", "type": "code_block", "line": start_line + 1, "language": language}
})
return chunks
def parse_rust_file(filepath: Path) -> List[Dict]:
"""Parse Rust file with caching."""
file_hash = get_file_hash(filepath)
cached_result = parse_rust_file_cached(str(filepath), file_hash)
return [{"text": text, "metadata": dict(meta)} for text, meta in cached_result]
@lru_cache(maxsize=1000)
def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple:
"""Cached Rust parsing with enhanced metadata for graph relationships.
ALWAYS returns a tuple of (text, metadata_tuple) entries. If the Rust helper
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 Rust file: {filepath}")
rust_helper = Path("./tools/target/release/parse_rust_ast")
if not rust_helper.is_file():
logger.warning(f"Rust parser binary not found: {rust_helper}. Using fallback whole-file chunk for {filepath}")
code = filepath.read_text(encoding="utf-8")
fallback_meta = {
"file": str(filepath),
"name": filepath.name,
"type": "file",
"line": 1,
"language": "rust",
}
return ((code, tuple(fallback_meta.items())),)
try:
code = filepath.read_text(encoding="utf-8")
lines = code.splitlines(keepends=True)
proc = subprocess.run(
[str(rust_helper), str(filepath)],
capture_output=True,
text=True,
check=True,
timeout=20
)
decls = json.loads(proc.stdout)
logger.info(f" AST helper found {len(decls)} declarations")
for i, d in enumerate(decls):
start_line = max(0, d.get("start_line", 1) - 1)
end_line = d.get("end_line", start_line + 1)
# guard the end_line to be at least start_line+1
if end_line <= start_line:
end_line = start_line + 1
# Make slice safe even if end_line > len(lines)
chunk_code = "".join(lines[start_line:end_line])
metadata = {
"file": str(filepath),
"name": d.get("name", "") or filepath.name,
"type": d.get("type_", "") or "unknown",
"line": start_line + 1,
"language": "rust",
}
# copy other fields defensively
if d.get("visibility"):
metadata["visibility"] = d.get("visibility")
if d.get("is_async"):
metadata["is_async"] = True
if d.get("is_unsafe"):
metadata["is_unsafe"] = True
if d.get("generics"):
metadata["generics"] = d.get("generics")
if d.get("traits"):
metadata["implements_traits"] = d.get("traits")
if d.get("fields"):
metadata["fields"] = d.get("fields")
if d.get("methods"):
metadata["methods"] = d.get("methods")
if d.get("return_type"):
metadata["return_type"] = d.get("return_type")
if d.get("parameters"):
metadata["parameters"] = d.get("parameters")
logger.info(f" Declaration {i+1}: {metadata.get('type', 'unknown')} {metadata.get('name', 'unnamed')}")
chunk_text = f"File: {filepath}\nType: {metadata.get('type')}\nName: {metadata.get('name')}\n"
if metadata.get("visibility"):
chunk_text += f"Visibility: {metadata.get('visibility')}\n"
if metadata.get("is_async"):
chunk_text += "Async: yes\n"
if metadata.get("is_unsafe"):
chunk_text += "Unsafe: yes\n"
if metadata.get("generics"):
chunk_text += f"Generics: {', '.join(metadata.get('generics', []))}\n"
if metadata.get("implements_traits"):
chunk_text += f"Implements: {', '.join(metadata.get('implements_traits', []))}\n"
if metadata.get("fields"):
chunk_text += f"Fields: {', '.join(metadata.get('fields', []))}\n"
chunk_text += f"Code:\n{chunk_code}"
chunks.append({
"text": chunk_text,
"metadata": metadata
})
except Exception as e:
logger.warning(f"Rust AST helper failed for {filepath}: {e}")
# If parser produced chunks, return them in the (text, metadata_tuple) format expected by parse_rust_file()
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 to keep counts consistent
logger.warning(f"❌ No chunks created from Rust AST 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.name,
"type": "file",
"line": 1,
"language": "rust",
}
return ((full_code, tuple(fallback_meta.items())),)
@lru_cache(maxsize=1000)
def parse_shell_file_cached(filepath_str: str) -> tuple:
"""Cached shell script parsing. Returns tuple for hashability."""
filepath = Path(filepath_str)
try:
content = filepath.read_text(encoding="utf-8")
chunks = []
lines = content.splitlines(keepends=True)
patterns = [
r'^(\w+)\s*\(\s*\)\s*\{',
r'^function\s+(\w+)\s*\{',
r'^function\s+(\w+)\s*\(\s*\)\s*\{',
r'^def\s+(\w+)\s*\(\s*\)\s*\{',
]
for i, line in enumerate(lines):
line_text = line.strip()
for pattern in patterns:
match = re.match(pattern, line_text)
if match:
func_name = match.group(1)
start_line = i
end_line = find_shell_function_end(lines, i)
func_code = "".join(lines[start_line:end_line + 1])
# Extract leading comments
leading_comments = extract_shell_comments(lines, start_line)
chunk_text = f"File: {filepath}\nType: function\nName: {func_name}\n"
if leading_comments:
chunk_text += f"Comments:\n{leading_comments}\n\n"
chunk_text += f"Code:\n{func_code}"
chunks.append({
"text": chunk_text,
"metadata": {
"file": str(filepath),
"name": func_name,
"type": "function",
"line": start_line + 1,
"language": "shell",
"has_braces": "{" in line_text
}
})
break # Move to next line after finding a match
# If no functions found, fall back to basic chunking
return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks) if chunks else tuple()
except Exception as e:
logger.warning(f"parse_shell_file failed {filepath}: {e}")
return tuple()
def parse_shell_file(filepath: Path) -> List[Dict]:
"""Enhanced shell script parsing with caching."""
cached_result = parse_shell_file_cached(str(filepath))
if not cached_result:
return chunk_text_file(filepath, "shell")
return [{"text": text, "metadata": dict(meta)} for text, meta in cached_result]
def extract_shell_comments(lines: List[str], function_line_index: int, max_lines: int = 10) -> str:
"""Extract leading comments above a shell function."""
comments = []
idx = function_line_index - 1
while idx >= 0 and len(comments) < max_lines:
line = lines[idx].strip()
if line.startswith('#'):
comments.insert(0, line[1:].strip()) # Remove the # and trim
idx -= 1
elif line == '':
# Allow one blank line between comments
if comments:
comments.insert(0, "")
idx -= 1
else:
break
else:
break
# Clean up: remove leading/trailing empty lines
while comments and comments[0] == "":
comments.pop(0)
while comments and comments[-1] == "":
comments.pop(-1)
return "\n".join(comments)
def find_shell_function_end(lines: List[str], start_line_index: int) -> int:
"""
Find the end of a shell function by tracking brace nesting.
Args:
lines: List of file lines
start_line_index: Starting line index (0-based) of the function
Returns:
Line index (0-based) where the function ends
"""
brace_count = 0
in_function = False
for i in range(start_line_index, len(lines)):
line = lines[i]
# Count opening and closing braces
for char in line:
if char == '{':
brace_count += 1
in_function = True
elif char == '}':
brace_count -= 1
# If we've returned to brace_count 0 and we were in a function, this is the end
if in_function and brace_count == 0:
return i
# Special case: shell functions without braces (single line)
if not in_function and i > start_line_index:
# Look for function end patterns
stripped = line.strip()
if (stripped.startswith(('function ', 'def ')) or
re.match(r'^\w+\(\s*\)', stripped) or
stripped.endswith(')') and '()' not in line):
# New function starting, so previous one ended
return i - 1
# If we never found the end, return the last line
return len(lines) - 1
# -----------------------------
# Detect language helper
# -----------------------------
def detect_language(filepath: Path) -> str:
ext = filepath.suffix.lower()
lang_map = {
".py": "python", ".go": "go", ".rs": "rust", ".ts": "typescript", ".tsx": "typescript",
".js": "javascript", ".jsx": "javascript", ".svelte": "svelte", ".java": "java",
".sql": "sql", ".pgsql": "sql", ".sh": "shell"
}
return lang_map.get(ext, "unknown")
# -----------------------------
# Indexing codebase
# -----------------------------
def index_codebase(codebase_path: Path) -> Tuple[List[str], List[Dict[str, Any]]]:
ragignore = load_ragignore(codebase_path)
all_texts = []
all_metadata = []
logger.info(f"Indexing codebase at {codebase_path}")
for filepath in codebase_path.rglob("*"):
if not filepath.is_file():
continue
rel = filepath.relative_to(codebase_path)
if ragignore.match_file(str(rel)):
continue
if filepath.suffix.lower() not in SUPPORTED_EXTENSIONS:
continue
language = detect_language(filepath)
if language == "python":
chunks = parse_python_file(filepath)
elif language == "go":
chunks = parse_go_file(filepath)
elif language == "java":
chunks = parse_java_file(filepath)
elif language == "svelte":
chunks = parse_svelte_file(filepath)
elif language == "sql":
chunks = parse_sql_file(filepath)
elif language == "shell":
chunks = parse_shell_file(filepath)#
elif language == "rust":
chunks = parse_rust_file(filepath)
else:
chunks = chunk_text_file(filepath, language)
for c in chunks:
all_texts.append(c["text"])
all_metadata.append(clean_metadata_for_chroma(c["metadata"]))
logger.info(f"Indexed {len(all_texts)} chunks")
return all_texts, all_metadata
def clean_metadata_for_chroma(metadata: Dict[str, Any]) -> Dict[str, Any]:
"""Clean metadata to only include str, int, float, bool values (no None)."""
cleaned = {}
for key, value in metadata.items():
if value is None:
cleaned[key] = "" # Convert None to empty string
elif isinstance(value, (str, int, float, bool)):
cleaned[key] = value
elif isinstance(value, list):
# Convert lists to comma-separated strings
cleaned[key] = ",".join(str(v) for v in value if v is not None)
else:
# Convert other types to strings
cleaned[key] = str(value)
return cleaned
def build_indexes():
global vectorstore, bm25, bm25_corpus, chunks_metadata, embeddings, index_build_time
logger.info("Building indexes with contextual weighting and batch embedding...")
start_time = time.time()
# -------------------------------------------------
# Step 1: Parse codebase and build enhanced graph
# -------------------------------------------------
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
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
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
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
if "implements_traits" in m and m["implements_traits"]:
m["implements"] = m["implements_traits"]
# ---- File node ---------------------------------------------
# Keep the same ID (`file::<path>`) 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).
file_node_id = f"file::{fpath}"
if file_node_id not in graph.graph.nodes:
graph.add_node(
file_node_id,
type="File",
name=fpath,
path=fpath,
lang=lang,
)
created_nodes[file_node_id] = {
"type": "File",
"name": fpath,
"path": fpath,
"lang": lang,
}
# Create node ID with language
node_id = f"{lang}::{node_type}::{fpath}::{name}"
# Base attributes for all nodes
base_attrs = {
"name": name,
"file": fpath,
"line": m.get("line"),
"lang": lang # Always include language
}
# Store for second-pass processing
created_nodes[node_id] = base_attrs.copy()
created_nodes[node_id]["type"] = node_type
# 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 - 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
graph.add_node(node_id, **base_attrs, type=node_type,
class_name=class_name,
return_type=m.get("return_type"),
parameters=m.get("parameters", []))
# Connect method to its 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
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")
# 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", [])
elif node_type == "trait":
# Traits
graph.add_node(node_id, **base_attrs, type=node_type,
methods=m.get("methods", []))
graph.add_edge(file_node_id, node_id, "contains")
created_nodes[node_id]["methods"] = m.get("methods", [])
elif node_type == "module":
# Modules
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")
# ---------- IMPORTS ----------
if node_type == "import" and "parameters" in m:
imports_raw = m["parameters"] # override metadata pathway
else:
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
graph._add_import_edges(file_node_id, imports)
# ---------- 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
graph._add_call_edges(node_id, calls, fpath)
# -------------------------------------------------
# Step 2: Second-pass relationship building
# -------------------------------------------------
logger.info("Building second-pass relationships...")
# Build type relationships
type_connections = build_type_relationships(graph, created_nodes)
logger.info(f"Added {type_connections} type relationships")
# Build Bevy-specific relationships
bevy_connections = build_bevy_relationships(graph, created_nodes)
logger.info(f"Added {bevy_connections} Bevy-specific relationships")
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
# -------------------------------------------------
os.environ["OLLAMA_HOST"] = OLLAMA_BASE_URL
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL, base_url=OLLAMA_BASE_URL)
# Generate base embeddings in batches
logger.info(f"Generating embeddings for {len(texts)} documents in batches of {EMBEDDING_BATCH_SIZE}...")
base_embeddings = batch_embed_documents(texts)
# Apply contextual weights
logger.info("Applying contextual weights...")
weighted_embeddings = apply_contextual_weights_to_embeddings(base_embeddings, metadatas)
# 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."
)
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
)
vectorstore = Chroma.from_texts(
texts=texts,
embedding=weighted_emb_func,
persist_directory=str(VECTOR_DB_PATH),
metadatas=metadatas
)
# Build BM25 index
logger.info("Building BM25 index...")
tokenized = [t.lower().split() for t in texts]
bm25 = BM25Okapi(tokenized)
bm25_corpus = texts
chunks_metadata = metadatas
# Save BM25 index and metadata
BM25_INDEX_PATH.write_text(
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")
def build_type_relationships(graph, created_nodes):
"""Build type relationships between nodes (parameter types, return types, etc.)"""
connections = 0
for node_id, node_data in created_nodes.items():
node_type = node_data.get("type")
# Handle function/method parameter and return types
if node_type in ["function", "method", "constructor"]:
connections += _connect_function_types(graph, node_id, node_data)
# Handle class/struct inheritance and implementation
elif node_type in ["class", "struct", "interface"]:
connections += _connect_class_relationships(graph, node_id, node_data)
# Handle impl blocks
elif node_type == "impl":
connections += _connect_impl_relationships(graph, node_id, node_data)
return connections
def _connect_function_types(graph, node_id, node_data):
"""Connect function nodes to their parameter and return types"""
connections = 0
# Connect return type
return_type = node_data.get("return_type")
if return_type and _is_meaningful_type(return_type):
return_type_clean = _clean_type_name(return_type)
for type_id, type_data in graph.find_nodes(name=return_type_clean):
graph.add_edge(node_id, type_id, "returns")
connections += 1
# Connect parameter types
parameters = node_data.get("parameters", [])
for param in parameters:
param_type = _extract_parameter_type(param)
if param_type and _is_meaningful_type(param_type):
param_type_clean = _clean_type_name(param_type)
for type_id, type_data in graph.find_nodes(name=param_type_clean):
graph.add_edge(node_id, type_id, "uses_parameter")
connections += 1
return connections
def _connect_class_relationships(graph, node_id, node_data):
"""Connect class/struct nodes to parent classes and interfaces"""
connections = 0
# Inheritance
extends = node_data.get("extends")
if extends and _is_meaningful_type(extends):
for parent_id, parent_data in graph.find_nodes(name=extends):
graph.add_edge(node_id, parent_id, "extends")
connections += 1
# Interface implementation
implements = node_data.get("implements", [])
for interface in implements:
if _is_meaningful_type(interface):
for interface_id, interface_data in graph.find_nodes(name=interface):
graph.add_edge(node_id, interface_id, "implements")
connections += 1
return connections
def _connect_impl_relationships(graph, node_id, node_data):
"""Connect impl blocks to their targets and traits"""
connections = 0
target = node_data.get("target")
if target and _is_meaningful_type(target):
# Connect to target type
for target_id, target_data in graph.find_nodes(name=target):
graph.add_edge(node_id, target_id, "implements_for")
connections += 1
# Connect to implemented traits
traits = node_data.get("traits", [])
for trait in traits:
if _is_meaningful_type(trait):
for trait_id, trait_data in graph.find_nodes(name=trait):
graph.add_edge(node_id, trait_id, "implements")
connections += 1
return connections
def build_bevy_relationships(graph, created_nodes):
"""Build Bevy-specific relationships"""
connections = 0
# Find all plugin structs and impls
plugin_structs = {}
plugin_impls = {}
for node_id, node_data in created_nodes.items():
node_type = node_data.get("type")
name = node_data.get("name", "")
# Collect plugin structs
if node_type in ["struct", "class"] and name.endswith("Plugin"):
plugin_structs[name] = node_id
# Collect plugin impls
elif node_type == "impl" and node_data.get("target", "").endswith("Plugin"):
plugin_impls[node_data["target"]] = node_id
# Connect plugin impls to their structs
for plugin_name, impl_id in plugin_impls.items():
if plugin_name in plugin_structs:
struct_id = plugin_structs[plugin_name]
graph.add_edge(impl_id, struct_id, "implements_plugin")
connections += 1
# Find build method in this impl
impl_methods = created_nodes[impl_id].get("methods", [])
for method_name in impl_methods:
if method_name == "build":
# Find the method node
method_id = f"{created_nodes[impl_id]['lang']}::method::{created_nodes[impl_id]['file']}::{plugin_name}.{method_name}"
if method_id in created_nodes:
graph.add_edge(impl_id, method_id, "defines_plugin")
connections += 1
# Connect systems to plugins (simplified - same file assumption)
for node_id, node_data in created_nodes.items():
if node_data.get("type") in ["function", "method"]:
# Simple heuristic for Bevy systems
function_name = node_data.get("name", "").lower()
if any(keyword in function_name for keyword in ['system', 'update', 'setup']):
file_path = node_data.get("file")
# Look for plugins in the same file
for plugin_id, plugin_data in created_nodes.items():
if (plugin_data.get("file") == file_path and
plugin_data.get("type") in ["struct", "class"] and
plugin_data.get("name", "").endswith("Plugin")):
graph.add_edge(plugin_id, node_id, "contains_system")
connections += 1
return connections
def _extract_parameter_type(param):
"""Extract type from parameter string, handling multiple formats"""
if not param:
return None
# Handle "type: name" format (TypeScript/Python)
if ':' in param:
parts = param.split(':')
if len(parts) >= 2:
return parts[1].strip()
# Handle "Type name" format (Rust/Java/C++)
if ' ' in param:
parts = param.split()
if len(parts) >= 2:
return parts[0].strip()
# Handle complex types with generics
if '<' in param and '>' in param:
# Extract the base type before generics
base_type = param.split('<')[0].strip()
return base_type
return param.strip()
def _clean_type_name(type_name):
"""Clean type names by removing common modifiers"""
if not type_name:
return ""
# Remove common modifiers
modifiers = ['&', 'mut ', 'const ', 'static ', 'pub ', 'private ', 'protected ']
cleaned = type_name
for mod in modifiers:
cleaned = cleaned.replace(mod, '')
# Remove trailing references/pointers
cleaned = cleaned.rstrip('*&')
return cleaned.strip()
def _is_meaningful_type(type_name):
"""Check if a type name is meaningful (not primitive, void, etc.)"""
if not type_name:
return False
meaningless_types = [
'void', 'int', 'i32', 'i64', 'f32', 'f64', 'bool', 'str', 'String',
'char', 'u8', 'u16', 'u32', 'u64', 'usize', 'isize',
'self', '&self', '&mut self', 'mut self', 'Self',
'Result', 'Option', 'Vec', 'String'
]
cleaned = _clean_type_name(type_name)
return (cleaned and
cleaned not in meaningless_types and
not cleaned.startswith('&') and
len(cleaned) > 1)
# -----------------------------
# Build / load indexes with contextual weighting and batch embedding
# -----------------------------
def batch_embed_documents(texts: List[str]) -> List[List[float]]:
"""Embed documents in batches for better performance."""
all_embeddings = []
total_batches = (len(texts) + EMBEDDING_BATCH_SIZE - 1) // EMBEDDING_BATCH_SIZE
for i in range(0, len(texts), EMBEDDING_BATCH_SIZE):
batch = texts[i:i + EMBEDDING_BATCH_SIZE]
batch_num = i // EMBEDDING_BATCH_SIZE + 1
logger.info(f"Embedding batch {batch_num}/{total_batches} ({len(batch)} docs)...")
batch_embeddings = embeddings.embed_documents(batch)
all_embeddings.extend(batch_embeddings)
return all_embeddings
def load_indexes():
global vectorstore, bm25, bm25_corpus, chunks_metadata, embeddings, index_build_time
logger.info("Loading indexes from disk...")
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL, base_url=OLLAMA_BASE_URL)
vectorstore = Chroma(persist_directory=str(VECTOR_DB_PATH), embedding_function=embeddings)
data = json.loads(BM25_INDEX_PATH.read_text(encoding="utf-8"))
bm25_corpus = data["corpus"]
chunks_metadata = data["metadata"]
bm25 = BM25Okapi([t.lower().split() for t in bm25_corpus])
index_build_time = BM25_INDEX_PATH.stat().st_mtime
logger.info("Indexes loaded.")
# -----------------------------
# Analytics Helper Functions
# -----------------------------
def count_by_language(metadata_list: List[Dict[str, Any]]) -> Dict[str, int]:
"""Count chunks by programming language."""
counts = {}
for meta in metadata_list:
lang = meta.get("language", "unknown")
counts[lang] = counts.get(lang, 0) + 1
return counts
def count_by_type(metadata_list: List[Dict[str, Any]]) -> Dict[str, int]:
"""Count chunks by type (function, class, etc.)."""
counts = {}
for meta in metadata_list:
typ = meta.get("type", "unknown")
counts[typ] = counts.get(typ, 0) + 1
return counts
def calculate_avg_chunk_size(corpus: List[str]) -> float:
"""Calculate average chunk size in characters."""
if not corpus:
return 0.0
return sum(len(text) for text in corpus) / len(corpus)
# -----------------------------
# Hybrid search 80/20 + normalization
# -----------------------------
def hybrid_search(query: str, k: int = 30) -> List[Tuple[str, Dict[str, Any], float]]:
with safe_search():
# Vector results (Chroma returns (doc, score) where lower score might be better depending on backend)
vector_results = vectorstore.similarity_search_with_score(query, k=50)
# Normalize vector scores so that higher is better in [0,1]
vec_values = [s for _, s in vector_results] if vector_results else [1.0]
vmin, vmax = (min(vec_values), max(vec_values)) if vector_results else (0.0, 1.0)
vector_scores = {}
for doc, s in vector_results:
# If Chroma returns distance-like (smaller better), invert; otherwise adapt
# We'll normalize via (vmax - s) / (vmax - vmin + eps)
eps = 1e-12
denom = (vmax - vmin) + eps
norm = (vmax - s) / denom
vector_scores[doc.page_content] = norm
# BM25
tokenized_query = query.lower().split()
bm25_scores = bm25.get_scores(tokenized_query)
top_idxs = sorted(range(len(bm25_scores)), key=lambda i: bm25_scores[i], reverse=True)[:50]
max_b = max((bm25_scores[i] for i in top_idxs), default=1.0)
bm25_doc_scores = {bm25_corpus[i]: (bm25_scores[i] / max_b if max_b > 0 else 0.0) for i in top_idxs}
# Fusion 80/20
all_docs = set(vector_scores.keys()) | set(bm25_doc_scores.keys())
fused = {}
for d in all_docs:
v = vector_scores.get(d, 0.0)
b = bm25_doc_scores.get(d, 0.0)
fused[d] = VECTOR_WEIGHT * v + BM25_WEIGHT * b
# Sort by fused score descending
sorted_items = sorted(fused.items(), key=lambda x: x[1], reverse=True)
# Build accurate mapping from document text to its true metadata
text_to_meta = {text: meta for text, meta in zip(bm25_corpus, chunks_metadata)}
results = []
for doc_text, score in sorted_items[:k]:
meta = text_to_meta.get(doc_text, {})
results.append((doc_text, meta, score))
return results
# -----------------------------
# Enhanced Reranker with Batch Processing
# -----------------------------
def rerank_with_ollama_enhanced(query: str, candidates: List[str], top_k: int = 5) -> List[Tuple[str, float]]:
"""Enhanced Qwen3 Reranker with batch processing and better scoring."""
if not candidates:
return []
# Use domain-specific instruction for code search
instruction = "Given a technical code search query, retrieve relevant code implementations, function definitions, or examples that directly address the query requirements"
batch_scores = []
# Process in smaller batches to avoid timeouts
batch_size = 5
for i in range(0, len(candidates), batch_size):
batch = candidates[i:i + batch_size]
batch_prompts = []
for chunk in batch:
system_prompt = '<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>\n'
user_prompt = (
f'<|im_start|>user\n'
f'<Instruct>: {instruction}\n'
f'<Query>: {query}\n'
f'<Document>: {chunk[:8000]}\n' # Limit document length
f'<|im_end|>\n'
f'<|im_start|>assistant\n<think>\n\n</think>\n\n'
)
batch_prompts.append(system_prompt + user_prompt)
# Score batch
for prompt, chunk in zip(batch_prompts, batch):
try:
resp = requests.post(
f"{OLLAMA_BASE_URL}/api/generate",
json={
"model": RERANKER_MODEL_OLLAMA,
"prompt": prompt,
"max_tokens": 3,
"temperature": 0.0,
"stream": False
},
timeout=20
)
j = resp.json()
response_text = j.get("response", "").strip().lower()
# Enhanced scoring with confidence levels
score = parse_reranker_response(response_text)
batch_scores.append((chunk, score))
except requests.exceptions.Timeout:
logger.warning("Reranker request timeout, assigning default score")
batch_scores.append((chunk, 0.0))
except Exception as e:
logger.warning(f"Rerank call failed: {e}")
batch_scores.append((chunk, 0.0))
# Sort by score descending and return top_k
batch_scores.sort(key=lambda x: x[1], reverse=True)
return batch_scores[:top_k]
def parse_reranker_response(response_text: str) -> float:
"""Parse Qwen3 Reranker response and convert to confidence score."""
response_lower = response_text.strip().lower()
# Exact matches from the model
if response_lower == "yes":
return 1.0
elif response_lower == "no":
return 0.0
# Handle variations and partial matches
yes_indicators = ["yes", "relevant", "correct", "matches", "appropriate", "suitable"]
no_indicators = ["no", "irrelevant", "incorrect", "unrelated", "inappropriate"]
yes_count = sum(1 for indicator in yes_indicators if indicator in response_lower)
no_count = sum(1 for indicator in no_indicators if indicator in response_lower)
if yes_count > no_count:
return 0.8 # Likely relevant but not confident
elif no_count > yes_count:
return 0.2 # Likely irrelevant but not confident
else:
# Ambiguous response
return 0.5
# -----------------------------
# Comment-aware reference search (ripgrep)
# -----------------------------
def get_context_with_comments(file_path: Path, match_line: int, max_context_lines: int = 5) -> str:
try:
lines = file_path.read_text(encoding="utf-8").splitlines()
except Exception:
return ""
idx = match_line - 1
context_lines = []
look = idx - 1
collected = 0
while look >= 0 and collected < max_context_lines:
line = lines[look]
stripped = line.strip()
if stripped.startswith(("#", "//", "--")) or stripped == "":
context_lines.insert(0, line)
look -= 1
collected += 1
else:
break
context_lines.append(lines[idx] if 0 <= idx < len(lines) else "")
return "\n".join(context_lines)
def search_symbol(symbol: str, root_dir: Path, top_k: int = 20) -> List[Dict[str, Any]]:
matches = []
rg_cmd = [
"rg",
"-n",
"-i",
"-C", "3",
"--no-ignore",
"--no-heading",
"--color", "never",
"--glob", "!node_modules",
"--glob", "!__pycache__",
"--glob", "!.git",
]
for ext in SUPPORTED_EXTENSIONS:
rg_cmd.extend(["--glob", f"*{ext}"])
rg_cmd.extend(["--", symbol, "."])
try:
proc = subprocess.run(
rg_cmd,
cwd=root_dir,
capture_output=True,
text=True,
timeout=30
)
if proc.returncode not in (0, 1):
logger.warning(f"ripgrep returned code {proc.returncode}: {proc.stderr}")
return []
for line in proc.stdout.splitlines():
if not line.strip() or line.startswith("--"):
continue
parts = line.split(":", 2)
if len(parts) < 3:
continue
file_rel, lineno_str, snippet = parts
try:
lineno = int(lineno_str)
except ValueError:
continue
matches.append({
"file": file_rel,
"line": lineno,
"context": snippet.strip()
})
except subprocess.TimeoutExpired:
logger.warning("ripgrep timed out")
except Exception as e:
logger.warning(f"ripgrep failed: {e}")
return matches
def _extract_imports(source: str) -> list[str]:
"""Return a list of fully‑qualified imports found in `source`."""
return [m.group(1).strip() for m in re.finditer(r'^\s*use\s+([^;]+);', source, re.MULTILINE)]
def _extract_calls(source: str) -> list[str]:
"""Return a list of called function names (bare name only)."""
return [m.group(1).split("::")[-1] for m in re.finditer(r'([A-Za-z_][A-Za-z0-9_:]*)\s*\(', source, re.MULTILINE)]
# -----------------------------
# Additional tool: read file lines with context
# -----------------------------
def read_file_lines(path: str, start: int = 1, end: Optional[int] = None) -> str:
"""
Read file lines with intelligent context inclusion.
- Includes requested lines
- Expands to include full function/class definition if within one
- Includes leading comments (docstrings, inline comments)
- Shows function signature, return type, decorators
"""
# Handle both absolute and relative paths
if path.startswith('/'):
p = Path(path)
else:
p = CODEBASE_PATH / Path(path)
# Security: Validate path is within codebase
try:
resolved = p.resolve()
codebase_resolved = CODEBASE_PATH.resolve()
if not resolved.is_relative_to(codebase_resolved):
return f"Error: Path '{path}' is outside the codebase directory"
except (ValueError, OSError) as e:
return f"Error: Invalid path '{path}': {e}"
try:
content = p.read_text(encoding="utf-8")
lines = content.splitlines()
if end is None:
end = start + 50 # reasonable default
# Clamp to valid range
start_i = max(1, start)
end_i = min(len(lines), end)
# Detect language
language = detect_language(p)
# Build context-aware line range
context_start, context_end, context_info = find_context_boundaries(
lines, start_i, end_i, language
)
# Extract the context lines
context_lines = lines[context_start - 1:context_end]
# Build output with metadata
output = []
output.append(f"File: {path}")
output.append(f"Language: {language}")
output.append(f"Requested: lines {start_i}-{end_i}")
output.append(f"Showing: lines {context_start}-{context_end} (with context)")
if context_info:
output.append(f"Context: {context_info}")
output.append("\n" + "="*60)
# Add line numbers
for i, line in enumerate(context_lines, start=context_start):
# Highlight the originally requested range
marker = ">>> " if start_i <= i <= end_i else " "
output.append(f"{marker}{i:4d} | {line}")
output.append("="*60)
return "\n".join(output)
except Exception as e:
return f"Failed to read {path}: {e}"
def find_context_boundaries(lines: List[str], start: int, end: int, language: str) -> Tuple[int, int, str]:
"""
Find intelligent context boundaries around the requested line range.
Returns: (context_start_line, context_end_line, description)
"""
# Convert to 0-based indexing for processing
start_idx = start - 1
end_idx = end - 1
context_info = []
# Expand upward to include leading comments and function/class headers
context_start = start
# Step 1: Include leading comments above start
comment_start = find_leading_comments_start(lines, start_idx, language)
if comment_start < start_idx:
context_start = comment_start + 1
context_info.append("leading comments")
# Step 2: Check if we're inside a function/class and include its signature
func_start, func_end, func_name, func_type = find_enclosing_function_or_class(
lines, start_idx, end_idx, language
)
if func_start is not None:
context_start = min(context_start, func_start + 1)
if func_end is not None and func_end > end_idx:
context_end = func_end + 1
else:
context_end = end
if func_name:
context_info.append(f"{func_type} '{func_name}'")
else:
context_end = end
# Step 3: Include decorators (Python) or annotations (Java)
if language == "python":
decorator_start = find_decorators_start(lines, context_start - 1)
if decorator_start < context_start - 1:
context_start = decorator_start + 1
context_info.append("decorators")
info_str = " + ".join(context_info) if context_info else "no additional context"
return context_start, context_end, info_str
def find_leading_comments_start(lines: List[str], line_idx: int, language: str) -> int:
"""Find the start of leading comments above the given line."""
comment_patterns = {
"python": ("#", '"""', "'''"),
"javascript": ("//", "/*"),
"typescript": ("//", "/*"),
"go": ("//", "/*"),
"java": ("//", "/*"),
"sql": ("--", "/*"),
"rust": ("//", "/*"),
}
patterns = comment_patterns.get(language, ("#", "//", "/*", "--"))
start_idx = line_idx
i = line_idx - 1
# Track if we're in a block comment
in_block = False
while i >= 0:
line = lines[i].strip()
# Empty lines are OK if we already have comments
if not line:
if i < line_idx - 1: # Allow one blank line between comments
i -= 1
continue
else:
break
# Check for block comment end (we're going backwards)
if "*/" in line:
in_block = True
start_idx = i
i -= 1
continue
# Check for block comment start
if in_block and ("/*" in line or "/**" in line):
start_idx = i
in_block = False
i -= 1
continue
# Check for single-line comments
is_comment = any(line.startswith(p) for p in patterns)
if is_comment or in_block:
start_idx = i
i -= 1
else:
break
return start_idx
def find_decorators_start(lines: List[str], line_idx: int) -> int:
"""Find Python decorators above the given line."""
i = line_idx - 1
while i >= 0:
line = lines[i].strip()
if line.startswith("@"):
i -= 1
elif not line: # Allow blank lines
i -= 1
else:
return i + 1
return 0
def find_enclosing_function_or_class(lines: List[str], start_idx: int, end_idx: int, language: str) -> Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]:
"""
Find the function or class that encloses the given line range.
Returns: (start_line_idx, end_line_idx, name, type) or (None, None, None, None)
"""
if language == "python":
return find_python_function_or_class(lines, start_idx, end_idx)
elif language in ("javascript", "typescript"):
return find_js_function(lines, start_idx, end_idx)
elif language == "go":
return find_go_function(lines, start_idx, end_idx)
elif language == "java":
return find_java_method_or_class(lines, start_idx, end_idx)
elif language == "sql":
return find_sql_function_or_block(lines, start_idx, end_idx)
return None, None, None, None
def find_python_function_or_class(lines: List[str], start_idx: int, end_idx: int) -> Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]:
"""Find enclosing Python function or class."""
# Look backwards for def/class with proper indentation
target_indent = None
for i in range(start_idx, -1, -1):
line = lines[i]
stripped = line.lstrip()
if stripped.startswith(("def ", "class ", "async def ")):
indent = len(line) - len(stripped)
# Check if this could be our enclosing scope
if target_indent is None or indent < target_indent:
# Extract name
match = re.match(r'(?:async\s+)?(?:def|class)\s+(\w+)', stripped)
if match:
name = match.group(1)
func_type = "class" if stripped.startswith("class") else "function"
# Find the end by looking for next same-or-lower indent non-empty line
end_i = find_python_block_end(lines, i, indent)
return i, end_i, name, func_type
return None, None, None, None
def find_python_block_end(lines: List[str], start_idx: int, base_indent: int) -> Optional[int]:
"""Find the end of a Python block based on indentation."""
for i in range(start_idx + 1, len(lines)):
line = lines[i]
stripped = line.lstrip()
if not stripped: # Skip empty lines
continue
indent = len(line) - len(stripped)
if indent <= base_indent and stripped and not stripped.startswith("#"):
return i - 1
return len(lines) - 1
def find_js_function(lines: List[str], start_idx: int, end_idx: int) -> Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]:
"""Find enclosing JavaScript/TypeScript function."""
for i in range(start_idx, -1, -1):
line = lines[i].strip()
# Match various JS function patterns
patterns = [
r'function\s+(\w+)',
r'(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s+)?\(',
r'(\w+)\s*\([^)]*\)\s*{', # Arrow functions
r'async\s+function\s+(\w+)',
]
for pattern in patterns:
match = re.search(pattern, line)
if match:
name = match.group(1)
end_i = find_brace_block_end(lines, i)
return i, end_i, name, "function"
return None, None, None, None
def find_go_function(lines: List[str], start_idx: int, end_idx: int) -> Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]:
"""Find enclosing Go function."""
for i in range(start_idx, -1, -1):
line = lines[i].strip()
# Match Go function: func (receiver) name(params) returnType {
match = re.match(r'func\s+(?:\([^)]+\)\s+)?(\w+)', line)
if match:
name = match.group(1)
end_i = find_brace_block_end(lines, i)
return i, end_i, name, "function"
return None, None, None, None
def find_java_method_or_class(lines: List[str], start_idx: int, end_idx: int) -> Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]:
"""Find enclosing Java method or class."""
for i in range(start_idx, -1, -1):
line = lines[i].strip()
# Match class
class_match = re.match(r'(?:public|private|protected)?\s*(?:static)?\s*class\s+(\w+)', line)
if class_match:
name = class_match.group(1)
end_i = find_brace_block_end(lines, i)
return i, end_i, name, "class"
# Match method
method_match = re.match(r'(?:public|private|protected)?\s*(?:static)?\s*(?:\w+(?:<[^>]+>)?)\s+(\w+)\s*\(', line)
if method_match:
name = method_match.group(1)
end_i = find_brace_block_end(lines, i)
return i, end_i, name, "method"
return None, None, None, None
def find_sql_function_or_block(lines: List[str], start_idx: int, end_idx: int) -> Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]:
"""Find enclosing SQL function or procedure."""
for i in range(start_idx, -1, -1):
line = lines[i].strip().upper()
if line.startswith(("CREATE FUNCTION", "CREATE OR REPLACE FUNCTION", "CREATE PROCEDURE")):
# Extract name
match = re.search(r'(?:FUNCTION|PROCEDURE)\s+(\w+)', line, re.IGNORECASE)
if match:
name = match.group(1)
# SQL functions typically end with $$ or END
end_i = find_sql_function_end(lines, i)
func_type = "procedure" if "PROCEDURE" in line else "function"
return i, end_i, name, func_type
return None, None, None, None
def find_sql_function_end(lines: List[str], start_idx: int) -> Optional[int]:
"""Find end of SQL function (looks for $$ or END;)."""
for i in range(start_idx + 1, len(lines)):
line = lines[i].strip().upper()
if "$$" in line or line.startswith("END;") or line == "END":
return i
return len(lines) - 1
def find_brace_block_end(lines: List[str], start_idx: int) -> Optional[int]:
"""Find the end of a brace-delimited block {}."""
brace_count = 0
started = False
for i in range(start_idx, len(lines)):
line = lines[i]
for char in line:
if char == '{':
brace_count += 1
started = True
elif char == '}':
brace_count -= 1
if started and brace_count == 0:
return i
return len(lines) - 1
# -----------------------------
# MCP Tools
# -----------------------------
@mcp.tool()
def health_check() -> str:
"""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,"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)"
# Load graph to check if it exists - FIX: Convert Path to string
graph = LocalGraph(root_path=str(CODEBASE_PATH)) # Convert Path to str
graph_loaded = graph.load()
status = {
"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,
"graph_enabled": True
},
"Tools": {
"NOTE": "THESE TOOLS ARE RESTRICTED BY .gitignore AS WELL AS .ragignore",
"semantic_rag_search": "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"
}
}
if bm25_corpus:
status["statistics"] = {
"total_chunks": len(bm25_corpus),
"by_language": count_by_language(chunks_metadata),
"by_type": count_by_type(chunks_metadata),
"avg_chunk_size_chars": round(calculate_avg_chunk_size(bm25_corpus), 1),
"index_age_hours": round((time.time() - index_build_time) / 3600, 1) if index_build_time > 0 else None
}
# Add graph statistics - FIX: Use graph.graph.nodes instead of graph.nodes
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."""
# FIX: Use graph.graph.nodes and graph.graph.edges instead of graph.nodes/graph.edges
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): # FIX: graph.graph.nodes
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): # FIX: graph.graph.edges
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): # FIX: graph.graph.nodes
lang = data.get('lang', 'unknown') # Changed from 'language' to 'lang' to match our new attribute
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'] # FIX: graph.graph.nodes
# Relationship density
avg_edges_per_node = edges / nodes if nodes > 0 else 0
# Most connected nodes (hubs)
degree_centrality = dict(graph.graph.degree()) # FIX: 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] # FIX: graph.graph.nodes
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 semantic_rag_search(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.
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: semantic_rag_search("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."
# FIX: Convert Path to string
graph = LocalGraph(root_path=str(CODEBASE_PATH))
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))
if ENABLE_RERANK and rerank:
try:
candidates = [t for t, m, s in hy]
rr = rerank_with_ollama_enhanced(query, candidates[:RERANK_TOP_N], top_k=top_k)
# Build enhanced results with graph context
results = []
for i, (chunk_text, score) in enumerate(rr, 1):
meta = next((m for t, m, s in hy if t == chunk_text), {})
results.append(_build_enhanced_result(chunk_text, meta, score, graph))
return _format_enhanced_results(results, query, "search")
except Exception as e:
logger.warning(f"Rerank step failed: {e}")
# Fall through to non-reranked results
# Non-reranked results with graph context
results = []
for i, (text, meta, score) in enumerate(hy[:top_k], 1):
results.append(_build_enhanced_result(text, meta, score, graph))
return _format_enhanced_results(results, query, "search")
except RuntimeError as e:
return f"Error: {e}"
except Exception as e:
logger.exception("Search failed")
return f"Search failed: {e}"
def _build_enhanced_result(chunk_text: str, meta: dict, score: float, graph: Optional[LocalGraph] = None) -> dict:
"""Build enhanced result with graph context and important metadata."""
# 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 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,
'type': entity_type,
'language': language,
'name': entity_name,
'score': score,
'content': clean_content,
'docstring': docstring,
'graph_context': graph_context
}
def _get_result_graph_context(meta: dict, graph: Optional[LocalGraph] = None) -> dict:
"""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:
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}"
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
rel_str = f"{rel_type}:{neighbor_name}@{neighbor_file}"
context['cross_file_deps'].append(rel_str)
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):
pred_data = graph.graph.nodes[predecessor]
pred_file = pred_data.get('file', '')
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:
file_neighbors = list(graph.neighbors(file_node_id))
if file_neighbors:
context['file_entities'] = []
for neighbor_id, neighbor_data in file_neighbors[:5]:
neighbor_type = neighbor_data.get('type', '')
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
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)
# Python style docstrings
elif '"""' in line or "'''" in line:
if not in_docstring:
in_docstring = True
else:
in_docstring = False
break
elif in_docstring:
doc_lines.append(stripped)
# Specific doc patterns in the metadata section
elif stripped.startswith('Doc:') and len(stripped) > 4:
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 ""
def _format_enhanced_results(results: List[Dict], query: str, result_type: str) -> str:
"""Format enhanced results with dense LLM-optimized context."""
if not results:
return f"# {result_type.title()}: {query}\nNo results found.\n"
lines = [
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', ''),
result.get('line', ''),
result.get('type', 'code'),
result.get('name', ''),
f"{result.get('score', 0.0):.2f}",
result.get('content', '')[:500], # More content tokens
context_str
]
# Efficient escaping
escaped_row = []
for field in row:
field_str = str(field)
if ',' in field_str or ';' in field_str:
escaped = field_str.replace('"', '\\"')
escaped_row.append(f'"{escaped}"')
else:
escaped_row.append(field_str)
lines.append(" " + ",".join(escaped_row))
return "\n".join(lines)
def _extract_clean_content(chunk_text: str) -> str:
"""Extract clean content from chunk text by removing duplicate metadata."""
lines = chunk_text.split('\n')
content_lines = []
# Skip the first few metadata lines (File:, Type:, Name:, etc.)
skip_metadata = True
for line in lines:
stripped = line.strip()
if skip_metadata:
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:']):
skip_metadata = False
content_lines.append(stripped)
else:
content_lines.append(stripped)
# Join and clean up
content = ' '.join(content_lines)
content = re.sub(r'\s+', ' ', content) # 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.
When to use - After identifying a symbol in a search hit, or when you need every use of a class/method.
When not to use - When you need context-rich code; use read_file_lines_tool after getting the locations.
Example - find_code_references("AuthService") → list of references."""
matches = search_symbol(symbol, CODEBASE_PATH)
if not matches:
return f"# References: {symbol}\nNo references found.\n"
# Build clean TOON format directly
results = []
for i, m in enumerate(matches[:top_k], 1):
results.append({
'file': m['file'],
'line': m['line'],
'context': m['context']
})
return _format_reference_results(results, symbol)
def _format_reference_results(results: List[Dict], symbol: str) -> str:
"""Format reference results in clean TOON format."""
if not results:
return f"# References: {symbol}\nNo references found.\n"
lines = [f"# References: {symbol}", f"references[{len(results)}]{{file,line,context}}:"]
for result in results:
row = [
result.get('file', ''),
result.get('line', ''),
result.get('context', '')[:300] # Reasonable limit
]
# Escape fields if needed
escaped_row = []
for field in row:
field_str = str(field)
if any(char in field_str for char in [',', '"', '\n', '\r']):
escaped = field_str.replace('"', '\\"')
escaped_row.append(f'"{escaped}"')
else:
escaped_row.append(field_str)
lines.append(" " + ",".join(escaped_row))
return "\n".join(lines)
@mcp.tool()
def read_file_lines_tool(path: str, start: int = 1, end: Optional[int] = None) -> str:
"""What it does - Reads a requested line range and auto-expands to include surrounding context (function body, docstring, comments).
When to use - To inspect a specific hit from find_code_references or a snippet from semantic_rag_search.
When not to use - For browsing the entire file or unrelated sections; use semantic_rag_search first to pinpoint relevant parts.
Example - read_file_lines_tool("src/auth/login.py", start=45, end=60) → full function with context.
Note: This tool treats / as '/home/popertots/Crussell/' so adjust paths accordingly"""
return EnhancedToon.file_content_results(read_file_lines(path, start, end), path)
@mcp.tool()
def init_repo(git_url: str) -> str:
"""What it does - Initialise the RAG stack with a brand-new Git repository."""
import shutil
from pathlib import Path
import subprocess
# ----- 1. Always wipe and recreate working_repo ---------------------
if CODEBASE_PATH.exists():
shutil.rmtree(CODEBASE_PATH)
logger.info(f"Wiped existing working repo: {CODEBASE_PATH}")
CODEBASE_PATH.mkdir(exist_ok=True)
# ----- 2. Clone fresh -----------------------------------------------
try:
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 {git_url}: {e.stderr.strip()}"
# ----- 3. Reset ALL global state ------------------------------------
global vectorstore, bm25, bm25_corpus, chunks_metadata, index_build_time
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:
if artefact_path.is_dir():
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 ----------------------------------------
try:
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 --------------------------------------------
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()
# 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:
"""What it does - Re-creates all embeddings, BM25, and metadata indexes after major code changes.
When to use - When the codebase has been pulled or refactored and you suspect stale search results, and only after explicit instruction to do so.
When not to use - On every query; it's expensive and unnecessary if the index is already up-to-date.
Example - rebuild_index() → '✅ Index rebuilt (13.2 s); 11,234 chunks'"""
with _startup_lock:
try:
# Clear LRU caches
parse_python_file_cached.cache_clear()
parse_sql_file_cached.cache_clear()
parse_go_file_cached.cache_clear()
parse_java_file_cached.cache_clear()
parse_svelte_file_cached.cache_clear()
parse_shell_file_cached.cache_clear()
LocalGraph(root_path=CODEBASE_PATH).clear()
t0 = time.time()
build_indexes()
dt = time.time() - t0
return f"✅ Index rebuilt successfully! (took {dt:.1f}s)\n\nStatistics:\n" + json.dumps({
"total_chunks": len(bm25_corpus),
"by_language": count_by_language(chunks_metadata),
"by_type": count_by_type(chunks_metadata)
}, indent=2)
except Exception as e:
logger.exception("Index rebuild failed")
return f"❌ Rebuild failed: {e}"
# -----------------------------
# Startup
# -----------------------------
def signal_handler(sig, frame):
"""Handle graceful shutdown on Ctrl+C."""
logger.info("\n=== Shutdown signal received ===")
# Clear any LRU caches to release memory
try:
parse_python_file_cached.cache_clear()
parse_sql_file_cached.cache_clear()
parse_go_file_cached.cache_clear()
parse_java_file_cached.cache_clear()
parse_svelte_file_cached.cache_clear()
parse_shell_file_cached.cache_clear()
logger.info("✓ Caches cleared")
except Exception as e:
logger.warning(f"Cache clearing error: {e}")
# Note: Chroma auto-persists, no explicit close needed
logger.info("✓ Indexes are already persisted to disk")
logger.info("Goodbye!\n")
# Force exit immediately to avoid thread hang
os._exit(0)
if __name__ == "__main__":
# Register signal handlers
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
try:
logger.info("="*60)
logger.info("🚀 MCP RAG Server is ready!")
logger.info("Press Ctrl+C to stop")
# Just run in stdio mode by default
mcp.run(transport='stdio')
except KeyboardInterrupt:
signal_handler(signal.SIGINT, None)
except Exception as e:
logger.exception("Fatal error during startup")
os._exit(1)
Reference in New Issue View Git Blame Copy Permalink