Files
BUGFS/BUGFS_mcp_server.py
2026-01-08 15:50:54 +09:00

783 lines
26 KiB
Python
Executable File

#!/usr/bin/env python3
"""
MCP server that crawls GitHub issues, embeds the reports, and recommends bug
locations by reusing the crawl/embedding utilities plus the RQ1 workflow.
Uses sqlite-vec for efficient vector similarity search with minimal memory footprint.
"""
from __future__ import annotations
import json
import logging
import os
import re
import sqlite3
import struct
import time
from collections import Counter, defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple
import numpy as np
import sqlite_vec
from dotenv import load_dotenv
from openai import OpenAI
from tqdm import tqdm
from crawl_issues_With_diffs import crawl_issues_with_diffs
try:
from fastmcp import FastMCP, Context
except ImportError as exc: # pragma: no cover - handled at runtime
raise SystemExit(
"The modelcontextprotocol package is required. "
"Install it with `pip install modelcontextprotocol`."
) from exc
load_dotenv()
# Configure logging
LOG_DIR = Path(os.getenv("BUGFS_LOG_DIR", "logs"))
LOG_DIR.mkdir(parents=True, exist_ok=True)
LOG_FILE = LOG_DIR / "bugfs_mcp.log"
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
handlers=[
logging.FileHandler(LOG_FILE, encoding="utf-8"),
],
)
logger = logging.getLogger("bugfs_mcp")
DEFAULT_MODEL = os.getenv("BUGFS_EMBED_MODEL", "text-embedding-3-small")
CHUNK_CHAR_LIMIT = int(os.getenv("BUGFS_CHUNK_CHAR_LIMIT", "3000"))
DATA_DIR = Path(os.getenv("BUGFS_DATA_DIR", "data"))
DATA_DIR.mkdir(parents=True, exist_ok=True)
# Embedding dimension for text-embedding-3-small
EMBEDDING_DIM = 1536
# Regex to extract hunk header: line info and context
HUNK_HEADER_REGEX = re.compile(
r"@@ -\d+,\d+ \+(?P<start>\d+),(?P<length>\d+) @@(?:\s+(?P<context>.+))?$",
flags=re.MULTILINE,
)
# Patterns to extract method/function name from the hunk context
# Python: def function_name
PYTHON_FUNC_PATTERN = re.compile(r"def\s+(?P<name>\w+)")
# Java/Groovy/Kotlin method: modifiers + return_type + method_name(
# e.g., "public void methodName(", "private CloseableHttpClient create("
JAVA_METHOD_PATTERN = re.compile(
r"(?:public|private|protected|static|final|abstract|synchronized|void|\w+)\s+(?P<name>\w+)\s*\("
)
# Class/interface/enum definition
CLASS_PATTERN = re.compile(r"(?:class|interface|enum)\s+(?P<name>\w+)")
def slugify_repo(owner: str, repo: str) -> str:
return f"{owner}_{repo}".replace("/", "_")
def ensure_text(value: Optional[str]) -> str:
return (value or "").strip()
def build_issue_text(issue: Dict[str, Any]) -> str:
title = ensure_text(issue.get("title"))
body = ensure_text(issue.get("body"))
text = f"{title}\n\n{body}".strip()
return text
def split_into_chunks(text: str, limit: int = CHUNK_CHAR_LIMIT) -> List[str]:
if not text:
return []
return [text[i : i + limit] for i in range(0, len(text), limit)]
def sanitize_issue_body(issue: Dict[str, Any], min_word_count: int = 1) -> Optional[str]:
text = build_issue_text(issue)
text = re.sub(r"\s+", " ", text)
if len(text.split()) < min_word_count:
return None
return text
def normalize_filename(path_str: str) -> str:
path_str = path_str.replace("\\", "/")
parts = [p for p in path_str.split("/") if p]
if not parts:
return path_str
return parts[-1]
@dataclass(frozen=True)
class IssueLocation:
start: int
length: int
function: str
file: str
@property
def label(self) -> str:
return f"{self.file}:{self.function}"
def extract_method_from_context(context: str) -> Optional[str]:
"""Extract method/function name from git diff hunk context."""
if not context:
return None
# Try Python function pattern
match = PYTHON_FUNC_PATTERN.search(context)
if match:
return match.group("name")
# Try Java/Groovy/Kotlin method pattern
match = JAVA_METHOD_PATTERN.search(context)
if match:
name = match.group("name")
# Filter out common false positives (keywords that look like method names)
if name not in ("if", "for", "while", "switch", "catch", "synchronized", "return", "new", "throw"):
return name
# Try class/interface/enum pattern
match = CLASS_PATTERN.search(context)
if match:
return match.group("name")
return None
def extract_locations(issue: Dict[str, Any]) -> List[IssueLocation]:
"""Extract file:function locations from git diff patches.
Parses the hunk header context (text after @@ ... @@) to find the enclosing
function/method name. Supports Python, Java, Groovy, Kotlin, and class definitions.
"""
locations: List[IssueLocation] = []
seen: set[tuple[str, str]] = set() # Avoid duplicates
for commit in issue.get("commits", []):
filename = normalize_filename(commit.get("filename", "unknown"))
patch = commit.get("patch")
if not patch:
continue
# Parse each hunk header
for match in HUNK_HEADER_REGEX.finditer(patch):
start = int(match.group("start"))
length = int(match.group("length"))
context = match.group("context") or ""
# Try to extract method/function name from context
method_name = extract_method_from_context(context)
if method_name:
key = (filename, method_name)
if key not in seen:
seen.add(key)
locations.append(
IssueLocation(start=start, length=length, function=method_name, file=filename)
)
return locations
class Embedder:
def __init__(self, model: str = DEFAULT_MODEL):
self.model = model
self._client: Optional[OpenAI] = None
@property
def client(self) -> OpenAI:
if self._client is None:
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
raise RuntimeError("OPENAI_API_KEY is not set. Please configure your OpenAI API key.")
self._client = OpenAI(api_key=api_key)
return self._client
def embed_text(self, text: str) -> np.ndarray:
chunks = split_into_chunks(text)
if not chunks:
raise ValueError("Cannot embed empty text.")
last_error: Optional[Exception] = None
for attempt in range(5):
try:
response = self.client.embeddings.create(model=self.model, input=chunks)
vectors = [np.array(item.embedding, dtype=np.float32) for item in response.data]
avg_vector = np.mean(vectors, axis=0)
return avg_vector.astype(np.float32)
except Exception as exc:
time.sleep(1 + attempt)
last_error = exc
raise RuntimeError(f"Failed to embed text after retries: {last_error}")
def serialize_f32(vector: np.ndarray) -> bytes:
"""Serialize a numpy float32 vector to bytes for sqlite-vec."""
return struct.pack(f"{len(vector)}f", *vector)
def normalize_vector(vector: np.ndarray) -> np.ndarray:
"""Normalize vector for cosine similarity."""
norm = np.linalg.norm(vector)
if norm > 0:
return vector / norm
return vector
class SqliteVecCorpus:
"""Efficient vector search using sqlite-vec with minimal memory footprint."""
def __init__(self, db_path: Path):
self.db_path = db_path
self._conn: Optional[sqlite3.Connection] = None
@property
def conn(self) -> sqlite3.Connection:
if self._conn is None:
self._conn = sqlite3.connect(str(self.db_path))
self._conn.enable_load_extension(True)
sqlite_vec.load(self._conn)
self._conn.enable_load_extension(False)
return self._conn
def close(self):
if self._conn:
self._conn.close()
self._conn = None
def _init_schema(self):
"""Initialize database schema."""
self.conn.executescript(f"""
-- Raw issues table (stores full issue data as JSON)
CREATE TABLE IF NOT EXISTS raw_issues (
id INTEGER PRIMARY KEY,
issue_id TEXT UNIQUE,
data TEXT NOT NULL,
created_at TEXT,
updated_at TEXT
);
-- Indexed issues table (for vector search)
CREATE TABLE IF NOT EXISTS issues (
id INTEGER PRIMARY KEY,
issue_id TEXT UNIQUE,
content TEXT,
locations TEXT
);
-- Vector index
CREATE VIRTUAL TABLE IF NOT EXISTS issue_vectors USING vec0(
id INTEGER PRIMARY KEY,
embedding float[{EMBEDDING_DIM}]
);
CREATE INDEX IF NOT EXISTS idx_issue_id ON issues(issue_id);
CREATE INDEX IF NOT EXISTS idx_raw_issue_id ON raw_issues(issue_id);
""")
self.conn.commit()
def count(self) -> int:
"""Get number of indexed issues."""
cursor = self.conn.execute("SELECT COUNT(*) FROM issues")
return cursor.fetchone()[0]
def count_raw_issues(self) -> int:
"""Get number of raw issues stored."""
cursor = self.conn.execute("SELECT COUNT(*) FROM raw_issues")
return cursor.fetchone()[0]
def save_raw_issues(self, issues: Sequence[Dict[str, Any]]) -> int:
"""Save raw issues to database. Returns number of issues saved."""
cursor = self.conn.cursor()
saved = 0
for issue in tqdm(issues, desc="Saving raw issues", unit="issue"):
issue_id = str(issue.get("id", ""))
if not issue_id:
continue
created_at = issue.get("created_at")
updated_at = issue.get("updated_at")
cursor.execute(
"""INSERT OR REPLACE INTO raw_issues (issue_id, data, created_at, updated_at)
VALUES (?, ?, ?, ?)""",
(issue_id, json.dumps(issue, ensure_ascii=False), created_at, updated_at)
)
saved += 1
if saved % 1000 == 0:
self.conn.commit()
self.conn.commit()
return saved
def load_raw_issues(self) -> List[Dict[str, Any]]:
"""Load all raw issues from database."""
cursor = self.conn.execute("SELECT data FROM raw_issues")
return [json.loads(row[0]) for row in cursor.fetchall()]
def get_raw_issue(self, issue_id: str) -> Optional[Dict[str, Any]]:
"""Get a single raw issue by ID."""
cursor = self.conn.execute(
"SELECT data FROM raw_issues WHERE issue_id = ?", (issue_id,)
)
row = cursor.fetchone()
return json.loads(row[0]) if row else None
def get_raw_issue_ids(self) -> set:
"""Get set of all raw issue IDs."""
cursor = self.conn.execute("SELECT issue_id FROM raw_issues")
return {row[0] for row in cursor.fetchall()}
def save_embedding(self, issue_id: str, embedding: List[float], model: str) -> None:
"""Save embedding for an issue."""
cursor = self.conn.cursor()
# Get existing data, update it, and save back
cursor.execute("SELECT data FROM raw_issues WHERE issue_id = ?", (issue_id,))
row = cursor.fetchone()
if row:
data = json.loads(row[0])
data["embedding"] = embedding
data["embedding_model"] = model
cursor.execute(
"UPDATE raw_issues SET data = ? WHERE issue_id = ?",
(json.dumps(data, ensure_ascii=False), issue_id)
)
self.conn.commit()
def get_embedded_issue_ids(self) -> set:
"""Get set of issue IDs that have embeddings."""
cursor = self.conn.execute(
"SELECT issue_id FROM raw_issues WHERE json_extract(data, '$.embedding') IS NOT NULL"
)
return {row[0] for row in cursor.fetchall()}
def iter_raw_issues(self, batch_size: int = 100):
"""Iterate over raw issues in batches to reduce memory usage."""
cursor = self.conn.execute("SELECT COUNT(*) FROM raw_issues")
total = cursor.fetchone()[0]
for offset in range(0, total, batch_size):
cursor = self.conn.execute(
"SELECT issue_id, data FROM raw_issues LIMIT ? OFFSET ?",
(batch_size, offset)
)
for row in cursor.fetchall():
yield row[0], json.loads(row[1])
def rebuild_vector_index(self, min_word_count: int = 1) -> int:
"""Rebuild vector index from raw_issues with embeddings."""
# Clear existing index
self.conn.execute("DELETE FROM issues")
self.conn.execute("DELETE FROM issue_vectors")
self.conn.commit()
batch_issues = []
batch_vectors = []
indexed = 0
for issue_id, issue in tqdm(
self.iter_raw_issues(batch_size=500),
desc="Building vector index",
unit="issue",
total=self.count_raw_issues()
):
embedding = issue.get("embedding")
if not embedding:
continue
# Handle double-encoded JSON strings
if isinstance(embedding, str):
try:
embedding = json.loads(embedding)
except json.JSONDecodeError:
continue
locations = extract_locations(issue)
if not locations:
continue
text = sanitize_issue_body(issue, min_word_count=min_word_count)
if not text:
continue
locations_as_dicts = [
{"start": loc.start, "length": loc.length, "function": loc.function, "file": loc.file}
for loc in locations
]
vec = normalize_vector(np.asarray(embedding, dtype=np.float32))
batch_issues.append((issue_id, text, json.dumps(locations_as_dicts)))
batch_vectors.append(serialize_f32(vec))
indexed += 1
if len(batch_issues) >= 1000:
self._batch_insert(batch_issues, batch_vectors)
batch_issues = []
batch_vectors = []
if batch_issues:
self._batch_insert(batch_issues, batch_vectors)
return indexed
@classmethod
def build_from_issues(
cls, db_path: Path, issues: Sequence[Dict[str, Any]], min_word_count: int = 1
) -> "SqliteVecCorpus":
"""Build sqlite-vec index from issues with embeddings."""
# Remove existing db if exists
if db_path.exists():
db_path.unlink()
corpus = cls(db_path)
corpus._init_schema()
batch_issues = []
batch_vectors = []
for issue in tqdm(issues, desc="Indexing issues", unit="issue"):
embedding = issue.get("embedding")
if not embedding:
continue
locations = extract_locations(issue)
if not locations:
continue
text = sanitize_issue_body(issue, min_word_count=min_word_count)
if not text:
continue
locations_as_dicts = [
{"start": loc.start, "length": loc.length, "function": loc.function, "file": loc.file}
for loc in locations
]
# Normalize embedding for cosine similarity
vec = normalize_vector(np.asarray(embedding, dtype=np.float32))
batch_issues.append((
str(issue["id"]),
text,
json.dumps(locations_as_dicts)
))
batch_vectors.append(serialize_f32(vec))
# Batch insert every 1000 records
if len(batch_issues) >= 1000:
corpus._batch_insert(batch_issues, batch_vectors)
batch_issues = []
batch_vectors = []
# Insert remaining
if batch_issues:
corpus._batch_insert(batch_issues, batch_vectors)
return corpus
def _batch_insert(self, issues: List[Tuple], vectors: List[bytes]):
"""Batch insert issues and vectors."""
cursor = self.conn.cursor()
# Insert issues
cursor.executemany(
"INSERT OR REPLACE INTO issues (issue_id, content, locations) VALUES (?, ?, ?)",
issues
)
# Get the rowids
issue_ids = [i[0] for i in issues]
placeholders = ",".join("?" * len(issue_ids))
cursor.execute(
f"SELECT id, issue_id FROM issues WHERE issue_id IN ({placeholders})",
issue_ids
)
id_map = {row[1]: row[0] for row in cursor.fetchall()}
# Insert vectors
vector_data = [(id_map[issues[i][0]], vectors[i]) for i in range(len(issues))]
cursor.executemany(
"INSERT OR REPLACE INTO issue_vectors (id, embedding) VALUES (?, ?)",
vector_data
)
self.conn.commit()
def add_issue(self, issue_id: str, content: str, locations: List[Dict], embedding: np.ndarray):
"""Add a single issue to the index."""
vec = normalize_vector(embedding.astype(np.float32))
cursor = self.conn.cursor()
cursor.execute(
"INSERT OR REPLACE INTO issues (issue_id, content, locations) VALUES (?, ?, ?)",
(issue_id, content, json.dumps(locations))
)
rowid = cursor.lastrowid
cursor.execute(
"INSERT OR REPLACE INTO issue_vectors (id, embedding) VALUES (?, ?)",
(rowid, serialize_f32(vec))
)
self.conn.commit()
def recommend(
self,
query_vector: np.ndarray,
neighbor_count: int = 10,
max_candidates: int = 10,
) -> Dict[str, Any]:
if self.count() == 0:
raise RuntimeError("No issues with embeddings and locations are available.")
# Normalize query vector for cosine similarity
query = normalize_vector(query_vector.astype(np.float32))
query_bytes = serialize_f32(query)
# Search using sqlite-vec
cursor = self.conn.execute("""
SELECT
i.issue_id,
i.locations,
v.distance
FROM issue_vectors v
JOIN issues i ON i.id = v.id
WHERE v.embedding MATCH ? AND k = ?
ORDER BY v.distance
""", (query_bytes, neighbor_count))
neighbors: List[Dict[str, Any]] = []
location_votes: Counter[str] = Counter()
location_sources: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
for row in cursor.fetchall():
issue_id, locations_json, distance = row
# sqlite-vec returns L2 distance, convert to similarity
# For normalized vectors: cosine_similarity = 1 - (distance^2 / 2)
similarity = float(1 - (distance ** 2) / 2)
neighbors.append({"issue_id": issue_id, "similarity": similarity})
locations = json.loads(locations_json)
for loc in locations:
label = f"{loc['file']}:{loc['function']}"
location_votes[label] += 1
location_sources[label].append({
"issue_id": issue_id,
"file": loc["file"],
"function": loc["function"],
"start": loc["start"],
"length": loc["length"],
"similarity": similarity,
})
recommendations: List[Dict[str, Any]] = []
for label, votes in sorted(location_votes.items(), key=lambda item: (-item[1], item[0]))[:max_candidates]:
sources = location_sources[label][:5]
recommendations.append({"location": label, "votes": votes, "sources": sources})
return {"neighbors": neighbors, "recommendations": recommendations}
class DatasetManager:
def __init__(self, data_dir: Path):
self.data_dir = data_dir
def db_path(self, slug: str, model: str = DEFAULT_MODEL) -> Path:
"""Path for sqlite-vec database."""
safe_model = model.replace("/", "_")
return self.data_dir / f"{slug}_{safe_model}.db"
def get_or_create_corpus(self, slug: str, model: str = DEFAULT_MODEL) -> SqliteVecCorpus:
"""Get or create a corpus for the given slug and model."""
db_file = self.db_path(slug, model)
corpus = SqliteVecCorpus(db_file)
corpus._init_schema()
return corpus
def crawl(self, owner: str, repo: str, token: str) -> Tuple[Path, int]:
slug = slugify_repo(owner, repo)
# Get existing issues from sqlite
db_file = self.db_path(slug)
corpus = self.get_or_create_corpus(slug)
existing_ids = corpus.get_raw_issue_ids()
# Convert to list format for crawler compatibility
existing = None
if existing_ids:
existing = corpus.load_raw_issues()
results = crawl_issues_with_diffs(
owner=owner,
repo=repo,
token=token,
output_path=str(db_file), # Not used for saving, just for logging
existing_results=existing,
)
# Save to sqlite
corpus.save_raw_issues(results)
corpus.close()
return db_file, len(results)
async def embed(
self, owner: str, repo: str, embedder: Embedder, ctx: Optional[Context] = None
) -> Dict[str, Any]:
slug = slugify_repo(owner, repo)
db_file = self.db_path(slug, embedder.model)
corpus = self.get_or_create_corpus(slug, embedder.model)
if corpus.count_raw_issues() == 0:
raise FileNotFoundError(f"No crawled issues found. Expected {db_file}.")
# Get already embedded issue IDs
embedded_ids = corpus.get_embedded_issue_ids()
total_issues = corpus.count_raw_issues()
embedded = 0
processed = 0
for issue_id, issue in corpus.iter_raw_issues(batch_size=100):
processed += 1
# Report progress to MCP client
if ctx:
await ctx.report_progress(progress=processed, total=total_issues)
# Skip already embedded
if issue_id in embedded_ids:
continue
text = sanitize_issue_body(issue)
if not text:
continue
vector = embedder.embed_text(text)
corpus.save_embedding(issue_id, vector.tolist(), embedder.model)
embedded += 1
logger.info(f"Embedded {processed}/{total_issues} (new: {embedded})")
# Rebuild vector index
logger.info(f"Building sqlite-vec index for {owner}/{repo}...")
if ctx:
await ctx.report_progress(progress=total_issues, total=total_issues)
indexed_count = corpus.rebuild_vector_index()
logger.info(f"sqlite-vec index saved to {db_file}")
corpus.close()
return {
"output_path": str(db_file),
"total_issues": total_issues,
"embedded_issues": len(embedded_ids) + embedded,
"indexed_issues": indexed_count,
"new_embeddings": embedded,
}
def load_corpus(self, owner: str, repo: str, model: str, min_word_count: int = 1) -> SqliteVecCorpus:
slug = slugify_repo(owner, repo)
db_file = self.db_path(slug, model)
if not db_file.exists():
raise FileNotFoundError(
f"No sqlite-vec database found at {db_file}. Please run embed_repo first."
)
logger.info(f"Loading sqlite-vec index from {db_file}")
corpus = SqliteVecCorpus(db_file)
# Rebuild index if empty but has raw issues with embeddings
if corpus.count() == 0 and corpus.count_raw_issues() > 0:
logger.info("Vector index empty, rebuilding from raw issues...")
corpus.rebuild_vector_index(min_word_count=min_word_count)
return corpus
dataset_manager = DatasetManager(DATA_DIR)
embedder = Embedder(DEFAULT_MODEL)
server = FastMCP("bugfs")
@server.tool()
def crawl_repo(
owner: str,
repo: str,
state: str = "closed",
) -> Dict[str, Any]:
"""Crawl GitHub issues plus associated commit diffs for the target repository."""
token = os.getenv("GITHUB_TOKEN")
if not token:
raise RuntimeError("GITHUB_TOKEN is required to crawl issues.")
path, count = dataset_manager.crawl(owner, repo, token)
return {"issues_path": str(path), "issues_crawled": count, "state": state}
@server.tool()
async def embed_repo(
owner: str,
repo: str,
model: str = "",
ctx: Context = None,
) -> Dict[str, Any]:
"""Embed every crawled issue (title + body) using the chosen embedding model."""
if model and model != embedder.model:
custom_embedder = Embedder(model)
else:
custom_embedder = embedder
result = await dataset_manager.embed(owner, repo, custom_embedder, ctx=ctx)
return result
@server.tool(name="run_BUGFS")
def run_BUGFS(
owner: str,
repo: str,
title: str,
body: str,
neighbors: int = 10,
max_candidates: int = 10,
model: str = "",
) -> Dict[str, Any]:
"""
Recommend bug locations (file:function) by retrieving similar historical issues.
"""
chosen_model = model if model else embedder.model
corpus = dataset_manager.load_corpus(owner, repo, chosen_model)
query_text = re.sub(r"\s+", " ", f"{title.strip()}\n\n{body.strip()}".strip()).strip()
if not query_text:
raise ValueError("title+body must contain text.")
query_embedder = embedder if chosen_model == embedder.model else Embedder(chosen_model)
query_vector = query_embedder.embed_text(query_text)
return corpus.recommend(query_vector, neighbor_count=neighbors, max_candidates=max_candidates)
def main() -> None:
import argparse
parser = argparse.ArgumentParser(description="BUGFS MCP Server")
parser.add_argument("--host", default="0.0.0.0", help="Host to bind (default: 0.0.0.0)")
parser.add_argument("--port", type=int, default=8000, help="Port to bind (default: 8000)")
args = parser.parse_args()
server.run(transport="sse", host=args.host, port=args.port)
if __name__ == "__main__":
main()