codebase-brain-mcp 0.1.0__py3-none-any.whl

codebase_brain_mcp/sources.py

@@ -0,0 +1,186 @@
+"""Find a repository on disk (cloning it first if it's a URL) and list the files worth indexing."""
+
+from __future__ import annotations
+
+import os
+import re
+import subprocess
+from dataclasses import dataclass
+from pathlib import Path
+
+LANGUAGES = {
+    ".py": "python",
+    ".pyi": "python",
+    ".md": "markdown",
+    ".mdx": "markdown",
+    ".rst": "text",
+    ".txt": "text",
+    ".js": "javascript",
+    ".jsx": "javascript",
+    ".mjs": "javascript",
+    ".cjs": "javascript",
+    ".ts": "typescript",
+    ".tsx": "typescript",
+    ".go": "go",
+    ".rs": "rust",
+    ".java": "java",
+    ".kt": "kotlin",
+    ".rb": "ruby",
+    ".php": "php",
+    ".c": "c",
+    ".h": "c",
+    ".cc": "cpp",
+    ".cpp": "cpp",
+    ".hpp": "cpp",
+    ".cs": "csharp",
+    ".swift": "swift",
+    ".sh": "shell",
+    ".bash": "shell",
+    ".sql": "sql",
+    ".html": "html",
+    ".css": "css",
+    ".scss": "css",
+    ".toml": "config",
+    ".yaml": "config",
+    ".yml": "config",
+    ".json": "config",
+    ".ini": "config",
+    ".cfg": "config",
+}
+NAMED_FILES = {"Dockerfile": "config", "Makefile": "config", "README": "text", "LICENSE": "text"}
+SKIP_DIRS = {
+    ".git",
+    "node_modules",
+    ".venv",
+    "venv",
+    "__pycache__",
+    "dist",
+    "build",
+    ".mypy_cache",
+    ".pytest_cache",
+    ".tox",
+    ".next",
+    "target",
+    "vendor",
+    ".idea",
+    ".vscode",
+}
+SKIP_FILES = {"package-lock.json", "yarn.lock", "pnpm-lock.yaml", "poetry.lock", "uv.lock", "Cargo.lock"}
+GIT_TIMEOUT_SECONDS = 300
+_URL = re.compile(r"^https://[A-Za-z0-9.-]+/[A-Za-z0-9._/-]+?(\.git)?/?$")
+
+
+class SourceError(ValueError):
+    """The source can't be used: missing path, bad URL, or a failed clone."""
+
+
+@dataclass(frozen=True)
+class SourceFile:
+    path: str
+    """Path relative to the repository root, always with forward slashes."""
+    language: str
+    text: str
+
+
+def is_url(source: str) -> bool:
+    return source.startswith(("https://", "http://", "git@", "ssh://"))
+
+
+def default_name(source: str) -> str:
+    """``https://github.com/tiangolo/fastapi.git`` and ``~/code/fastapi`` both become ``fastapi``."""
+    tail = source.rstrip("/").rsplit("/", 1)[-1]
+    tail = tail.removesuffix(".git")
+    name = re.sub(r"[^A-Za-z0-9._-]", "-", tail).strip("-.")
+    if not name:
+        raise SourceError(f"Can't work out a name from {source!r}; pass one explicitly.")
+    return name
+
+
+def resolve(source: str, clone_dir: Path) -> Path:
+    """The repository's root on disk. URLs are shallow-cloned (or refreshed) under ``clone_dir``."""
+    if not is_url(source):
+        root = Path(source).expanduser().resolve()
+        if not root.is_dir():
+            raise SourceError(f"{source!r} is not a directory.")
+        return root
+    if not _URL.match(source):
+        # Only plain https URLs: no ssh, no git's ext:: transports, no options smuggled into the URL.
+        raise SourceError("Only https:// repository URLs are supported, e.g. https://github.com/owner/repo.")
+    target = clone_dir / default_name(source)
+    # Blobless clones keep the full commit history (needed for explain_history) but fetch file
+    # contents only for the checkout and for the commits we actually inspect.
+    if (target / ".git").is_dir():
+        if run_git(target, "rev-parse", "--is-shallow-repository").strip() == "true":
+            run_git(target, "fetch", "--unshallow", "--filter=blob:none", "origin")
+        else:
+            run_git(target, "fetch", "--filter=blob:none", "origin")
+        run_git(target, "reset", "--hard", "FETCH_HEAD")
+    else:
+        clone_dir.mkdir(parents=True, exist_ok=True)
+        run_git(clone_dir, "clone", "--filter=blob:none", "--", source, str(target))
+    return target
+
+
+def head_commit(root: Path) -> str | None:
+    try:
+        return run_git(root, "rev-parse", "HEAD").strip() or None
+    except SourceError:
+        return None
+
+
+def iter_files(root: Path, max_bytes: int) -> list[SourceFile]:
+    """Text files worth indexing, sorted by path. Honours .gitignore when ``root`` is a git checkout."""
+    files = []
+    for rel in _candidate_paths(root):
+        language = language_for(rel)
+        if language is None:
+            continue
+        full = root / rel
+        try:
+            if not full.is_file() or full.is_symlink() or full.stat().st_size > max_bytes:
+                continue
+            raw = full.read_bytes()
+        except OSError:
+            continue
+        if b"\x00" in raw[:8192]:
+            continue  # binary
+        files.append(SourceFile(path=rel, language=language, text=raw.decode("utf-8", errors="replace")))
+    return sorted(files, key=lambda f: f.path)
+
+
+def language_for(rel: str) -> str | None:
+    name = rel.rsplit("/", 1)[-1]
+    if name in SKIP_FILES or name.endswith((".min.js", ".min.css", ".map")):
+        return None
+    if name in NAMED_FILES:
+        return NAMED_FILES[name]
+    return LANGUAGES.get(os.path.splitext(name)[1].lower())
+
+
+def _candidate_paths(root: Path) -> list[str]:
+    try:
+        out = run_git(root, "ls-files", "-z", "--cached", "--others", "--exclude-standard")
+        paths = [p for p in out.split("\0") if p]
+        return [p for p in paths if not SKIP_DIRS.intersection(p.split("/")[:-1])]
+    except SourceError:
+        pass
+    paths = []
+    for dirpath, dirnames, filenames in os.walk(root):
+        dirnames[:] = sorted(d for d in dirnames if d not in SKIP_DIRS and not d.startswith("."))
+        rel_dir = Path(dirpath).relative_to(root).as_posix()
+        for filename in filenames:
+            paths.append(filename if rel_dir == "." else f"{rel_dir}/{filename}")
+    return paths
+
+
+def run_git(cwd: Path, *args: str) -> str:
+    env = {**os.environ, "GIT_TERMINAL_PROMPT": "0", "GIT_ALLOW_PROTOCOL": "https:file"}
+    try:
+        done = subprocess.run(
+            ["git", *args], cwd=cwd, env=env, capture_output=True, text=True, timeout=GIT_TIMEOUT_SECONDS, check=False
+        )
+    except (OSError, subprocess.TimeoutExpired) as exc:
+        raise SourceError(f"git {args[0]} failed: {exc}") from exc
+    if done.returncode != 0:
+        raise SourceError(f"git {args[0]} failed: {done.stderr.strip()[:300]}")
+    return done.stdout

