first commit
This commit is contained in:
jsahn
Executable
+782
@@ -0,0 +1,782 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user