This commit is contained in:
2025-11-14 23:36:33 +00:00
parent c4cdc0e688
commit e706853d0c
4 changed files with 88 additions and 626 deletions
-124
View File
@@ -1,124 +0,0 @@
# compact_toon.py
"""
Ultra-compact Toon format translator for maximum token efficiency
"""
import re
from typing import Dict, List, Any, Optional
class CompactToon:
"""Minimal Toon format translator for maximum token savings."""
@staticmethod
def search_results(result_string: str, query: str = "") -> str:
"""Convert search results to minimal Toon format."""
if not result_string or "No results found" in result_string:
return result_string
results = result_string.split("---")
toon_blocks = []
for result in results:
if not result.strip():
continue
toon_block = CompactToon._minimal_result(result.strip())
if toon_block:
toon_blocks.append(toon_block)
return "\n\n".join(toon_blocks)
@staticmethod
def _minimal_result(result_text: str) -> Optional[str]:
"""Create minimal Toon block for a single result."""
try:
# Fast regex extraction
file_match = re.search(r"File:\s*(.+?)\n", result_text)
line_match = re.search(r"Line:\s*(\d+)", result_text)
type_match = re.search(r"Type:\s*(.+?)\n", result_text)
score_match = re.search(r"score:\s*([\d.]+)", result_text)
if not file_match:
return None
file_path = file_match.group(1).strip()
line_num = line_match.group(1) if line_match else "1"
# Extract content efficiently
lines = result_text.split('\n')
content_lines = []
in_content = False
for line in lines:
if in_content:
content_lines.append(line)
elif line.strip() and not any(line.startswith(x) for x in
['###', 'File:', 'Line:', 'Type:', 'Language:']):
in_content = True
content_lines.append(line)
content = '\n'.join(content_lines).strip()
# Build minimal Toon format
toon_lines = []
toon_lines.append(f"file://{file_path}:{line_num}")
if type_match:
toon_lines.append(f" type: {type_match.group(1).strip()}")
if score_match:
toon_lines.append(f" score: {score_match.group(1)}")
toon_lines.append(" content: |")
# Add content with minimal processing
for line in content.split('\n'):
toon_lines.append(f" {line}")
return "\n".join(toon_lines)
except Exception:
return None
@staticmethod
def references(result_string: str, symbol: str = "") -> str:
"""Convert reference results to minimal format."""
if not result_string or "No references found" in result_string:
return result_string
results = result_string.split("---")
toon_blocks = []
for result in results:
if not result.strip():
continue
toon_block = CompactToon._minimal_reference(result.strip())
if toon_block:
toon_blocks.append(toon_block)
return "\n\n".join(toon_blocks)
@staticmethod
def _minimal_reference(result_text: str) -> Optional[str]:
"""Create minimal reference block."""
try:
file_match = re.search(r"File:\s*(.+?)\n", result_text)
line_match = re.search(r"Line:\s*(\d+)", result_text)
if not file_match or not line_match:
return None
# Extract context efficiently
context_match = re.search(r"Context:\s*(.+)", result_text, re.DOTALL)
context = context_match.group(1).strip() if context_match else ""
toon_lines = []
toon_lines.append(f"file://{file_match.group(1).strip()}:{line_match.group(1)}")
toon_lines.append(" type: reference")
toon_lines.append(" content: |")
for line in context.split('\n'):
toon_lines.append(f" {line}")
return "\n".join(toon_lines)
except Exception:
return None
+88 -161
View File
@@ -1547,96 +1547,108 @@ def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple:
logger.info(f"🔧 Parsing Rust file: {filepath}")
# The binary built from `tools/src/main.rs` lives in
# `tools/src/target/release/rust_parser` it is *always* shipped
# with the release.
rust_helper = Path("./tools/src/target/release/rust_parser")
# Always use the helper; if it is missing, the whole Rust parsing
# step will simply return an empty tuple, which is the safest
# fallback.
logger.info(f"🔧 Using Rust parser binary: {rust_helper}")
try:
code = filepath.read_text(encoding="utf-8")
lines = code.splitlines(keepends=True)
# Try using rust-analyzer AST if available
rust_helper = Path("./tools/parse_rust_ast")
if rust_helper.exists():
logger.info(f" Using Rust AST helper: {rust_helper}")
try:
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")
# Capture the import paths for each file they are specific to the file
# and will be attached to every chunk emitted for this file.
import_paths = []
for node in ast.walk(tree):
if isinstance(node, ItemUse):
import_paths.append(use_tree_to_string(node.tree))
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)
chunk_code = "".join(lines[start_line:end_line])
metadata["imports"] = import_paths
# Build comprehensive metadata
metadata = {
"file": str(filepath),
"name": d.get("name", ""),
"type": d.get("type_", ""),
"line": start_line + 1,
"language": "rust",
}
# Run the Rust binary.
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")
# Add Rust-specific metadata
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")
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)
chunk_code = "".join(lines[start_line:end_line])
# Log what we found
logger.info(f" Declaration {i+1}: {metadata.get('type', 'unknown')} {metadata.get('name', 'unnamed')}")
# Build comprehensive metadata
metadata = {
"file": str(filepath),
"name": d.get("name", ""),
"type": d.get("type_", ""),
"line": start_line + 1,
"language": "rust",
}
chunk_text = f"File: {filepath}\nType: {d.get('type_')}\nName: {d.get('name')}\n"
# Add Rust-specific metadata
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")
# Add Rust-specific details to text
if d.get("visibility"):
chunk_text += f"Visibility: {d.get('visibility')}\n"
if d.get("is_async"):
chunk_text += "Async: yes\n"
if d.get("is_unsafe"):
chunk_text += "Unsafe: yes\n"
if d.get("generics"):
chunk_text += f"Generics: {d.get('generics')}\n"
if d.get("traits"):
chunk_text += f"Implements: {', '.join(d.get('traits', []))}\n"
if d.get("fields"):
chunk_text += f"Fields: {', '.join(d.get('fields', []))}\n"
# Log what we found
logger.info(f" Declaration {i+1}: {metadata.get('type', 'unknown')} {metadata.get('name', 'unnamed')}")
chunk_text += f"Code:\n{chunk_code}"
chunk_text = f"File: {filepath}\nType: {d.get('type_')}\nName: {d.get('name')}\n"
chunks.append({
"text": chunk_text,
"metadata": metadata
})
# Add Rust-specific details to text
if d.get("visibility"):
chunk_text += f"Visibility: {d.get('visibility')}\n"
if d.get("is_async"):
chunk_text += "Async: yes\n"
if d.get("is_unsafe"):
chunk_text += "Unsafe: yes\n"
if d.get("generics"):
chunk_text += f"Generics: {d.get('generics')}\n"
if d.get("traits"):
chunk_text += f"Implements: {', '.join(d.get('traits', []))}\n"
if d.get("fields"):
chunk_text += f"Fields: {', '.join(d.get('fields', []))}\n"
if chunks:
logger.info(f"✅ Successfully parsed {len(chunks)} chunks from {filepath}")
return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)
else:
logger.warning(f"❌ AST helper returned declarations but no chunks were created")
chunk_text += f"Code:\n{chunk_code}"
except Exception as e:
logger.warning(f"Rust AST helper failed for {filepath}: {e}; falling back to regex")
chunks.append({
"text": chunk_text,
"metadata": metadata
})
if chunks:
logger.info(f"✅ Successfully parsed {len(chunks)} chunks from {filepath}")
return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)
else:
logger.warning(f"❌ AST helper returned declarations but no chunks were created")
except Exception as e:
logger.warning(f"parse_rust_file failed {filepath}: {e}")
logger.warning(f"Rust AST helper failed for {filepath}: {e}")
final_result = tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)
logger.info(f"📦 Final result: {len(final_result)} chunks with metadata")
@@ -2044,27 +2056,6 @@ def build_indexes():
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 == "impl":
# Implementation blocks
target = m.get("target", "unknown")
node_id = f"{lang}::{node_type}::{fpath}::{target}"
graph.add_node(node_id, **base_attrs, type=node_type,
target=target,
traits=m.get("traits", []),
methods=m.get("methods", []))
# Connect impl to file
graph.add_edge(file_node_id, node_id, "contains")
# Store for second-pass
created_nodes[node_id] = base_attrs.copy()
created_nodes[node_id]["type"] = node_type
created_nodes[node_id]["target"] = target
created_nodes[node_id]["traits"] = m.get("traits", [])
created_nodes[node_id]["methods"] = m.get("methods", [])
elif node_type == "trait":
# Traits
graph.add_node(node_id, **base_attrs, type=node_type,
@@ -2094,37 +2085,7 @@ def build_indexes():
else:
imports = imports_raw
for imp in imports:
# Try to find an existing node for this import
candidate_id = None
for node_id, data in created_nodes.items():
if data.get("type") in {"module", "import"} and data.get("name") == imp:
candidate_id = node_id
break
if not candidate_id:
# No existing node found - create a placeholder
candidate_id = f"import::{imp}"
if candidate_id not in graph.graph.nodes:
# Double-check for a module node with this name
found = False
for n_id, n_data in created_nodes.items():
if n_data.get("name") == imp and n_data.get("type") == "module":
candidate_id = n_id
found = True
break
if not found:
# Create a generic import node
graph.add_node(
candidate_id,
type="Import",
name=imp,
lang="unknown",
file="",
)
graph.add_edge(file_node_id, candidate_id, "imports")
# Create the connection
graph.add_edge(file_node_id, candidate_id, "imports")
graph._add_import_edges(file_node_id, imports)
# ---------- FUNCTION CALLS ----------
calls_raw = m.get("calls", [])
@@ -2134,40 +2095,7 @@ def build_indexes():
else:
calls = calls_raw
for call in calls:
# Try to find the function/method being called
target_id = None
for node_id, data in created_nodes.items():
if data.get("type") in {"function", "method", "constructor"} and data.get("name") == call:
target_id = node_id
break
if not target_id:
# No existing function found - create a placeholder
target_id = f"call::{call}"
if target_id not in graph.graph.nodes:
# Double-check for a matching function
found = False
for n_id, n_data in created_nodes.items():
if (
n_data.get("name") == call
and n_data.get("type") in {"function", "method", "constructor"}
):
target_id = n_id
found = True
break
if not found:
# Create a generic function call node
graph.add_node(
target_id,
type="FunctionCall",
name=call,
lang="unknown",
file="",
)
graph.add_edge(node_id, target_id, "calls")
# Create the connection
graph.add_edge(node_id, target_id, "calls")
graph._add_call_edges(node_id, calls, fpath)
# -------------------------------------------------
@@ -3732,7 +3660,6 @@ if __name__ == "__main__":
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
-86
View File
@@ -1,86 +0,0 @@
import re
import ast
import sqlglot
from typing import Dict, List, Any
from pathlib import Path
def extract_sql_schema(sql_code: str) -> Dict[str, List[str]]:
"""
Extracts table -> columns mapping from SQL code using sqlglot.
Returns a canonical schema.
"""
schema: Dict[str, List[str]] = {}
try:
statements = sqlglot.parse(sql_code, read="postgres")
except Exception:
return schema
for stmt in statements:
if stmt.key and stmt.key.upper() == "CREATE":
for table in stmt.find_all(sqlglot.exp.Create):
try:
tname = table.this.this
cols = []
for coldef in table.find_all(sqlglot.exp.ColumnDef):
cname = getattr(coldef.this, "name", None)
if cname:
cols.append(cname)
if tname and cols:
schema[tname] = cols
except Exception:
continue
return schema
def extract_python_structure(code: str) -> Dict[str, List[str]]:
"""Return module structure: functions and classes."""
try:
tree = ast.parse(code)
except SyntaxError:
return {}
funcs = [n.name for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)]
classes = [n.name for n in ast.walk(tree) if isinstance(n, ast.ClassDef)]
return {"functions": funcs, "classes": classes}
def extract_go_structure(code: str) -> Dict[str, List[str]]:
"""Simple regex-based Go structure detection (lightweight)."""
funcs = re.findall(r"func\s+([A-Z]\w+)", code)
structs = re.findall(r"type\s+(\w+)\s+struct", code)
return {"functions": funcs, "structs": structs}
def extract_rust_structure(code: str) -> Dict[str, List[str]]:
"""Heuristic Rust index."""
structs = re.findall(r"struct\s+(\w+)", code)
traits = re.findall(r"trait\s+(\w+)", code)
funcs = re.findall(r"fn\s+(\w+)", code)
return {"functions": funcs, "structs": structs, "traits": traits}
def extract_svelte_structure(code: str) -> Dict[str, List[str]]:
"""Minimal Svelte export/prop finder."""
exports = re.findall(r"export\s+let\s+(\w+)", code)
funcs = re.findall(r"function\s+(\w+)", code)
return {"props": exports, "functions": funcs}
def build_quick_index(language: str, code: str, filepath: Path) -> Dict[str, Any]:
"""Dispatch to the appropriate structure extractor."""
data = {}
if language == "sql":
data = extract_sql_schema(code)
elif language == "python":
data = extract_python_structure(code)
elif language == "go":
data = extract_go_structure(code)
elif language == "rust":
data = extract_rust_structure(code)
elif language == "svelte":
data = extract_svelte_structure(code)
return {
"text": f"Quick index for {filepath.name}:\n{data}",
"metadata": {"language": language, "file": str(filepath), "type": "index"}
}
-255
View File
@@ -1,255 +0,0 @@
#!/usr/bin/env python3
"""
HTTP server for MCP codebase RAG with REST API
"""
import uvicorn
import logging
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional
import signal
import sys
import json
from enhanced_toon import EnhancedToon
from mcp_codebase import (
startup,
search_codebase as _search_codebase,
find_code_references as _find_code_references,
read_file_lines,
rebuild_index as _rebuild_index
)
logger = logging.getLogger("rag-mcp")
def signal_handler(sig, frame):
"""Handle graceful shutdown on Ctrl+C."""
logger.info("\n=== HTTP Server Shutdown ===")
logger.info("Goodbye!\n")
sys.exit(0)
# Request/Response models
class SearchRequest(BaseModel):
query: str
top_k: int = 5
rerank: bool = True
class ReferenceRequest(BaseModel):
symbol: str
top_k: int = 20
class ReadFileRequest(BaseModel):
path: str
start: int = 1
end: Optional[int] = None
class SymbolRequest(BaseModel):
symbol: str
def create_app():
"""Create FastAPI app with MCP tool endpoints"""
app = FastAPI(
title="MCP Codebase RAG Server",
description="Codebase search and analysis via MCP tools over HTTP",
version="1.0.0"
)
# Add CORS middleware for remote access
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Configure this appropriately for production
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Health check endpoint
@app.get("/")
async def root():
return {
"status": "running",
"service": "MCP Codebase RAG Server",
"version": "1.0.0",
"endpoints": {
"docs": "/docs",
"health": "/health",
"tools": "/tools",
"search": "POST /search",
"references": "POST /references",
"read_file": "POST /read_file",
"rebuild": "POST /rebuild"
}
}
@app.get("/health")
async def health():
"""Get server health and statistics"""
try:
# Import here to avoid circular import issues
from mcp_codebase import health_check
result = health_check()
# Parse JSON string to dict for better API response
try:
result_dict = json.loads(result)
return result_dict
except:
return {"result": result}
except Exception as e:
logger.exception("Health check failed")
raise HTTPException(status_code=500, detail=str(e))
@app.get("/tools")
async def list_tools():
"""List available MCP tools"""
return {
"tools": [
{
"name": "search_codebase",
"description": "Search the entire codebase using hybrid RAG",
"endpoint": "POST /search",
"parameters": {
"query": "string (required)",
"top_k": "int (default: 5)",
"rerank": "bool (default: true)"
}
},
{
"name": "find_code_references",
"description": "Find all references to a specific symbol",
"endpoint": "POST /references",
"parameters": {
"symbol": "string (required)",
"top_k": "int (default: 20)"
}
},
{
"name": "read_file_lines",
"description": "Read specific lines from a file with context",
"endpoint": "POST /read_file",
"parameters": {
"path": "string (required)",
"start": "int (default: 1)",
"end": "int (optional)"
}
},
{
"name": "rebuild_index",
"description": "Force rebuild of search indexes",
"endpoint": "POST /rebuild",
"parameters": {}
}
]
}
@app.post("/search")
async def search(request: SearchRequest):
"""Search codebase - returns enhanced Toon format"""
try:
result = _search_codebase(
query=request.query,
top_k=request.top_k,
rerank=request.rerank
)
enhanced_result = EnhancedToon.search_results(result, request.query)
return {
"result": enhanced_result,
"query": request.query,
"top_k": request.top_k
}
except Exception as e:
logger.exception("Search failed")
raise HTTPException(status_code=500, detail=str(e))
@app.post("/references")
async def references(request: ReferenceRequest):
"""Find all references to a symbol - returns enhanced Toon format"""
try:
result = _find_code_references(
symbol=request.symbol,
top_k=request.top_k
)
enhanced_result = EnhancedToon.reference_results(result, request.symbol)
return {
"result": enhanced_result,
"symbol": request.symbol,
"top_k": request.top_k
}
except Exception as e:
logger.exception("Reference search failed")
raise HTTPException(status_code=500, detail=str(e))
@app.post("/read_file")
async def read_file(request: ReadFileRequest):
"""Read file lines with context - returns enhanced Toon format"""
try:
result = read_file_lines(
path=request.path,
start=request.start,
end=request.end
)
enhanced_result = EnhancedToon.file_content_results(result, request.path)
return {
"result": enhanced_result,
"path": request.path,
"start_line": request.start,
"end_line": request.end
}
except Exception as e:
logger.exception("Read file failed")
raise HTTPException(status_code=500, detail=str(e))
@app.post("/rebuild")
async def rebuild():
"""Rebuild search indexes"""
try:
result = _rebuild_index()
# For rebuild, we might not need enhanced format as it's usually a simple status message
return {"result": result}
except Exception as e:
logger.exception("Rebuild failed")
raise HTTPException(status_code=500, detail=str(e))
return app
if __name__ == "__main__":
# Register signal handlers
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
try:
# Initialize the indexes (this will load or build them)
startup()
logger.info("="*60)
logger.info("🚀 Starting HTTP server on: http://0.0.0.0:8000")
logger.info("📚 REST API endpoints:")
logger.info(" GET / - Service info")
logger.info(" GET /health - Server health & stats")
logger.info(" GET /tools - List available tools")
logger.info(" POST /search - Search codebase")
logger.info(" POST /references - Find symbol references")
logger.info(" POST /read_file - Read file with context")
logger.info(" POST /rebuild - Rebuild indexes")
logger.info("📖 API docs available at: http://0.0.0.0:8000/docs")
logger.info("⏹️ Press Ctrl+C to stop")
logger.info("="*60)
# Create and run app
app = create_app()
uvicorn.run(
app,
host="0.0.0.0",
port=8000,
log_level="info",
access_log=True
)
except KeyboardInterrupt:
signal_handler(signal.SIGINT, None)
except Exception as e:
logger.exception("Fatal error during HTTP server startup")
sys.exit(1)