codrspot-processor-mcp 0.1.0__py3-none-any.whl

codepreproc_client/layer2_tooling/retrieval/graph_walk.py

@@ -0,0 +1,47 @@
+"""Dependency expansion over the AST graph."""
+
+from __future__ import annotations
+
+from collections import deque
+from collections.abc import Callable
+
+from codepreproc_client.layer2_tooling.condensation.ast_graph import AstGraph
+from contracts.dtos import CodeChunk
+
+
+class GraphWalker:
+    """Expand from seed chunks to nearby callers and callees."""
+
+    def __init__(self, max_chunks: int = 30) -> None:
+        self.max_chunks = max_chunks
+
+    def expand(
+        self,
+        seed_chunks: list[CodeChunk],
+        ast_graph: AstGraph,
+        chunk_loader: Callable[[list[str]], list[CodeChunk]],
+        depth: int,
+    ) -> list[CodeChunk]:
+        """Return extra chunks up to a fixed graph depth."""
+
+        symbol_map, callees_map, callers_map = ast_graph.load_for_walk_cached()
+
+        seen = {chunk.chunk_id for chunk in seed_chunks}
+        queue = deque((chunk.chunk_id, chunk.symbol, 0) for chunk in seed_chunks)
+        discovered: list[str] = []
+
+        while queue and len(discovered) < self.max_chunks:
+            chunk_id, symbol, level = queue.popleft()
+            if level >= depth:
+                continue
+            neighbors = callees_map.get(chunk_id, []) + callers_map.get(symbol, [])
+            for neighbor in neighbors:
+                if neighbor in seen:
+                    continue
+                seen.add(neighbor)
+                discovered.append(neighbor)
+                nb_sym = symbol_map.get(neighbor) or symbol
+                queue.append((neighbor, nb_sym, level + 1))
+                if len(discovered) >= self.max_chunks:
+                    break
+        return chunk_loader(discovered)

codepreproc_client/layer2_tooling/retrieval/hybrid.py