codebase_brain_mcp/store.py

@@ -0,0 +1,349 @@
+"""One SQLite file per indexed repository: files, chunks, their vectors, and an FTS5 keyword index."""
+
+from __future__ import annotations
+
+import json
+import re
+import sqlite3
+import threading
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+
+from .chunking import Chunk
+from .history import Commit, PullRequest
+
+SCHEMA = """
+CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);
+CREATE TABLE IF NOT EXISTS files (path TEXT PRIMARY KEY, sha TEXT NOT NULL, language TEXT NOT NULL);
+CREATE TABLE IF NOT EXISTS chunks (
+    id INTEGER PRIMARY KEY,
+    path TEXT NOT NULL,
+    start_line INTEGER NOT NULL,
+    end_line INTEGER NOT NULL,
+    kind TEXT NOT NULL,
+    symbol TEXT,
+    text TEXT NOT NULL,
+    vector BLOB NOT NULL,
+    params TEXT,
+    accepts_any_keyword INTEGER NOT NULL DEFAULT 0
+);
+CREATE INDEX IF NOT EXISTS chunks_path ON chunks (path);
+CREATE INDEX IF NOT EXISTS chunks_symbol ON chunks (symbol);
+CREATE TABLE IF NOT EXISTS refs (chunk_id INTEGER NOT NULL, name TEXT NOT NULL);
+CREATE INDEX IF NOT EXISTS refs_name ON refs (name);
+CREATE INDEX IF NOT EXISTS refs_chunk ON refs (chunk_id);
+CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5 (body, symbol, path, tokenize = 'porter unicode61');
+CREATE TABLE IF NOT EXISTS commits (
+    sha TEXT PRIMARY KEY,
+    author TEXT NOT NULL,
+    date TEXT NOT NULL,
+    subject TEXT NOT NULL,
+    body TEXT NOT NULL,
+    prs TEXT NOT NULL
+);
+CREATE VIRTUAL TABLE IF NOT EXISTS commits_fts USING fts5 (subject, body, sha UNINDEXED, tokenize = 'porter unicode61');
+CREATE TABLE IF NOT EXISTS pull_requests (number INTEGER PRIMARY KEY, data TEXT NOT NULL);
+"""
+SCHEMA_VERSION = "3"
+"""Bump when the index layout changes; older indexes are rebuilt automatically."""
+_WORD = re.compile(r"[A-Za-z0-9_]+")
+_CAMEL = re.compile(r"[A-Z]+(?![a-z])|[A-Z]?[a-z]+|\d+")
+
+
+@dataclass(frozen=True)
+class StoredChunk:
+    id: int
+    path: str
+    start_line: int
+    end_line: int
+    kind: str
+    symbol: str | None
+    text: str
+
+    @property
+    def citation(self) -> str:
+        return f"{self.path}:{self.start_line}-{self.end_line}"
+
+
+def identifier_terms(text: str) -> str:
+    """``getUserById`` and ``get_user_by_id`` both also yield ``get user by id``, so keywords match either style."""
+    extra = []
+    for word in _WORD.findall(text):
+        parts = [p.lower() for piece in word.split("_") for p in _CAMEL.findall(piece)]
+        if len(parts) > 1:
+            extra.extend(parts)
+    return " ".join(extra)
+
+
+def _glob_literal(text: str) -> str:
+    """Escape GLOB wildcards. GLOB, unlike LIKE, is case-sensitive, which symbol names need."""
+    return re.sub(r"([*?\[])", r"[\1]", text)
+
+
+def _match_expression(query: str) -> str:
+    """An FTS5 OR-query of the query's words and identifier parts, each quoted so nothing is FTS syntax."""
+    terms = {t.lower() for t in _WORD.findall(query)} | set(identifier_terms(query).split())
+    return " OR ".join(f'"{t}"' for t in sorted(t for t in terms if len(t) > 1))
+
+
+class Store:
+    def __init__(self, path: Path):
+        path.parent.mkdir(parents=True, exist_ok=True)
+        self.path = path
+        self._db = sqlite3.connect(path, check_same_thread=False)
+        self._db.row_factory = sqlite3.Row
+        self._upgrade()
+        self._db.executescript(SCHEMA)
+        self._lock = threading.Lock()
+        self._matrix: tuple[np.ndarray, np.ndarray] | None = None
+
+    def close(self) -> None:
+        self._db.close()
+
+    def _upgrade(self) -> None:
+        """Indexes from older versions are dropped (they're rebuilt on the next index run)."""
+        tables = {r[0] for r in self._db.execute("SELECT name FROM sqlite_master WHERE type = 'table'")}
+        if "meta" not in tables:
+            return
+        row = self._db.execute("SELECT value FROM meta WHERE key = 'schema_version'").fetchone()
+        if row and row[0] == SCHEMA_VERSION:
+            return
+        with self._db:
+            for table in ("chunks_fts", "chunks", "files", "refs", "commits_fts", "commits"):
+                self._db.execute(f"DROP TABLE IF EXISTS {table}")
+            self._db.execute("DELETE FROM meta WHERE key IN ('embed_model', 'history_head')")
+
+    @property
+    def schema_version(self) -> str:
+        return SCHEMA_VERSION
+
+    # ---- metadata -------------------------------------------------------------------------
+    def get_meta(self) -> dict[str, str]:
+        return {row["key"]: row["value"] for row in self._db.execute("SELECT key, value FROM meta")}
+
+    def set_meta(self, **values: str) -> None:
+        with self._lock, self._db:
+            self._db.executemany("INSERT OR REPLACE INTO meta (key, value) VALUES (?, ?)", values.items())
+
+    # ---- files and chunks -----------------------------------------------------------------
+    def file_hashes(self) -> dict[str, str]:
+        return {row["path"]: row["sha"] for row in self._db.execute("SELECT path, sha FROM files")}
+
+    def replace_file(self, path: str, sha: str, language: str, chunks: list[Chunk], vectors: np.ndarray) -> None:
+        with self._lock, self._db:
+            self._delete_chunks(path)
+            for chunk, vector in zip(chunks, vectors, strict=True):
+                cur = self._db.execute(
+                    "INSERT INTO chunks (path, start_line, end_line, kind, symbol, text, vector, params,"
+                    " accepts_any_keyword) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
+                    (
+                        chunk.path,
+                        chunk.start_line,
+                        chunk.end_line,
+                        chunk.kind,
+                        chunk.symbol,
+                        chunk.text,
+                        np.asarray(vector, dtype=np.float32).tobytes(),
+                        None if chunk.params is None else json.dumps(list(chunk.params)),
+                        int(chunk.accepts_any_keyword),
+                    ),
+                )
+                self._db.executemany(
+                    "INSERT INTO refs (chunk_id, name) VALUES (?, ?)", [(cur.lastrowid, name) for name in chunk.refs]
+                )
+                body = f"{chunk.text}\n{identifier_terms(chunk.text)}"
+                self._db.execute(
+                    "INSERT INTO chunks_fts (rowid, body, symbol, path) VALUES (?, ?, ?, ?)",
+                    (cur.lastrowid, body, chunk.symbol or "", f"{path} {identifier_terms(path)}"),
+                )
+            self._db.execute(
+                "INSERT OR REPLACE INTO files (path, sha, language) VALUES (?, ?, ?)", (path, sha, language)
+            )
+            self._matrix = None
+
+    def remove_file(self, path: str) -> None:
+        with self._lock, self._db:
+            self._delete_chunks(path)
+            self._db.execute("DELETE FROM files WHERE path = ?", (path,))
+            self._matrix = None
+
+    def clear(self) -> None:
+        with self._lock, self._db:
+            self._db.execute("DELETE FROM chunks_fts")
+            self._db.execute("DELETE FROM refs")
+            self._db.execute("DELETE FROM chunks")
+            self._db.execute("DELETE FROM files")
+            self._db.execute("DELETE FROM commits_fts")
+            self._db.execute("DELETE FROM commits")
+            self._matrix = None
+
+    def _delete_chunks(self, path: str) -> None:
+        self._db.execute("DELETE FROM chunks_fts WHERE rowid IN (SELECT id FROM chunks WHERE path = ?)", (path,))
+        self._db.execute("DELETE FROM refs WHERE chunk_id IN (SELECT id FROM chunks WHERE path = ?)", (path,))
+        self._db.execute("DELETE FROM chunks WHERE path = ?", (path,))
+
+    def counts(self) -> dict[str, int]:
+        files = self._db.execute("SELECT COUNT(*) FROM files").fetchone()[0]
+        chunks = self._db.execute("SELECT COUNT(*) FROM chunks").fetchone()[0]
+        return {"files": files, "chunks": chunks}
+
+    def languages(self) -> dict[str, int]:
+        rows = self._db.execute("SELECT language, COUNT(*) AS n FROM files GROUP BY language ORDER BY n DESC")
+        return {row["language"]: row["n"] for row in rows}
+
+    def get_chunks(self, ids: list[int]) -> dict[int, StoredChunk]:
+        if not ids:
+            return {}
+        marks = ",".join("?" * len(ids))
+        rows = self._db.execute(
+            f"SELECT id, path, start_line, end_line, kind, symbol, text FROM chunks WHERE id IN ({marks})",  # noqa: S608 (only "?" placeholders)
+            ids,
+        )
+        return {row["id"]: StoredChunk(**dict(row)) for row in rows}
+
+    def find_symbol(self, symbol: str, path: str | None = None) -> list[StoredChunk]:
+        """Chunks whose symbol is ``symbol`` or ends with ``.symbol`` (so ``unsign`` finds ``Signer.unsign``)."""
+        query = (
+            "SELECT id, path, start_line, end_line, kind, symbol, text FROM chunks"  # noqa: S608 (fixed SQL)
+            " WHERE (symbol = ? OR symbol GLOB ?)" + (" AND path = ?" if path else "") + " ORDER BY start_line"
+        )
+        params = [symbol, f"*.{_glob_literal(symbol)}"] + ([path] if path else [])
+        return [StoredChunk(**dict(row)) for row in self._db.execute(query, params)]
+
+    # ---- structure (impact and drift) -----------------------------------------------------
+    def chunks_referencing(self, name: str, limit: int = 200) -> list[StoredChunk]:
+        rows = self._db.execute(
+            "SELECT DISTINCT c.id, c.path, c.start_line, c.end_line, c.kind, c.symbol, c.text FROM refs r"
+            " JOIN chunks c ON c.id = r.chunk_id WHERE r.name = ? ORDER BY c.path, c.start_line LIMIT ?",
+            (name, limit),
+        )
+        return [StoredChunk(**dict(row)) for row in rows]
+
+    def symbols(self) -> list[tuple[str, str]]:
+        """Every defined (symbol, path) pair."""
+        rows = self._db.execute(
+            "SELECT DISTINCT symbol, path FROM chunks WHERE symbol IS NOT NULL"
+            " AND kind IN ('function', 'class', 'method')"
+        )
+        return [(row["symbol"], row["path"]) for row in rows]
+
+    def signatures(self, short_name: str) -> list[tuple[str, str, list[str], bool]]:
+        """(symbol, path, params, accepts_any_keyword) for functions named ``short_name``."""
+        rows = self._db.execute(
+            "SELECT symbol, path, params, accepts_any_keyword FROM chunks WHERE params IS NOT NULL"
+            " AND (symbol = ? OR symbol GLOB ?)",
+            (short_name, f"*.{_glob_literal(short_name)}"),
+        )
+        return [(r["symbol"], r["path"], json.loads(r["params"]), bool(r["accepts_any_keyword"])) for r in rows]
+
+    def file_paths(self) -> set[str]:
+        return {row["path"] for row in self._db.execute("SELECT path FROM files")}
+
+    def doc_chunks(self, path_prefix: str | None = None) -> list[StoredChunk]:
+        query = (
+            "SELECT c.id, c.path, c.start_line, c.end_line, c.kind, c.symbol, c.text FROM chunks c"
+            " JOIN files f ON f.path = c.path WHERE f.language IN ('markdown', 'text')"
+        )
+        params: list[str] = []
+        if path_prefix:
+            query += " AND c.path LIKE ?"
+            params.append(path_prefix.replace("%", "") + "%")
+        return [StoredChunk(**dict(row)) for row in self._db.execute(query + " ORDER BY c.path, c.start_line", params)]
+
+    def code_mentions(self, term: str) -> bool:
+        """Whether ``term`` appears anywhere in non-documentation files (a plain substring check)."""
+        row = self._db.execute(
+            "SELECT 1 FROM chunks c JOIN files f ON f.path = c.path"
+            " WHERE f.language NOT IN ('markdown', 'text') AND instr(c.text, ?) > 0 LIMIT 1",
+            (term,),
+        ).fetchone()
+        return row is not None
+
+    # ---- history --------------------------------------------------------------------------
+    def add_commits(self, commits: list[Commit]) -> int:
+        added = 0
+        with self._lock, self._db:
+            for commit in commits:
+                cur = self._db.execute(
+                    "INSERT OR IGNORE INTO commits (sha, author, date, subject, body, prs) VALUES (?, ?, ?, ?, ?, ?)",
+                    (commit.sha, commit.author, commit.date, commit.subject, commit.body, json.dumps(commit.prs)),
+                )
+                if cur.rowcount:
+                    added += 1
+                    self._db.execute(
+                        "INSERT INTO commits_fts (subject, body, sha) VALUES (?, ?, ?)",
+                        (f"{commit.subject} {identifier_terms(commit.subject)}", commit.body, commit.sha),
+                    )
+        return added
+
+    def commit_count(self) -> int:
+        return self._db.execute("SELECT COUNT(*) FROM commits").fetchone()[0]
+
+    def search_commits(self, query: str, limit: int) -> list[Commit]:
+        match = _match_expression(query)
+        if not match:
+            return []
+        rows = self._db.execute(
+            "SELECT c.sha, c.author, c.date, c.subject, c.body, c.prs FROM commits_fts f"
+            " JOIN commits c ON c.sha = f.sha WHERE commits_fts MATCH ?"
+            " ORDER BY bm25(commits_fts, 3.0, 1.0) LIMIT ?",
+            (match, limit),
+        )
+        return [
+            Commit(
+                sha=r["sha"],
+                author=r["author"],
+                date=r["date"],
+                subject=r["subject"],
+                body=r["body"],
+                prs=tuple(json.loads(r["prs"])),
+            )
+            for r in rows
+        ]
+
+    def cached_pull_request(self, number: int) -> PullRequest | None:
+        row = self._db.execute("SELECT data FROM pull_requests WHERE number = ?", (number,)).fetchone()
+        return PullRequest(**json.loads(row["data"])) if row else None
+
+    def cache_pull_request(self, pr: PullRequest) -> None:
+        if pr.title is None:
+            return  # don't cache failed lookups; they may work next time
+        with self._lock, self._db:
+            self._db.execute(
+                "INSERT OR REPLACE INTO pull_requests (number, data) VALUES (?, ?)",
+                (pr.number, json.dumps(pr.__dict__)),
+            )
+
+    # ---- retrieval ------------------------------------------------------------------------
+    def keyword_search(self, query: str, limit: int) -> list[int]:
+        """Chunk ids ranked by BM25. Symbol and path matches count more than body matches."""
+        match = _match_expression(query)
+        if not match:
+            return []
+        rows = self._db.execute(
+            "SELECT rowid FROM chunks_fts WHERE chunks_fts MATCH ? ORDER BY bm25(chunks_fts, 1.0, 4.0, 2.0) LIMIT ?",
+            (match, limit),
+        )
+        return [row[0] for row in rows]
+
+    def vector_search(self, query_vector: np.ndarray, limit: int) -> list[int]:
+        """Chunk ids ranked by cosine similarity (vectors are normalised, so a dot product)."""
+        ids, matrix = self._vectors()
+        if not len(ids):
+            return []
+        scores = matrix @ query_vector.astype(np.float32)
+        top = np.argsort(-scores)[:limit]
+        return [int(ids[i]) for i in top]
+
+    def _vectors(self) -> tuple[np.ndarray, np.ndarray]:
+        if self._matrix is None:
+            rows = self._db.execute("SELECT id, vector FROM chunks ORDER BY id").fetchall()
+            if rows:
+                ids = np.array([row["id"] for row in rows])
+                matrix = np.stack([np.frombuffer(row["vector"], dtype=np.float32) for row in rows])
+            else:
+                ids, matrix = np.array([], dtype=int), np.zeros((0, 0), dtype=np.float32)
+            self._matrix = (ids, matrix)
+        return self._matrix