projectmind-mcp 0.9.0__py3-none-any.whl

annotations.py

@@ -0,0 +1,315 @@
+"""
+File annotations — AI-authored summaries that make keyword search semantic.
+
+The design inverts the usual RAG architecture: instead of a small embedding
+model guessing what code means, the *client* LLM (which actually understands
+the code) writes a short natural-language summary + keywords per file. Those
+annotations are indexed into the BM25 corpus and served as a cheap query tier,
+so "where is authentication handled?" matches the summary text with plain
+keyword search — no vector database required.
+
+Storage: `.ai/annotations.json` (atomic writes, UTF-8, project-relative posix
+keys). Each entry remembers the file's mtime at annotation time so stale
+annotations (file changed since) can be surfaced for re-annotation.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import threading
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import config
+from incremental_indexing import atomic_write
+from logger import get_logger
+
+logger = get_logger()
+
+ANNOTATIONS_FILENAME = "annotations.json"
+ANNOTATIONS_VERSION = 1
+
+# Cap on files returned by unannotated scans (safety on huge repos)
+_SCAN_CAP = 20000
+
+_TOKEN_RE = re.compile(r"[a-z0-9_]+")
+
+
+def _annotations_path() -> Path:
+    return config.AI_DIR / ANNOTATIONS_FILENAME
+
+
+def _rel_posix(path: str | Path) -> str:
+    """Normalize any path to a project-relative posix string."""
+    p = Path(path)
+    try:
+        if p.is_absolute():
+            return p.resolve().relative_to(config.PROJECT_ROOT.resolve()).as_posix()
+        return p.as_posix()
+    except Exception:
+        return str(path).replace("\\", "/")
+
+
+def _tokenize(text: str) -> set[str]:
+    return {t for t in _TOKEN_RE.findall(text.lower()) if len(t) >= 2}
+
+
+@dataclass
+class Annotation:
+    path: str  # project-relative posix
+    summary: str
+    keywords: list[str] = field(default_factory=list)
+    annotated_mtime: float = 0.0
+    updated_at: str = ""
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "summary": self.summary,
+            "keywords": self.keywords,
+            "annotated_mtime": self.annotated_mtime,
+            "updated_at": self.updated_at,
+        }
+
+    def search_text(self) -> str:
+        """Text indexed into BM25 for this annotation."""
+        stem = Path(self.path).stem
+        return f"{self.path} {stem} {self.summary} {' '.join(self.keywords)}"
+
+    @property
+    def doc_id(self) -> str:
+        return f"annot::{self.path}"
+
+
+class AnnotationStore:
+    """Thread-safe store over `.ai/annotations.json`."""
+
+    def __init__(self, path: Path | None = None):
+        self._path = path if path is not None else _annotations_path()
+        self._lock = threading.RLock()
+        self._files: dict[str, Annotation] = {}
+        self._loaded_mtime: float = -1.0
+        self._load()
+
+    # -- persistence --------------------------------------------------------
+
+    def _load(self) -> None:
+        with self._lock:
+            if not self._path.exists():
+                self._files = {}
+                self._loaded_mtime = -1.0
+                return
+            try:
+                data = json.loads(self._path.read_text(encoding="utf-8"))
+                files = {}
+                for rel, entry in (data.get("files") or {}).items():
+                    files[rel] = Annotation(
+                        path=rel,
+                        summary=str(entry.get("summary", "")),
+                        keywords=[str(k) for k in entry.get("keywords", [])],
+                        annotated_mtime=float(entry.get("annotated_mtime", 0.0) or 0.0),
+                        updated_at=str(entry.get("updated_at", "")),
+                    )
+                self._files = files
+                self._loaded_mtime = self._path.stat().st_mtime
+            except Exception as e:
+                logger.warning(f"Could not load annotations: {e}")
+                self._files = {}
+
+    def refresh_if_changed(self) -> None:
+        """Re-read from disk when another process/session updated the file."""
+        with self._lock:
+            try:
+                current = self._path.stat().st_mtime if self._path.exists() else -1.0
+            except OSError:
+                return
+            if current != self._loaded_mtime:
+                self._load()
+
+    def _save(self) -> None:
+        with self._lock:
+            data = {
+                "version": ANNOTATIONS_VERSION,
+                "files": {rel: a.to_dict() for rel, a in sorted(self._files.items())},
+            }
+            atomic_write(self._path, json.dumps(data, indent=2, ensure_ascii=False))
+            try:
+                self._loaded_mtime = self._path.stat().st_mtime
+            except OSError:
+                pass
+
+    # -- API ----------------------------------------------------------------
+
+    def set(self, path: str, summary: str, keywords: list[str] | None = None) -> Annotation:
+        """Save/replace one file's annotation. Path must exist inside the project."""
+        target = config.validate_path(path)
+        if not target.is_file():
+            raise ValueError(f"Not a file: {path}")
+        rel = _rel_posix(target)
+        try:
+            mtime = target.stat().st_mtime
+        except OSError:
+            mtime = 0.0
+        ann = Annotation(
+            path=rel,
+            summary=summary.strip(),
+            keywords=[k.strip() for k in (keywords or []) if k.strip()],
+            annotated_mtime=mtime,
+            updated_at=datetime.now().isoformat(timespec="seconds"),
+        )
+        with self._lock:
+            self._files[rel] = ann
+            self._save()
+        return ann
+
+    def remove(self, path: str) -> bool:
+        rel = _rel_posix(path)
+        with self._lock:
+            if rel in self._files:
+                del self._files[rel]
+                self._save()
+                return True
+            return False
+
+    def get(self, path: str) -> Annotation | None:
+        self.refresh_if_changed()
+        return self._files.get(_rel_posix(path))
+
+    def all(self) -> dict[str, Annotation]:
+        self.refresh_if_changed()
+        with self._lock:
+            return dict(self._files)
+
+    def count(self) -> int:
+        self.refresh_if_changed()
+        return len(self._files)
+
+    def is_stale(self, ann: Annotation) -> bool:
+        """True when the file changed after it was annotated (or vanished)."""
+        target = config.PROJECT_ROOT / ann.path
+        try:
+            return target.stat().st_mtime != ann.annotated_mtime
+        except OSError:
+            return True
+
+    def unannotated(self, limit: int = 20) -> tuple[list[str], list[str], int]:
+        """
+        Scan indexable code files and report annotation coverage.
+
+        Returns:
+            (missing, stale, total_scanned): relative posix paths of files with
+            no annotation, files whose annotation is stale, and the number of
+            code files scanned.
+        """
+        self.refresh_if_changed()
+        missing: list[str] = []
+        stale: list[str] = []
+        scanned = 0
+        root = config.PROJECT_ROOT
+
+        for dirpath, dirnames, filenames in os.walk(root):
+            dirnames[:] = [d for d in dirnames if not config.is_dir_ignored(d)]
+            for fname in filenames:
+                fp = Path(dirpath) / fname
+                suffix = fp.suffix.lower()
+                # Annotations target source code, not configs/docs
+                if suffix not in config.CODE_EXTENSIONS:
+                    continue
+                scanned += 1
+                if scanned > _SCAN_CAP:
+                    return missing, stale, scanned
+                rel = _rel_posix(fp)
+                ann = self._files.get(rel)
+                if ann is None:
+                    if len(missing) < limit:
+                        missing.append(rel)
+                elif self.is_stale(ann):
+                    if len(stale) < limit:
+                        stale.append(rel)
+        return missing, stale, scanned
+
+    # -- search integration --------------------------------------------------
+
+    def as_documents(self) -> tuple[list[str], list[str], list[dict[str, Any]]]:
+        """Annotation docs for the BM25 corpus: (ids, texts, metadatas)."""
+        self.refresh_if_changed()
+        ids: list[str] = []
+        texts: list[str] = []
+        metas: list[dict[str, Any]] = []
+        with self._lock:
+            for ann in self._files.values():
+                if not ann.summary:
+                    continue
+                ids.append(ann.doc_id)
+                texts.append(ann.search_text())
+                metas.append(
+                    {
+                        "source": ann.path,
+                        "symbol_type": "annotation",
+                        "symbol_name": Path(ann.path).stem,
+                    }
+                )
+        return ids, texts, metas
+
+    def search(self, query: str, n: int = 8) -> list[dict[str, Any]]:
+        """
+        Cheap keyword-overlap scoring over annotations (query-tier friendly:
+        no BM25 matrix or vector store needed, runs in microseconds).
+        """
+        self.refresh_if_changed()
+        q_tokens = _tokenize(query)
+        if not q_tokens:
+            return []
+        scored: list[tuple[float, Annotation]] = []
+        with self._lock:
+            for ann in self._files.values():
+                path_tokens = _tokenize(ann.path)
+                kw_tokens = _tokenize(" ".join(ann.keywords))
+                sum_tokens = _tokenize(ann.summary)
+                score = 0.0
+                for tok in q_tokens:
+                    if tok in kw_tokens:
+                        score += 3.0
+                    if tok in path_tokens:
+                        score += 2.0
+                    if tok in sum_tokens:
+                        score += 1.0
+                matched = sum(
+                    1 for t in q_tokens if t in kw_tokens or t in path_tokens or t in sum_tokens
+                )
+                score += 2.0 * (matched / len(q_tokens))
+                if matched:
+                    scored.append((score, ann))
+        scored.sort(key=lambda x: x[0], reverse=True)
+        top = scored[:n]
+        max_score = top[0][0] if top else 1.0
+        return [
+            {
+                "source": ann.path,
+                "score": round(0.5 + 0.5 * (s / max_score), 4) if max_score else 0.5,
+                "tier": "L0_annot",
+                "snippet": ann.summary[:300],
+                "extra": {"keywords": ann.keywords, "updated_at": ann.updated_at},
+            }
+            for s, ann in top
+        ]
+
+
+# Module-level singleton (per project root; reset via get_store cache key)
+_store: AnnotationStore | None = None
+_store_path: Path | None = None
+_store_lock = threading.Lock()
+
+
+def get_store() -> AnnotationStore:
+    """Returns the AnnotationStore for the current project root."""
+    global _store, _store_path
+    with _store_lock:
+        current = _annotations_path()
+        if _store is None or _store_path != current:
+            _store = AnnotationStore(current)
+            _store_path = current
+        return _store