@@ -0,0 +1,95 @@
+"""Hybrid retrieval using dense plus BM25 fusion."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from codepreproc_client.layer2_tooling.condensation.bm25_store import BM25Store
+from codepreproc_client.layer2_tooling.condensation.embeddings import Embedder
+from codepreproc_client.layer2_tooling.condensation.qdrant_store import QdrantStore
+from contracts.dtos import ProjectConfig, RetrievalConfig, ScoredChunk
+
+if TYPE_CHECKING:
+    from codepreproc_client.layer2_tooling.condensation.ast_graph import AstGraph
+
+
+class HybridRetriever:
+    """Fuse dense and lexical retrieval with RRF."""
+
+    def __init__(self, embedder: Embedder, qdrant_store: QdrantStore, bm25_store: BM25Store, config: RetrievalConfig) -> None:
+        self.embedder = embedder
+        self.qdrant_store = qdrant_store
+        self.bm25_store = bm25_store
+        self.config = config
+
+    async def retrieve(
+        self,
+        project: ProjectConfig,
+        query: str,
+        overlay_collection: str | None,
+        *,
+        collection_name: str | None = None,
+        user_collection_name: str | None = None,
+        monitor=None,
+        ast_graph: "AstGraph | None" = None,
+        md_harness_top_k: int = 3,
+    ) -> list[ScoredChunk]:
+        """Retrieve candidate chunks, optionally prepending directory-scoped MD harness."""
+
+        query_vec = (await self.embedder.encode_batch([query], monitor=monitor, stage_name="retrieve"))[0]
+        dense_results = []
+        active_collection = collection_name or project.qdrant_collection
+        if active_collection:
+            dense_results = self.qdrant_store.search(active_collection, query_vec, self.config.dense_top_k)
+        if user_collection_name and self.qdrant_store.collection_exists(user_collection_name):
+            dense_results.extend(self.qdrant_store.search(user_collection_name, query_vec, self.config.dense_top_k))
+        if overlay_collection:
+            dense_results.extend(self.qdrant_store.search(overlay_collection, query_vec, self.config.dense_top_k))
+        bm25_results = self.bm25_store.search(query, self.config.bm25_top_k)
+        results = self._rrf_fuse(dense_results, bm25_results)[: self.config.rerank_top_k]
+
+        if ast_graph is not None and md_harness_top_k > 0:
+            harness = self._get_md_harness(ast_graph, results, top_k=md_harness_top_k)
+            if harness:
+                existing_ids = {item.chunk.chunk_id for item in results}
+                new_harness = [h for h in harness if h.chunk.chunk_id not in existing_ids]
+                results = new_harness + results
+
+        return results
+
+    def _get_md_harness(self, ast_graph: "AstGraph", results: list[ScoredChunk], top_k: int = 3) -> list[ScoredChunk]:
+        """Fetch directory-scoped MD knowledge chunks for the retrieved results."""
+        file_paths = [item.chunk.file_path for item in results if item.chunk.language != 'markdown']
+        if not file_paths:
+            return []
+        dir_paths = self._dir_hierarchy(file_paths)
+        chunk_ids = ast_graph.get_md_chunk_ids_for_dirs(dir_paths)
+        if not chunk_ids:
+            return []
+        md_chunks = self.bm25_store.get_by_ids(chunk_ids[:top_k])
+        return [ScoredChunk(chunk=c, score=1.0, source="md_harness") for c in md_chunks]
+
+    @staticmethod
+    def _dir_hierarchy(file_paths: list[str]) -> list[str]:
+        """Return all unique ancestor directory paths for the given files."""
+        dirs: set[str] = set()
+        for fp in file_paths:
+            parts = Path(fp).parent.parts
+            dirs.add('.')
+            for i in range(1, len(parts) + 1):
+                dirs.add('/'.join(parts[:i]))
+        return sorted(dirs)
+
+    @staticmethod
+    def _rrf_fuse(*rank_lists: list[ScoredChunk], k: int = 60) -> list[ScoredChunk]:
+        scores: dict[str, float] = defaultdict(float)
+        chunks: dict[str, ScoredChunk] = {}
+        for rank_list in rank_lists:
+            for rank, item in enumerate(rank_list, start=1):
+                scores[item.chunk.chunk_id] += 1.0 / (k + rank)
+                chunks[item.chunk.chunk_id] = item
+        fused = [ScoredChunk(chunk=chunks[key].chunk, score=value, source="rrf") for key, value in scores.items()]
+        fused.sort(key=lambda item: item.score, reverse=True)
+        return fused

codepreproc_client/layer2_tooling/retrieval/rerank_worker.py

@@ -0,0 +1,86 @@
+"""Cross-encoder reranker subprocess worker.
+
+Launched as ``python -u -m codepreproc_client.layer2_tooling.retrieval.rerank_worker``.
+Communicates with the parent via stdin/stdout JSON lines.
+
+Protocol
+--------
+Parent → worker (stdin):
+    First line: {"model_name": "...", "device": "..."}
+    Subsequent: {"type": "rerank", "request_id": "...", "pairs": [[q, t], ...], "batch_size": N}
+             or {"type": "shutdown"}
+
+Worker → parent (stdout):
+    {"type": "worker_ready"}
+    {"type": "worker_failed", "error": "..."}
+    {"type": "rerank_done", "request_id": "...", "scores": [...]}
+    {"type": "failed", "request_id": "...", "error": "..."}
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+
+
+def _send(event: dict) -> None:
+    sys.stdout.write(json.dumps(event) + "\n")
+    sys.stdout.flush()
+
+
+def main() -> None:
+    config_line = sys.stdin.readline()
+    if not config_line:
+        return
+    try:
+        config = json.loads(config_line)
+    except json.JSONDecodeError as exc:
+        _send({"type": "worker_failed", "error": f"bad config line: {exc}"})
+        return
+
+    model_name = config["model_name"]
+    device = config["device"]
+
+    try:
+        from sentence_transformers import CrossEncoder
+
+        model = CrossEncoder(model_name, device=device)
+        _send({"type": "worker_ready"})
+    except Exception as exc:
+        _send({"type": "worker_failed", "error": str(exc)})
+        return
+
+    for raw_line in sys.stdin:
+        line = raw_line.strip()
+        if not line:
+            continue
+        try:
+            command = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+
+        if command["type"] == "shutdown":
+            return
+        if command["type"] != "rerank":
+            continue
+
+        request_id = command["request_id"]
+        pairs = command["pairs"]
+        batch_size = max(1, int(command["batch_size"]))
+
+        try:
+            while True:
+                try:
+                    scores = model.predict(pairs, batch_size=batch_size, show_progress_bar=False)
+                    break
+                except RuntimeError as exc:
+                    if "out of memory" not in str(exc).lower() or batch_size == 1:
+                        raise
+                    batch_size = max(1, batch_size // 2)
+            _send({"type": "rerank_done", "request_id": request_id, "scores": [float(s) for s in scores]})
+        except Exception as exc:
+            _send({"type": "failed", "request_id": request_id, "error": str(exc)})
+
+
+if __name__ == "__main__":
+    main()

codepreproc_client/layer2_tooling/retrieval/reranker.py

@@ -0,0 +1,213 @@
+"""Cross-encoder reranker.
+
+On non-Windows platforms the model is loaded in a thread pool via
+``asyncio.to_thread``.  On Windows, loading PyTorch models inside a thread
+pool deadlocks due to the Windows DLL Loader Lock, so we launch a dedicated
+``subprocess.Popen`` pointing at ``codepreproc_client.layer2_tooling.retrieval.rerank_worker`` which
+loads the model in its own main thread.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import queue
+import subprocess
+import sys
+import threading
+import uuid
+from typing import Any, ClassVar
+
+from codepreproc_client.layer1_business.mcp.action_monitor import ToolExecutionError
+from contracts.dtos import DefaultsConfig, ScoredChunk
+
+
+class Reranker:
+    """Lazy reranker singleton."""
+
+    _instance: ClassVar["Reranker" | None] = None
+
+    def __init__(self, defaults: DefaultsConfig) -> None:
+        self.defaults = defaults
+        self._use_subprocess_worker = os.name == "nt"
+        # --- in-process state (non-Windows) ---
+        self._model = None
+        # --- subprocess state (Windows) ---
+        self._proc: subprocess.Popen | None = None  # type: ignore[type-arg]
+        self._proc_event_queue: queue.Queue | None = None
+        self._reader_thread: threading.Thread | None = None
+        self._worker_lock = asyncio.Lock()
+
+    @classmethod
+    def get(cls, defaults: DefaultsConfig) -> "Reranker":
+        """Return singleton instance."""
+
+        if cls._instance is None:
+            cls._instance = cls(defaults)
+        return cls._instance
+
+    # ------------------------------------------------------------------
+    # Public interface
+    # ------------------------------------------------------------------
+
+    async def rerank(self, query: str, chunks: list[ScoredChunk], top_k: int) -> list[ScoredChunk]:
+        """Rerank without blocking the event loop."""
+
+        if not chunks:
+            return []
+        async with self._worker_lock:
+            if self._use_subprocess_worker:
+                return await self._rerank_subprocess(query, chunks, top_k)
+            return await asyncio.to_thread(self.rerank_sync, query, chunks, top_k)
+
+    def rerank_sync(self, query: str, chunks: list[ScoredChunk], top_k: int) -> list[ScoredChunk]:
+        """Rerank synchronously (non-Windows / direct call)."""
+
+        if not chunks:
+            return []
+        model = self._ensure_model()
+        pairs = [(query, f"{item.chunk.file_path}::{item.chunk.symbol}\n{item.chunk.body[:1500]}") for item in chunks]
+        batch_size = min(16, max(1, len(pairs)))
+        while batch_size >= 1:
+            try:
+                scores = model.predict(pairs, batch_size=batch_size, show_progress_bar=False)
+                break
+            except RuntimeError as exc:
+                if "out of memory" not in str(exc).lower() or batch_size == 1:
+                    raise
+                batch_size = max(1, batch_size // 2)
+        reranked = [ScoredChunk(chunk=item.chunk, score=float(score), source="rerank") for item, score in zip(chunks, scores, strict=True)]
+        reranked.sort(key=lambda item: item.score, reverse=True)
+        return reranked[:top_k]
+
+    # ------------------------------------------------------------------
+    # In-process model (non-Windows)
+    # ------------------------------------------------------------------
+
+    def _ensure_model(self):
+        if self._model is None:
+            from sentence_transformers import CrossEncoder
+
+            self._model = CrossEncoder(self.defaults.reranker_model, device=self.defaults.reranker_device)
+        return self._model
+
+    # ------------------------------------------------------------------
+    # Subprocess worker (Windows)
+    # ------------------------------------------------------------------
+
+    async def _rerank_subprocess(self, query: str, chunks: list[ScoredChunk], top_k: int) -> list[ScoredChunk]:
+        await self._ensure_subprocess_worker()
+        pairs = [[query, f"{item.chunk.file_path}::{item.chunk.symbol}\n{item.chunk.body[:1500]}"] for item in chunks]
+        batch_size = min(16, max(1, len(pairs)))
+        request_id = uuid.uuid4().hex
+        self._send_subprocess_command(
+            {"type": "rerank", "request_id": request_id, "pairs": pairs, "batch_size": batch_size}
+        )
+        while True:
+            event = await self._wait_subprocess_event()
+            event_type = event["type"]
+            if event.get("request_id") != request_id and event_type not in ("worker_failed", "worker_exited"):
+                continue
+            if event_type == "rerank_done":
+                scores = event["scores"]
+                reranked = [
+                    ScoredChunk(chunk=item.chunk, score=float(score), source="rerank")
+                    for item, score in zip(chunks, scores, strict=True)
+                ]
+                reranked.sort(key=lambda item: item.score, reverse=True)
+                return reranked[:top_k]
+            if event_type == "failed":
+                self._restart_subprocess_worker()
+                raise ToolExecutionError(
+                    "rerank_failed",
+                    event.get("error", "unknown rerank error")[:500],
+                    details={"stage": "rerank/graph"},
+                )
+            if event_type in ("worker_failed", "worker_exited"):
+                self._restart_subprocess_worker()
+                raise ToolExecutionError(
+                    "rerank_worker_failed",
+                    event.get("error", "worker exited unexpectedly")[:500],
+                    details={"stage": "rerank/graph"},
+                )
+
+    async def _ensure_subprocess_worker(self) -> None:
+        if self._proc is not None and self._proc.poll() is None:
+            return
+        self._restart_subprocess_worker()
+        self._proc_event_queue = queue.Queue()
+        self._proc = subprocess.Popen(
+            [sys.executable, "-u", "-m", "codepreproc_client.layer2_tooling.retrieval.rerank_worker"],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.DEVNULL,
+        )
+        config_line = (
+            json.dumps({"model_name": self.defaults.reranker_model, "device": self.defaults.reranker_device}) + "\n"
+        )
+        self._proc.stdin.write(config_line.encode())
+        self._proc.stdin.flush()
+        self._reader_thread = threading.Thread(target=self._read_subprocess_stdout, daemon=True)
+        self._reader_thread.start()
+        # Wait for worker_ready
+        while True:
+            event = await self._wait_subprocess_event()
+            event_type = event["type"]
+            if event_type == "worker_ready":
+                return
+            if event_type in ("worker_failed", "worker_exited"):
+                self._restart_subprocess_worker()
+                raise ToolExecutionError(
+                    "rerank_worker_failed",
+                    event.get("error", "worker exited unexpectedly")[:500],
+                    details={"stage": "rerank/graph"},
+                )
+
+    def _read_subprocess_stdout(self) -> None:
+        try:
+            for raw_line in self._proc.stdout:
+                line = raw_line.decode("utf-8", errors="replace").strip()
+                if line:
+                    try:
+                        self._proc_event_queue.put(json.loads(line))
+                    except json.JSONDecodeError:
+                        pass
+        except Exception:
+            pass
+        finally:
+            if self._proc_event_queue is not None:
+                self._proc_event_queue.put({"type": "worker_exited"})
+
+    async def _wait_subprocess_event(self) -> dict[str, Any]:
+        while True:
+            try:
+                return await asyncio.to_thread(self._proc_event_queue.get, True, 1.0)
+            except queue.Empty:
+                if self._proc is None or self._proc.poll() is not None:
+                    raise ToolExecutionError(
+                        "rerank_worker_failed",
+                        "reranker subprocess exited unexpectedly",
+                        details={"stage": "rerank/graph"},
+                    ) from None
+
+    def _send_subprocess_command(self, command: dict) -> None:
+        line = (json.dumps(command) + "\n").encode()
+        self._proc.stdin.write(line)
+        self._proc.stdin.flush()
+
+    def _restart_subprocess_worker(self) -> None:
+        if self._proc is not None:
+            try:
+                self._proc.stdin.close()
+            except Exception:
+                pass
+            if self._proc.poll() is None:
+                self._proc.terminate()
+                try:
+                    self._proc.wait(timeout=2)
+                except subprocess.TimeoutExpired:
+                    self._proc.kill()
+        self._proc = None
+        self._proc_event_queue = None
+        self._reader_thread = None

codepreproc_client/layer2_tooling/watcher/fs_watcher.py

@@ -0,0 +1,3 @@
+"""Filesystem watcher stub for phase 2."""
+
+# TODO: implement in phase 2