2413 lines
91 KiB
Python
2413 lines
91 KiB
Python
import os
|
||
import signal
|
||
import sys
|
||
import json
|
||
import ast
|
||
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
|
||
from dataclasses import dataclass
|
||
import time
|
||
import hashlib
|
||
|
||
import pathspec
|
||
import requests
|
||
import numpy as np
|
||
|
||
from mcp.server.fastmcp import FastMCP
|
||
import uvicorn
|
||
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
|
||
|
||
os.environ["ANONYMIZED_TELEMETRY"] = "False"
|
||
|
||
# -----------------------------
|
||
# Configuration
|
||
# -----------------------------
|
||
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
|
||
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(os.environ.get("CODEBASE_PATH", ".")) # change as needed
|
||
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.2, # Class definitions
|
||
"python_function": 1.15, # Function definitions
|
||
|
||
# Go weights
|
||
"go_type": 1.2, # Type definitions
|
||
"go_function": 1.15, # Function definitions
|
||
|
||
# Java weights
|
||
"java_class": 1.2, # Class definitions
|
||
"java_method": 1.15, # Method definitions
|
||
|
||
# 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()
|
||
|
||
# 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
|
||
elif language == "python":
|
||
if chunk_type == "classdef":
|
||
return CONTEXTUAL_WEIGHTS.get("python_class", 1.0)
|
||
elif chunk_type in ("functiondef", "asyncfunctiondef"):
|
||
return CONTEXTUAL_WEIGHTS.get("python_function", 1.0)
|
||
|
||
# Go-specific weights
|
||
elif language == "go":
|
||
if chunk_type in ("type", "struct", "interface"):
|
||
return CONTEXTUAL_WEIGHTS.get("go_type", 1.0)
|
||
elif chunk_type in ("func", "method"):
|
||
return CONTEXTUAL_WEIGHTS.get("go_function", 1.0)
|
||
|
||
# Java-specific weights
|
||
elif language == "java":
|
||
if chunk_type == "class":
|
||
return CONTEXTUAL_WEIGHTS.get("java_class", 1.0)
|
||
elif chunk_type == "method":
|
||
return CONTEXTUAL_WEIGHTS.get("java_method", 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 = []
|
||
|
||
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. Returns tuple for hashability."""
|
||
filepath = Path(filepath_str)
|
||
chunks = []
|
||
try:
|
||
code = filepath.read_text(encoding="utf-8")
|
||
tree = ast.parse(code)
|
||
lines = code.splitlines(keepends=True)
|
||
for node in ast.walk(tree):
|
||
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])
|
||
chunks.append({
|
||
"text": f"File: {filepath}\nType: {type(node).__name__}\nName: {getattr(node,'name', '<anon>')}\nDocstring: {docstring}\n\nCode:\n{chunk_code}",
|
||
"metadata": {
|
||
"file": str(filepath),
|
||
"name": getattr(node, "name", ""),
|
||
"type": type(node).__name__,
|
||
"line": node.lineno,
|
||
"language": "python",
|
||
}
|
||
})
|
||
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 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)
|
||
"""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:
|
||
# find position of this statement in the original SQL (using running search_pos)
|
||
start_idx = sql_code.find(stmt, search_pos)
|
||
if start_idx == -1:
|
||
# fallback: try from beginning
|
||
start_idx = sql_code.find(stmt)
|
||
if start_idx == -1:
|
||
# cannot locate, approximate line number as last line
|
||
start_line = len(lines)
|
||
else:
|
||
start_line = sql_code[:start_idx].count("\n") + 1
|
||
search_pos = start_idx + len(stmt)
|
||
|
||
# include leading comments above this statement
|
||
leading_comments = include_leading_comments(lines, start_line)
|
||
|
||
# Trim and normalize statement
|
||
stmt_clean = stmt.strip()
|
||
stmt_text_for_embed = (leading_comments + "\n\n" + stmt_clean) if leading_comments else stmt_clean
|
||
|
||
# Try to parse with sqlglot to extract tables / functions / ctes metadata
|
||
meta: Dict[str, Any] = {
|
||
"file": str(filepath),
|
||
"name": None,
|
||
"type": None,
|
||
"line": start_line,
|
||
"language": "sql"
|
||
}
|
||
try:
|
||
parsed = sqlglot.parse_one(stmt_clean, read="postgres")
|
||
# stmt type
|
||
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"
|
||
|
||
# table names (if any)
|
||
tables = [t.this for t in parsed.find_all(sqlglot.exp.Table)]
|
||
meta["tables"] = [t for t in tables if isinstance(t, str)]
|
||
# ctes
|
||
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"] = ctes
|
||
|
||
# functions called in the statement (names)
|
||
funcs = []
|
||
for f in parsed.find_all(sqlglot.exp.Func):
|
||
try:
|
||
name = f.name
|
||
if name:
|
||
funcs.append(name)
|
||
except Exception:
|
||
pass
|
||
meta["functions"] = funcs
|
||
|
||
# choose a sensible name for the metadata
|
||
if meta.get("tables"):
|
||
meta["name"] = meta["tables"][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\nComments:\n{leading_comments}\n\nSQL:\n{stmt_clean}" if leading_comments else f"File: {filepath}\nType: {meta.get('type')}\nName: {meta.get('name')}\n\nSQL:\n{stmt_clean}"
|
||
|
||
chunks.append({
|
||
"text": chunk_text,
|
||
"metadata": meta
|
||
})
|
||
|
||
continue
|
||
except Exception:
|
||
# sqlglot failed on this statement — fall back to heuristics (still keep comments)
|
||
pass
|
||
|
||
# Fallback: simple heuristics if sqlglot fails
|
||
# Determine an approximate type by the first token
|
||
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\nComments:\n{leading_comments}\n\nSQL:\n{stmt_clean}" if leading_comments else f"File: {filepath}\nType: {meta['type']}\nName: {meta.get('name')}\n\nSQL:\n{stmt_clean}"
|
||
|
||
chunks.append({
|
||
"text": chunk_text,
|
||
"metadata": meta
|
||
})
|
||
pass
|
||
|
||
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. Returns tuple for hashability."""
|
||
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", "")
|
||
|
||
# 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"
|
||
chunk_text += f"Code:\n{chunk_code}"
|
||
|
||
chunks.append({
|
||
"text": chunk_text,
|
||
"metadata": {
|
||
"file": str(filepath),
|
||
"name": name,
|
||
"type": typ,
|
||
"receiver": receiver,
|
||
"line": start + 1,
|
||
"language": "go",
|
||
}
|
||
})
|
||
|
||
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 (same as before)
|
||
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)
|
||
|
||
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}"
|
||
|
||
chunks.append({
|
||
"text": chunk_text,
|
||
"metadata": {
|
||
"file": str(filepath),
|
||
"name": func_name,
|
||
"type": "function",
|
||
"line": start_line + 1,
|
||
"language": "go",
|
||
}
|
||
})
|
||
|
||
# Structs
|
||
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)
|
||
|
||
chunk_text = f"File: {filepath}\nType: struct\nName: {struct_name}\n"
|
||
if comments:
|
||
chunk_text += f"Comments:\n{comments}\n\n"
|
||
chunk_text += f"Code:\n{struct_code}"
|
||
|
||
chunks.append({
|
||
"text": chunk_text,
|
||
"metadata": {
|
||
"file": str(filepath),
|
||
"name": struct_name,
|
||
"type": "struct",
|
||
"line": start_line + 1,
|
||
"language": "go",
|
||
}
|
||
})
|
||
|
||
# Interfaces
|
||
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)
|
||
|
||
chunk_text = f"File: {filepath}\nType: interface\nName: {interface_name}\n"
|
||
if comments:
|
||
chunk_text += f"Comments:\n{comments}\n\n"
|
||
chunk_text += f"Code:\n{interface_code}"
|
||
|
||
chunks.append({
|
||
"text": chunk_text,
|
||
"metadata": {
|
||
"file": str(filepath),
|
||
"name": interface_name,
|
||
"type": "interface",
|
||
"line": start_line + 1,
|
||
"language": "go",
|
||
}
|
||
})
|
||
|
||
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, file_hash: 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)
|
||
chunks.append({
|
||
"text": f"File: {filepath}\nType: Class\nName: {node.name}\nDoc:\n{doc}\nCode:\n{chunk_text}",
|
||
"metadata": {
|
||
"file": str(filepath),
|
||
"name": node.name,
|
||
"type": "class",
|
||
"line": start_line + 1,
|
||
"language": "java"
|
||
}
|
||
})
|
||
|
||
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
|
||
chunks.append({
|
||
"text": f"File: {filepath}\nType: Method\nClass: {class_name or 'unknown'}\nName: {node.name}\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"
|
||
}
|
||
})
|
||
|
||
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)
|
||
for declarator in node.declarators:
|
||
chunks.append({
|
||
"text": f"File: {filepath}\nType: Field\nName: {declarator.name}\nDoc:\n{doc}\nCode:\n{chunk_text}",
|
||
"metadata": {
|
||
"file": str(filepath),
|
||
"name": declarator.name,
|
||
"type": "field",
|
||
"line": start_line + 1,
|
||
"language": "java"
|
||
}
|
||
})
|
||
|
||
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"
|
||
}
|
||
})
|
||
|
||
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."""
|
||
file_hash = get_file_hash(filepath)
|
||
cached_result = parse_java_file_cached(str(filepath), file_hash)
|
||
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. Returns tuple for hashability."""
|
||
filepath = Path(filepath_str)
|
||
chunks = []
|
||
try:
|
||
text = filepath.read_text(encoding="utf-8")
|
||
lines = text.splitlines(keepends=True)
|
||
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
|
||
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
|
||
chunk_text = f"File: {filepath}\nType: {d.get('type')}\nName: {d.get('name')}\nDoc: {doc}\nCode:\n{chunk_code}"
|
||
chunks.append({
|
||
"text": chunk_text,
|
||
"metadata": {
|
||
"file": str(filepath),
|
||
"name": d.get("name"),
|
||
"type": d.get("type"),
|
||
"line": abs_start,
|
||
"language": "typescript",
|
||
"block_type": "script"
|
||
}
|
||
})
|
||
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"))
|
||
finally:
|
||
tmp.unlink(missing_ok=True)
|
||
else:
|
||
chunks.append(_make_script_chunk(filepath, script_content, script_start_line, "script"))
|
||
idx += 1
|
||
|
||
# Style blocks
|
||
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
|
||
chunks.append({
|
||
"text": f"File: {filepath}\nType: style\n{style_content}",
|
||
"metadata": {
|
||
"file": str(filepath),
|
||
"name": f"style_block_{idx}",
|
||
"type": "style",
|
||
"line": style_start_line,
|
||
"language": "css",
|
||
"block_type": "style"
|
||
}
|
||
})
|
||
idx += 1
|
||
|
||
# Markup
|
||
markup = re.sub(r"<script[^>]*>.*?</script>", "", text, flags=re.DOTALL)
|
||
markup = re.sub(r"<style[^>]*>.*?</style>", "", markup, flags=re.DOTALL).strip()
|
||
if markup:
|
||
chunks.append({
|
||
"text": f"File: {filepath}\nType: markup\n{markup}",
|
||
"metadata": {
|
||
"file": str(filepath),
|
||
"name": "markup",
|
||
"type": "markup",
|
||
"line": 1,
|
||
"language": "svelte",
|
||
"block_type": "markup"
|
||
}
|
||
})
|
||
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 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]
|
||
|
||
|
||
def _make_script_chunk(filepath: Path, content: str, start_line: int, block_type: str = "script") -> Dict:
|
||
return {
|
||
"text": f"File: {filepath}\nType: {block_type}\n{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
|
||
}
|
||
}
|
||
|
||
# -----------------------------
|
||
# 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
|
||
|
||
@lru_cache(maxsize=1000)
|
||
def parse_shell_file_cached(filepath_str: str, file_hash: 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."""
|
||
file_hash = get_file_hash(filepath)
|
||
cached_result = parse_shell_file_cached(str(filepath), file_hash)
|
||
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)
|
||
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
|
||
|
||
# -----------------------------
|
||
# Build / load indexes with contextual weighting
|
||
# -----------------------------
|
||
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()
|
||
|
||
texts, metadatas = index_codebase(CODEBASE_PATH)
|
||
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, weighted_embs):
|
||
self.weighted_embs = weighted_embs
|
||
self.idx = 0
|
||
|
||
def embed_documents(self, texts):
|
||
# Return the pre-computed weighted embeddings
|
||
start = self.idx
|
||
end = start + len(texts)
|
||
result = self.weighted_embs[start:end]
|
||
self.idx = end
|
||
return result
|
||
|
||
def batch_embed_documents(texts: List[str]) -> List[List[float]]:
|
||
# Embed documents in batches
|
||
all_embeddings = []
|
||
for i in range(0, len(texts), EMBEDDING_BATCH_SIZE):
|
||
batch = texts[i:i + EMBEDDING_BATCH_SIZE]
|
||
embeddings = embeddings.embed_documents(batch)
|
||
all_embeddings.extend(embeddings)
|
||
return all_embeddings
|
||
|
||
def embed_query(self, text):
|
||
# For queries, use the base embedding model (no weighting)
|
||
return embeddings.embed_query(text)
|
||
|
||
weighted_emb_func = WeightedEmbeddingFunction(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")
|
||
|
||
# -----------------------------
|
||
# 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
|
||
|
||
# -----------------------------
|
||
# 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")
|
||
|
||
# Limit expansion to reasonable bounds (max 100 lines of context)
|
||
max_context = 100
|
||
if context_end - context_start > max_context:
|
||
context_end = context_start + max_context
|
||
context_info.append("truncated to 100 lines")
|
||
|
||
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).
|
||
When to use – Before any other query, to confirm the index is current.
|
||
When not to use – After you already know the server is healthy; it adds no value.
|
||
Example – health_check() → { "status":"ready","total_chunks":11234,"index_age_hours":4.7 }"""
|
||
try:
|
||
with _startup_lock:
|
||
status = {
|
||
"status": "ready" if vectorstore is not None else "initializing",
|
||
"codebase": str(CODEBASE_PATH),
|
||
"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
|
||
},
|
||
"Tools": {
|
||
"NOTE": "THESE TOOLS ARE RESTRICTED BY .gitignore AS WELL AS .ragignore",
|
||
"search_codebase": search_codebase.__doc__,
|
||
"find_code_references": find_code_references.__doc__,
|
||
"read_file_lines_tool": read_file_lines_tool.__doc__
|
||
}
|
||
}
|
||
|
||
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
|
||
}
|
||
|
||
return json.dumps(status, indent=2)
|
||
except Exception as e:
|
||
return json.dumps({"status": "error", "message": str(e)}, indent=2)
|
||
|
||
@mcp.tool()
|
||
def search_codebase(query: str, top_k: int = 5, rerank: bool = True) -> str:
|
||
"""What it does – Hybrid RAG search that returns semantically ranked snippets with file/line info from natural language queries.
|
||
When to use – For open‑ended questions or unknown patterns (“how does auth work?”). Natural language and verbose queries ONLY.
|
||
When not to use – For a single symbol or exact phrase; use find_code_references then read instead to avoid a large result set.
|
||
Rerank false is good enough for 65% of searches and is VERY fast, rerank true is much slower but good enough for 95% of searches. Best to use rerank true if rerank false didn't quite help enough.
|
||
Example – search_codebase("user authentication python", top_k=5, rerank=True) → 5 relevant snippets."""
|
||
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 clean TOON format directly
|
||
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({
|
||
'file': meta.get('file', ''),
|
||
'line': meta.get('line', ''),
|
||
'type': meta.get('type', 'code'),
|
||
'score': score,
|
||
'content': _extract_clean_content(chunk_text)
|
||
})
|
||
|
||
return _format_toon_results(results, query, "search")
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Rerank step failed: {e}")
|
||
# Fall through to non-reranked results
|
||
|
||
# Non-reranked results - build clean TOON format directly
|
||
results = []
|
||
for i, (text, meta, score) in enumerate(hy[:top_k], 1):
|
||
results.append({
|
||
'file': meta.get('file', ''),
|
||
'line': meta.get('line', ''),
|
||
'type': meta.get('type', 'code'),
|
||
'score': score,
|
||
'content': _extract_clean_content(text)
|
||
})
|
||
|
||
return _format_toon_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 _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()
|
||
|
||
|
||
def _format_toon_results(results: List[Dict], query: str, result_type: str) -> str:
|
||
"""Format results in clean TOON format."""
|
||
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,score,content}}:"]
|
||
|
||
for result in results:
|
||
row = [
|
||
result.get('file', ''),
|
||
result.get('line', ''),
|
||
f"{result.get('score', 0.0):.2f}",
|
||
result.get('content', '')[:500] # 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 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 search_codebase.
|
||
When not to use – For browsing the entire file or unrelated sections; use search_codebase 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 or update an existing one.
|
||
|
||
The function will:
|
||
1. Clone the repo if it does not already exist locally,
|
||
2. Pull the latest changes if it does,
|
||
3. Reset the global `CODEBASE_PATH` to the local clone,
|
||
4. Remove any stale index artefacts,
|
||
5. Re‑build the vector/BM25 indexes with the same lock‑protected flow used by `rebuild_index`,
|
||
6. Return a human‑readable status message.
|
||
|
||
Next steps:
|
||
• Call `health_check` to confirm that the server is ready.
|
||
• Use `search_codebase` or `find_code_references` to explore the new code.
|
||
• Use `read_file_lines_tool` to inspect any hit in detail.
|
||
"""
|
||
import shutil
|
||
from pathlib import Path
|
||
import subprocess
|
||
|
||
# ----- 1. Determine where to clone ---------------------------------
|
||
repo_name = git_url.split('/')[-1]
|
||
if repo_name.endswith('.git'):
|
||
repo_name = repo_name[:-4]
|
||
clone_dir = Path.cwd() / repo_name
|
||
|
||
# ----- 2. Clone / pull -----------------------------------------------
|
||
try:
|
||
if clone_dir.exists():
|
||
# Existing repo → pull
|
||
subprocess.run(
|
||
["git", "-C", str(clone_dir), "pull"],
|
||
check=True,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.PIPE,
|
||
text=True,
|
||
)
|
||
else:
|
||
# New repo → clone
|
||
subprocess.run(
|
||
["git", "clone", git_url],
|
||
cwd=str(Path.cwd()),
|
||
check=True,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.PIPE,
|
||
text=True,
|
||
)
|
||
except subprocess.CalledProcessError as e:
|
||
return f"❌ Failed to clone/pull {git_url}: {e.stderr.strip()}"
|
||
|
||
# ----- 3. Update the global CODEBASE_PATH --------------------------------
|
||
global CODEBASE_PATH
|
||
CODEBASE_PATH = clone_dir
|
||
|
||
# ----- 4. Remove stale indices ----------------------------------------
|
||
for artefact in ("chroma_db", "bm25_index.json", "embeddings"):
|
||
artefact_path = CODEBASE_PATH / artefact
|
||
if artefact_path.exists():
|
||
try:
|
||
if artefact_path.is_dir():
|
||
shutil.rmtree(artefact_path)
|
||
else:
|
||
artefact_path.unlink()
|
||
except Exception as e:
|
||
logger.warning(f"Could not delete {artefact_path}: {e}")
|
||
|
||
# ----- 5. Re‑build the index (protected by the startup lock) ------------
|
||
try:
|
||
rebuild_index()
|
||
except Exception as e:
|
||
return f"❌ Index rebuild failed: {e}"
|
||
|
||
# ----- 6. Report success -----------------------------------------------
|
||
branch_name = subprocess.check_output(
|
||
["git", "-C", str(CODEBASE_PATH), "rev-parse", "--abbrev-ref", "HEAD"],
|
||
text=True,
|
||
).strip()
|
||
git_hash = subprocess.check_output(
|
||
["git", "-C", str(CODEBASE_PATH), "rev-parse", "HEAD"],
|
||
text=True,
|
||
).strip()
|
||
|
||
return f"✅ Initialized repository at {CODEBASE_PATH} on branch {branch_name} with hash {git_hash}"
|
||
|
||
|
||
@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()
|
||
|
||
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("Use 'python serve_http.py' for HTTP server")
|
||
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)
|