coderay 1.0.0__py3-none-any.whl

coderay/pipeline/watcher.py

@@ -0,0 +1,318 @@
+from __future__ import annotations
+
+import logging
+import threading
+import time
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+import pathspec
+from watchdog.events import (
+    EVENT_TYPE_CREATED,
+    EVENT_TYPE_DELETED,
+    EVENT_TYPE_MODIFIED,
+    EVENT_TYPE_MOVED,
+    FileSystemEvent,
+    FileSystemMovedEvent,
+)
+from watchdog.observers import Observer
+from watchdog.observers.polling import PollingObserver
+
+from coderay.chunking.registry import get_supported_extensions
+from coderay.vcs.git import load_gitignore
+
+logger = logging.getLogger(__name__)
+
+
+class _DebouncedHandler:
+    """Accumulates filesystem events and flushes after a quiet window."""
+
+    def __init__(
+        self,
+        repo_root: Path,
+        index_dir: Path,
+        gitignore_spec: pathspec.PathSpec,
+        supported_extensions: set[str],
+        debounce_seconds: float,
+        branch_switch_threshold: int,
+        extra_exclude: list[str],
+        on_batch: Callable[[set[str], set[str]], None],
+    ) -> None:
+        """Initialize the debounced event handler."""
+        self._repo_root = repo_root
+        self._index_dir = index_dir
+        self._gitignore = gitignore_spec
+        self._extensions = supported_extensions
+        self._debounce = debounce_seconds
+        self._threshold = branch_switch_threshold
+        self._on_batch = on_batch
+
+        extra_spec = pathspec.PathSpec.from_lines("gitignore", extra_exclude)
+        self._extra_spec = extra_spec
+
+        self._lock = threading.Lock()
+        self._changed: set[str] = set()
+        self._removed: set[str] = set()
+        self._timer: threading.Timer | None = None
+
+    # -- public (called from watchdog observer thread) -----------------
+
+    def on_event(self, event: FileSystemEvent) -> None:
+        """Handle a single filesystem event from watchdog."""
+        if event.is_directory:
+            return
+
+        paths = self._event_paths(event)
+        for abs_path in paths:
+            rel = self._relative(abs_path)
+            if rel is None or not self._should_index(rel, abs_path):
+                continue
+
+            with self._lock:
+                if event.event_type == EVENT_TYPE_DELETED:
+                    self._removed.add(rel)
+                    self._changed.discard(rel)
+                elif event.event_type == EVENT_TYPE_MOVED:
+                    assert isinstance(event, FileSystemMovedEvent)
+                    old_rel = self._relative(event.src_path)
+                    if old_rel:
+                        self._removed.add(old_rel)
+                        self._changed.discard(old_rel)
+                    self._changed.add(rel)
+                    self._removed.discard(rel)
+                else:
+                    self._changed.add(rel)
+                    self._removed.discard(rel)
+
+            self._reset_timer()
+
+    def flush_now(self) -> None:
+        """Force-flush any pending events (used during shutdown)."""
+        with self._lock:
+            if self._timer is not None:
+                self._timer.cancel()
+                self._timer = None
+        self._flush()
+
+    @property
+    def pending_count(self) -> int:
+        """Number of accumulated events not yet flushed."""
+        with self._lock:
+            return len(self._changed) + len(self._removed)
+
+    # -- internal ------------------------------------------------------
+
+    def _event_paths(self, event: FileSystemEvent) -> list[str]:
+        """Extract relevant absolute path(s) from an event."""
+        if isinstance(event, FileSystemMovedEvent):
+            return [event.dest_path]
+        return [event.src_path]
+
+    def _relative(self, abs_path: str) -> str | None:
+        """Convert an absolute path to a repo-relative string, or None."""
+        try:
+            return str(Path(abs_path).relative_to(self._repo_root))
+        except ValueError:
+            return None
+
+    def _should_index(self, rel_path: str, abs_path: str) -> bool:
+        """Return True if the path is indexable (right extension, not ignored)."""
+        parts = Path(rel_path).parts
+        if ".git" in parts:
+            return False
+
+        try:
+            Path(abs_path).relative_to(self._index_dir)
+            return False
+        except ValueError:
+            pass
+
+        if Path(abs_path).suffix not in self._extensions:
+            return False
+
+        if self._gitignore.match_file(rel_path):
+            return False
+
+        if self._extra_spec.match_file(rel_path):
+            return False
+
+        return True
+
+    def _reset_timer(self) -> None:
+        """Cancel existing timer and start a new debounce window."""
+        with self._lock:
+            if self._timer is not None:
+                self._timer.cancel()
+            self._timer = threading.Timer(self._debounce, self._flush)
+            self._timer.daemon = True
+            self._timer.start()
+
+    def _flush(self) -> None:
+        """Drain accumulated paths and invoke the batch callback."""
+        with self._lock:
+            changed = self._changed.copy()
+            removed = self._removed.copy()
+            self._changed.clear()
+            self._removed.clear()
+            self._timer = None
+
+        if not changed and not removed:
+            return
+
+        total = len(changed) + len(removed)
+        if total >= self._threshold:
+            logger.info(
+                "Branch switch detected (%d files); delegating to full sync",
+                total,
+            )
+
+        try:
+            self._on_batch(changed, removed)
+        except Exception:
+            logger.exception("Batch update failed")
+
+
+class FileWatcher:
+    """Watches a repository for changes and triggers index updates."""
+
+    def __init__(
+        self,
+        repo_root: Path,
+        index_dir: Path,
+        config: dict[str, Any] | None = None,
+        on_batch: Callable[[set[str], set[str]], None] | None = None,
+        *,
+        use_polling: bool = False,
+    ) -> None:
+        """Initialize the file watcher."""
+        from coderay.core.config import load_config
+
+        self._repo_root = repo_root.resolve()
+        self._index_dir = index_dir.resolve()
+        self._config = config or load_config(index_dir)
+        self._on_batch = on_batch
+        self._use_polling = use_polling
+
+        watch_cfg = self._config.get("watch") or {}
+        self._debounce = float(watch_cfg.get("debounce_seconds", 2))
+        self._threshold = int(watch_cfg.get("branch_switch_threshold", 50))
+        self._extra_exclude: list[str] = list(watch_cfg.get("exclude_patterns") or [])
+
+        self._observer: Observer | PollingObserver | None = None
+        self._handler: _DebouncedHandler | None = None
+        self._update_count = 0
+
+    @property
+    def update_count(self) -> int:
+        """Number of batch updates executed since start."""
+        return self._update_count
+
+    def start(self) -> None:
+        """Start the filesystem observer (non-blocking)."""
+        gitignore_spec = load_gitignore(self._repo_root)
+        extensions = get_supported_extensions()
+
+        batch_fn = self._on_batch or self._default_batch
+        self._handler = _DebouncedHandler(
+            repo_root=self._repo_root,
+            index_dir=self._index_dir,
+            gitignore_spec=gitignore_spec,
+            supported_extensions=extensions,
+            debounce_seconds=self._debounce,
+            branch_switch_threshold=self._threshold,
+            extra_exclude=self._extra_exclude,
+            on_batch=batch_fn,
+        )
+
+        adapter = _WatchdogAdapter(self._handler)
+        if self._use_polling:
+            self._observer = PollingObserver(timeout=1)
+        else:
+            self._observer = Observer()
+
+        self._observer.schedule(adapter, str(self._repo_root), recursive=True)
+        self._observer.daemon = True
+
+        try:
+            self._observer.start()
+        except SystemError:
+            logger.warning("Native observer failed; falling back to polling")
+            self._observer = PollingObserver(timeout=1)
+            self._observer.schedule(adapter, str(self._repo_root), recursive=True)
+            self._observer.daemon = True
+            self._observer.start()
+
+        logger.info(
+            "Watching %s (debounce=%.1fs, extensions=%d, observer=%s)",
+            self._repo_root,
+            self._debounce,
+            len(extensions),
+            type(self._observer).__name__,
+        )
+
+    def stop(self) -> None:
+        """Stop the observer and flush remaining events."""
+        if self._handler is not None:
+            self._handler.flush_now()
+        if self._observer is not None:
+            self._observer.stop()
+            self._observer.join(timeout=5)
+        logger.info("Watcher stopped (%d updates total)", self._update_count)
+
+    def wait(self, timeout: float | None = None) -> None:
+        """Block until the observer exits or timeout elapses."""
+        if self._observer is not None:
+            self._observer.join(timeout=timeout)
+
+    def _default_batch(self, changed: set[str], removed: set[str]) -> None:
+        """Default callback: acquire lock and run Indexer.update_paths."""
+        from coderay.core.lock import acquire_indexer_lock
+        from coderay.pipeline.indexer import Indexer
+
+        total = len(changed) + len(removed)
+        t0 = time.time()
+
+        try:
+            with acquire_indexer_lock(self._index_dir, timeout=30):
+                indexer = Indexer(
+                    self._repo_root,
+                    self._index_dir,
+                    config=self._config,
+                )
+                if total >= self._threshold:
+                    result = indexer.update_incremental()
+                else:
+                    result = indexer.update_paths(
+                        changed=sorted(changed),
+                        removed=sorted(removed),
+                    )
+            elapsed = time.time() - t0
+            self._update_count += 1
+            logger.info(
+                "Update #%d: %s (%.2fs) [%d changed, %d removed]",
+                self._update_count,
+                result,
+                elapsed,
+                len(changed),
+                len(removed),
+            )
+        except Exception:
+            logger.exception("Failed to update index")
+
+
+class _WatchdogAdapter:
+    """Bridges watchdog's event handler protocol to our ``_DebouncedHandler``."""
+
+    def __init__(self, handler: _DebouncedHandler) -> None:
+        self._handler = handler
+
+    def dispatch(self, event: FileSystemEvent) -> None:
+        """Forward relevant events to the debounced handler."""
+        if event.event_type in (
+            EVENT_TYPE_CREATED,
+            EVENT_TYPE_MODIFIED,
+            EVENT_TYPE_DELETED,
+            EVENT_TYPE_MOVED,
+        ):
+            self._handler.on_event(event)