ast_splitter.py

@@ -0,0 +1,382 @@
+from __future__ import annotations
+
+import threading
+from pathlib import Path
+from typing import Any
+
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+from config import CHUNK_OVERLAP, CHUNK_SIZE
+from logger import get_logger
+
+logger = get_logger()
+
+LANGUAGE_MAP: dict[str, str] = {
+    ".py": "python",
+    ".js": "javascript",
+    ".jsx": "javascript",
+    ".ts": "typescript",
+    ".tsx": "tsx",
+    ".java": "java",
+    ".go": "go",
+    ".rs": "rust",
+    ".rb": "ruby",
+}
+
+TOP_LEVEL_NODES: dict[str, list[str]] = {
+    "python": ["function_definition", "class_definition", "decorated_definition"],
+    "javascript": [
+        "function_declaration",
+        "class_declaration",
+        "export_statement",
+        "lexical_declaration",
+    ],
+    "typescript": [
+        "function_declaration",
+        "class_declaration",
+        "export_statement",
+        "interface_declaration",
+        "type_alias_declaration",
+    ],
+    "tsx": [
+        "function_declaration",
+        "class_declaration",
+        "export_statement",
+        "interface_declaration",
+    ],
+    "java": ["class_declaration", "interface_declaration", "enum_declaration"],
+    "go": ["function_declaration", "method_declaration", "type_declaration"],
+    "rust": ["function_item", "impl_item", "struct_item", "enum_item", "trait_item"],
+    "ruby": ["method", "class", "module"],
+}
+
+METHOD_NODES: dict[str, list[str]] = {
+    "python": ["function_definition"],
+    "javascript": ["method_definition", "function_declaration", "arrow_function"],
+    "typescript": ["method_definition", "method_signature", "function_declaration"],
+    "tsx": ["method_definition", "method_signature", "function_declaration"],
+    "java": ["method_declaration", "constructor_declaration"],
+    "go": ["method_declaration", "function_declaration"],
+    "rust": ["function_item"],
+    "ruby": ["method"],
+}
+
+NAME_FIELD: dict[str, str] = {
+    "python": "name",
+    "javascript": "name",
+    "typescript": "name",
+    "tsx": "name",
+    "java": "name",
+    "go": "name",
+    "rust": "name",
+    "ruby": "name",
+}
+
+_parsers: dict[str, Any] = {}
+
+# tree-sitter Parser objects are NOT thread-safe: concurrent parse() calls on
+# a shared parser can crash the interpreter. All parser creation and parsing
+# is serialized through this lock (parsing is cheap next to embedding).
+_parse_lock = threading.RLock()
+
+
+def _get_parser(language: str) -> Any | None:
+    with _parse_lock:
+        return _get_parser_locked(language)
+
+
+def _get_parser_locked(language: str) -> Any | None:
+    if language in _parsers:
+        return _parsers[language]
+
+    try:
+        from tree_sitter import Language, Parser
+
+        if language == "python":
+            import tree_sitter_python as mod
+        elif language == "javascript":
+            import tree_sitter_javascript as mod
+        elif language == "typescript":
+            import tree_sitter_typescript as mod
+
+            lang = Language(mod.language_typescript())
+            parser = Parser(lang)
+            _parsers[language] = parser
+            return parser
+        elif language == "tsx":
+            import tree_sitter_typescript as mod
+
+            lang = Language(mod.language_tsx())
+            parser = Parser(lang)
+            _parsers[language] = parser
+            return parser
+        elif language == "java":
+            import tree_sitter_java as mod
+        elif language == "go":
+            import tree_sitter_go as mod
+        elif language == "rust":
+            import tree_sitter_rust as mod
+        elif language == "ruby":
+            import tree_sitter_ruby as mod
+        else:
+            return None
+
+        lang = Language(mod.language())
+        parser = Parser(lang)
+        _parsers[language] = parser
+        return parser
+
+    except Exception as e:
+        logger.warning(f"Could not load tree-sitter parser for {language}: {e}")
+        _parsers[language] = None
+        return None
+
+
+def _get_node_name(node: Any, source: bytes) -> str:
+    for child in node.children:
+        if child.type == "identifier" or child.type == "name":
+            return child.text.decode("utf-8", errors="replace")
+    if node.start_byte < len(source):
+        first_line = source[node.start_byte : node.start_byte + 80].decode(
+            "utf-8", errors="replace"
+        )
+        return first_line.split("\n")[0].strip()[:60]
+    return "unknown"
+
+
+def _extract_class_chunks(
+    class_node: Any,
+    source: bytes,
+    language: str,
+    file_path: Path,
+    text_splitter: RecursiveCharacterTextSplitter,
+) -> list[dict[str, Any]]:
+    chunks = []
+    class_name = _get_node_name(class_node, source)
+    method_types = METHOD_NODES.get(language, [])
+
+    class_text = class_node.text.decode("utf-8", errors="replace")
+    body_node = None
+    for child in class_node.children:
+        if child.type in ("block", "class_body", "declaration_list"):
+            body_node = child
+            break
+
+    if body_node is None:
+        chunks.extend(
+            _make_text_chunks(
+                class_text,
+                str(file_path),
+                "class",
+                class_name,
+                None,
+                class_node.start_point[0] + 1,
+                class_node.end_point[0] + 1,
+                text_splitter,
+            )
+        )
+        return chunks
+
+    has_methods = False
+    for child in body_node.children:
+        if child.type in method_types:
+            has_methods = True
+            method_name = _get_node_name(child, source)
+            method_text = child.text.decode("utf-8", errors="replace")
+            context_prefix = f"# Class: {class_name}\n"
+            full_text = context_prefix + method_text
+            chunks.extend(
+                _make_text_chunks(
+                    full_text,
+                    str(file_path),
+                    "method",
+                    method_name,
+                    class_name,
+                    child.start_point[0] + 1,
+                    child.end_point[0] + 1,
+                    text_splitter,
+                )
+            )
+
+    if not has_methods:
+        chunks.extend(
+            _make_text_chunks(
+                class_text,
+                str(file_path),
+                "class",
+                class_name,
+                None,
+                class_node.start_point[0] + 1,
+                class_node.end_point[0] + 1,
+                text_splitter,
+            )
+        )
+
+    return chunks
+
+
+def _make_text_chunks(
+    text: str,
+    source_path: str,
+    symbol_type: str,
+    symbol_name: str,
+    class_name: str | None,
+    line_start: int,
+    line_end: int,
+    text_splitter: RecursiveCharacterTextSplitter,
+) -> list[dict[str, Any]]:
+    if len(text) <= CHUNK_SIZE:
+        return [
+            {
+                "text": text,
+                "metadata": {
+                    "source": source_path,
+                    "symbol_type": symbol_type,
+                    "symbol_name": symbol_name,
+                    "class_name": class_name or "",
+                    "line_start": line_start,
+                    "line_end": line_end,
+                    "chunk_index": 0,
+                },
+            }
+        ]
+
+    sub_chunks = text_splitter.split_text(text)
+    result = []
+    for i, sub in enumerate(sub_chunks):
+        result.append(
+            {
+                "text": sub,
+                "metadata": {
+                    "source": source_path,
+                    "symbol_type": symbol_type,
+                    "symbol_name": symbol_name,
+                    "class_name": class_name or "",
+                    "line_start": line_start,
+                    "line_end": line_end,
+                    "chunk_index": i,
+                },
+            }
+        )
+    return result
+
+
+class ASTSplitter:
+    def __init__(self) -> None:
+        self._text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP
+        )
+
+    def split(self, content: str, file_path: Path) -> list[dict[str, Any]]:
+        language = LANGUAGE_MAP.get(file_path.suffix.lower())
+        if language:
+            parser = _get_parser(language)
+            if parser:
+                try:
+                    return self._split_by_ast(content, language, parser, file_path)
+                except Exception as e:
+                    logger.warning(f"AST split failed for {file_path}, falling back: {e}")
+
+        return self._split_by_text(content, file_path)
+
+    def _split_by_ast(
+        self, content: str, language: str, parser: Any, file_path: Path
+    ) -> list[dict[str, Any]]:
+        source = content.encode("utf-8")
+        with _parse_lock:
+            tree = parser.parse(source)
+        root = tree.root_node
+
+        top_types = TOP_LEVEL_NODES.get(language, [])
+        chunks: list[dict[str, Any]] = []
+        covered_ranges: list[tuple[int, int]] = []
+
+        for node in root.children:
+            actual_node = node
+            if node.type == "decorated_definition":
+                for child in node.children:
+                    if child.type in ("function_definition", "class_definition"):
+                        actual_node = child
+                        break
+
+            if node.type == "export_statement":
+                for child in node.children:
+                    if child.type in top_types:
+                        actual_node = child
+                        break
+
+            if actual_node.type not in top_types and node.type not in top_types:
+                continue
+
+            node_text = node.text.decode("utf-8", errors="replace") if node.text else ""
+            if not node_text.strip():
+                continue
+
+            covered_ranges.append((node.start_byte, node.end_byte))
+
+            if actual_node.type == "class_definition" or actual_node.type == "class_declaration":
+                chunks.extend(
+                    _extract_class_chunks(
+                        actual_node, source, language, file_path, self._text_splitter
+                    )
+                )
+            else:
+                symbol_name = _get_node_name(actual_node, source)
+                chunks.extend(
+                    _make_text_chunks(
+                        node_text,
+                        str(file_path),
+                        "function",
+                        symbol_name,
+                        None,
+                        node.start_point[0] + 1,
+                        node.end_point[0] + 1,
+                        self._text_splitter,
+                    )
+                )
+
+        if not chunks:
+            return self._split_by_text(content, file_path)
+
+        module_lines = []
+        for node in root.children:
+            is_covered = any(s <= node.start_byte < e for s, e in covered_ranges)
+            if not is_covered and node.text:
+                line = node.text.decode("utf-8", errors="replace").strip()
+                if line:
+                    module_lines.append(line)
+
+        if module_lines:
+            module_text = "\n".join(module_lines)
+            chunks.extend(
+                _make_text_chunks(
+                    module_text,
+                    str(file_path),
+                    "module",
+                    "module_level",
+                    None,
+                    1,
+                    root.end_point[0] + 1,
+                    self._text_splitter,
+                )
+            )
+
+        return chunks
+
+    def _split_by_text(self, content: str, file_path: Path) -> list[dict[str, Any]]:
+        sub_chunks = self._text_splitter.split_text(content)
+        return [
+            {
+                "text": chunk,
+                "metadata": {
+                    "source": str(file_path),
+                    "symbol_type": "text",
+                    "symbol_name": "",
+                    "class_name": "",
+                    "line_start": 0,
+                    "line_end": 0,
+                    "chunk_index": i,
+                },
+            }
+            for i, chunk in enumerate(sub_chunks)
+        ]