coderay/retrieval/__init__.py

@@ -0,0 +1,3 @@
+from coderay.retrieval.search import Retrieval
+
+__all__ = ["Retrieval"]

coderay/retrieval/boosting.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from typing import Any
+
+DEFAULT_PENALTIES: list[dict[str, Any]] = [
+    {"pattern": r"(^|/)tests?/", "factor": 0.5},
+    {"pattern": r"(^|/)test_[^/]+\.py$", "factor": 0.5},
+    {"pattern": r"(^|/)(mock|fixture|conftest)", "factor": 0.4},
+    {"pattern": r"(^|/)(generated|vendor|third_party)/", "factor": 0.3},
+    {"pattern": r"(^|/)docs?/", "factor": 0.6},
+    {"pattern": r"(^|/)examples?/", "factor": 0.7},
+]
+
+DEFAULT_BONUSES: list[dict[str, Any]] = [
+    {"pattern": r"(^|/)src/", "factor": 1.1},
+    {"pattern": r"(^|/)lib/", "factor": 1.1},
+    {"pattern": r"(^|/)app/", "factor": 1.1},
+]
+
+
+@dataclass
+class BoostRule:
+    """Single path-based score multiplier rule (regex pattern + factor)."""
+
+    regex: re.Pattern[str]
+    factor: float
+
+
+@dataclass
+class StructuralBooster:
+    """Applies path-based score multipliers to search results."""
+
+    penalties: list[BoostRule] = field(default_factory=list)
+    bonuses: list[BoostRule] = field(default_factory=list)
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> StructuralBooster:
+        """Build from config dict. Falls back to defaults if not specified."""
+        search_cfg = config.get("search") or {}
+        boost_cfg = search_cfg.get("boost_rules") or {}
+
+        raw_penalties = boost_cfg.get("penalties") or DEFAULT_PENALTIES
+        raw_bonuses = boost_cfg.get("bonuses") or DEFAULT_BONUSES
+
+        return cls(
+            penalties=[
+                BoostRule(regex=re.compile(r["pattern"]), factor=r["factor"])
+                for r in raw_penalties
+            ],
+            bonuses=[
+                BoostRule(regex=re.compile(r["pattern"]), factor=r["factor"])
+                for r in raw_bonuses
+            ],
+        )
+
+    def _compute_multiplier(self, path: str) -> float:
+        multiplier = 1.0
+        for rule in self.penalties:
+            if rule.regex.search(path):
+                multiplier *= rule.factor
+        for rule in self.bonuses:
+            if rule.regex.search(path):
+                multiplier *= rule.factor
+        return multiplier
+
+    def boost(self, results: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        """Apply path-based score multipliers and re-sort results."""
+        boosted = []
+        for r in results:
+            row = dict(r)
+            path = row.get("path", "")
+            mult = self._compute_multiplier(path)
+            raw_score = row.get("score", 0.0)
+            row["raw_score"] = raw_score
+            row["score"] = raw_score * mult
+            boosted.append(row)
+        boosted.sort(key=lambda r: r["score"], reverse=True)
+        return boosted

coderay/retrieval/search.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import logging
+import time
+from pathlib import Path
+from typing import Any
+
+from coderay.core.config import get_embedding_dimensions, load_config
+from coderay.embedding.base import Embedder, load_embedder_from_config
+from coderay.graph.builder import load_graph
+from coderay.retrieval.boosting import StructuralBooster
+from coderay.state.machine import IndexMeta
+from coderay.state.version import check_index_version
+from coderay.storage.lancedb import Store, index_exists
+
+logger = logging.getLogger(__name__)
+
+
+class Retrieval:
+    """Query interface for the semantic index."""
+
+    def __init__(
+        self,
+        index_dir: str | Path,
+        config: dict[str, Any] | None = None,
+        embedder: Embedder | None = None,
+    ) -> None:
+        """Initialize retrieval for the given index."""
+        self.index_dir = Path(index_dir)
+        self._config = config or load_config(self.index_dir)
+        self._explicit_embedder = embedder
+        self._lazy_embedder: Embedder | None = None
+        self._dimensions = get_embedding_dimensions(self._config)
+        self._booster = StructuralBooster.from_config(self._config)
+        self._store: Store | None = None
+        check_index_version(self.index_dir)
+
+    @property
+    def _embedder(self) -> Embedder:
+        if self._explicit_embedder is not None:
+            return self._explicit_embedder
+        if self._lazy_embedder is None:
+            self._lazy_embedder = load_embedder_from_config(self._config)
+        return self._lazy_embedder
+
+    @property
+    def config(self) -> dict[str, Any]:
+        return self._config
+
+    def _get_store(self) -> Store:
+        if self._store is None:
+            self._store = Store(self.index_dir, dimensions=self._dimensions)
+        return self._store
+
+    def search(
+        self,
+        query: str,
+        current_state: IndexMeta,
+        *,
+        top_k: int = 10,
+        path_prefix: str | None = None,
+        language: str | None = None,
+    ) -> list[dict[str, Any]]:
+        """Semantic search over the index."""
+        if not index_exists(self.index_dir):
+            logger.warning("No index at %s", self.index_dir)
+            return []
+
+        if current_state.is_incomplete() or current_state.is_in_progress():
+            raise RuntimeError("Meta in progress; index might be stale")
+
+        store = self._get_store()
+
+        t0 = time.perf_counter()
+        query_vectors = self._embedder.embed([query])
+        logger.info("Query embed took %.3fs", time.perf_counter() - t0)
+
+        if not query_vectors:
+            return []
+        t1 = time.perf_counter()
+        results = store.search(
+            query_embedding=query_vectors[0],
+            top_k=top_k,
+            path_prefix=path_prefix,
+            language=language,
+            query_text=query,
+        )
+        logger.info("Vector search took %.3fs", time.perf_counter() - t1)
+
+        return self._booster.boost(results)
+
+    def load_graph(self) -> list[dict[str, Any]]:
+        """Load graph edges from index_dir/graph.json. [] if missing."""
+        graph = load_graph(self.index_dir)
+        if graph is None:
+            return []
+        data = graph.to_dict()
+        return data.get("edges", [])
+
+    def chunk_count(self) -> int:
+        """Total number of chunks in the index. Returns 0 if no index."""
+        if not index_exists(self.index_dir):
+            return 0
+        return self._get_store().chunk_count()
+
+    def list_chunks(
+        self,
+        *,
+        limit: int = 500,
+        path_prefix: str | None = None,
+    ) -> list[dict[str, Any]]:
+        """List indexed chunks (no vectors). For inspection/debugging."""
+        if not index_exists(self.index_dir):
+            return []
+        return self._get_store().list_chunks(limit=limit, path_prefix=path_prefix)
+
+    def chunks_by_path(self) -> dict[str, int]:
+        """Return mapping of file path -> chunk count. Empty if no index."""
+        if not index_exists(self.index_dir):
+            return {}
+        return self._get_store().chunks_by_path()

coderay/skeleton/extractor.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+from coderay.chunking.registry import get_language_for_file
+
+logger = logging.getLogger(__name__)
+
+
+def extract_skeleton(
+    path: str | Path,
+    content: str,
+    language: str | None = None,
+) -> str:
+    """Extract the skeleton of a source file (signatures, no bodies)."""
+    path_str = str(path)
+    lang_cfg = get_language_for_file(path_str) if language is None else None
+    if lang_cfg is None and language:
+        from coderay.chunking.registry import LANGUAGE_REGISTRY
+
+        lang_cfg = LANGUAGE_REGISTRY.get(language)
+    if lang_cfg is None:
+        lang_cfg = get_language_for_file(path_str)
+    if lang_cfg is None:
+        return content
+
+    try:
+        parser = lang_cfg.get_parser()
+    except Exception:
+        return content
+
+    source_bytes = content.encode("utf-8")
+    tree = parser.parse(source_bytes)
+    lines: list[str] = []
+    _visit_skeleton(tree.root_node, source_bytes, lang_cfg.name, lines, depth=0)
+    return "\n".join(lines)
+
+
+def _node_text(node, source_bytes: bytes) -> str:
+    return source_bytes[node.start_byte : node.end_byte].decode(
+        "utf-8", errors="replace"
+    )
+
+
+def _get_docstring(node, source_bytes: bytes) -> str | None:
+    """Extract docstring from the first child of the body block, if present."""
+    if not hasattr(node, "children"):
+        return None
+    body = None
+    for child in node.children:
+        if child.type in ("block", "statement_block"):
+            body = child
+            break
+    if body is None:
+        return None
+    for child in body.children:
+        if child.type == "expression_statement":
+            for sub in child.children:
+                if sub.type == "string":
+                    return _node_text(sub, source_bytes)
+        break
+    return None
+
+
+def _get_signature_line(node, source_bytes: bytes) -> str:
+    """Get everything up to and including the colon/opening brace."""
+    text = _node_text(node, source_bytes)
+    for delimiter in (":\n", "{\n", ":\r\n", "{\r\n"):
+        idx = text.find(delimiter)
+        if idx >= 0:
+            return text[: idx + 1]
+    first_line = text.split("\n")[0]
+    return first_line
+
+
+def _visit_skeleton(
+    node,
+    source_bytes: bytes,
+    language: str,
+    lines: list[str],
+    depth: int,
+) -> None:
+    indent = "    " * depth
+    ntype = node.type
+
+    if ntype in ("import_statement", "import_from_statement", "import_declaration"):
+        lines.append(indent + _node_text(node, source_bytes).strip())
+        return
+
+    if ntype in (
+        "function_definition",
+        "function_declaration",
+        "method_declaration",
+        "method_definition",
+    ):
+        sig = _get_signature_line(node, source_bytes).strip()
+        lines.append(indent + sig)
+        docstring = _get_docstring(node, source_bytes)
+        if docstring:
+            lines.append(indent + "    " + docstring)
+        lines.append(indent + "    ...")
+        lines.append("")
+        return
+
+    if ntype in (
+        "class_definition",
+        "class_declaration",
+        "interface_declaration",
+        "type_alias_declaration",
+        "type_declaration",
+    ):
+        sig = _get_signature_line(node, source_bytes).strip()
+        lines.append(indent + sig)
+        docstring = _get_docstring(node, source_bytes)
+        if docstring:
+            lines.append(indent + "    " + docstring)
+        for child in node.children:
+            if child.type in ("block", "class_body", "statement_block"):
+                for member in child.children:
+                    _visit_skeleton(member, source_bytes, language, lines, depth + 1)
+        lines.append("")
+        return
+
+    if ntype in ("expression_statement",) and depth == 0:
+        text = _node_text(node, source_bytes).strip()
+        if text.startswith(('"""', "'''", '"', "'")):
+            lines.append(indent + text)
+        elif "=" in text:
+            lines.append(indent + text)
+        return
+
+    if ntype in ("export_statement", "lexical_declaration") and depth == 0:
+        text = _node_text(node, source_bytes)
+        first_line = text.split("\n")[0].strip()
+        lines.append(indent + first_line)
+        return
+
+    for child in node.children:
+        _visit_skeleton(child, source_bytes, language, lines, depth)

coderay/state/__init__.py

@@ -0,0 +1,8 @@
+from coderay.state.machine import (
+    CurrentRun,
+    IndexMeta,
+    MetaState,
+    StateMachine,
+)
+
+__all__ = ["CurrentRun", "IndexMeta", "MetaState", "StateMachine"]