code-context-engine 0.4.0__py3-none-any.whl

context_engine/event_bus.py

@@ -0,0 +1,24 @@
+"""Simple async event bus for inter-module communication."""
+import asyncio
+from collections import defaultdict
+from typing import Any, Callable, Coroutine
+
+
+Handler = Callable[[Any], Coroutine[Any, Any, None]]
+
+
+class EventBus:
+    def __init__(self) -> None:
+        self._handlers: dict[str, list[Handler]] = defaultdict(list)
+
+    def subscribe(self, event: str, handler: Handler) -> None:
+        self._handlers[event].append(handler)
+
+    def unsubscribe(self, event: str, handler: Handler) -> None:
+        handlers = self._handlers.get(event, [])
+        if handler in handlers:
+            handlers.remove(handler)
+
+    async def emit(self, event: str, data: Any = None) -> None:
+        for handler in self._handlers.get(event, []):
+            await handler(data)

context_engine/indexer/chunker.py

@@ -0,0 +1,147 @@
+"""AST-aware code chunking using tree-sitter."""
+import hashlib
+
+import tree_sitter_python as tspython
+import tree_sitter_javascript as tsjavascript
+import tree_sitter_typescript as tstypescript
+import tree_sitter_php as tsphp
+import tree_sitter_go as tsgo
+import tree_sitter_rust as tsrust
+import tree_sitter_java as tsjava
+from tree_sitter import Language, Parser
+
+from context_engine.models import Chunk, ChunkType
+
+_FUNCTION_TYPES = {
+    "function_definition", "function_declaration",  # Python, PHP, JS
+    "method_definition", "method_declaration",       # JS/TS, PHP/Go/Java
+    "arrow_function",                                # JS/TS
+    "function_item",                                 # Rust
+}
+_CLASS_TYPES = {
+    "class_definition", "class_declaration",  # Python, JS/TS, PHP, Java
+    "type_declaration",                        # Go (struct/interface)
+    "struct_item", "impl_item", "enum_item",   # Rust
+}
+_IMPORT_TYPES = {
+    "import_statement", "import_from_statement",  # Python
+    "import_declaration",                          # TypeScript, Go, Java
+    "use_declaration",                             # PHP, Rust
+}
+
+_LANGUAGES = {
+    "python": Language(tspython.language()),
+    "javascript": Language(tsjavascript.language()),
+    "typescript": Language(tstypescript.language_typescript()),
+    "tsx": Language(tstypescript.language_tsx()),
+    "php": Language(tsphp.language_php()),
+    "go": Language(tsgo.language()),
+    "rust": Language(tsrust.language()),
+    "java": Language(tsjava.language()),
+}
+
+
+class Chunker:
+    def __init__(self) -> None:
+        self._parsers: dict[str, Parser] = {}
+
+    def _get_parser(self, language: str) -> Parser | None:
+        if language not in _LANGUAGES:
+            return None
+        if language not in self._parsers:
+            parser = Parser(_LANGUAGES[language])
+            self._parsers[language] = parser
+        return self._parsers[language]
+
+    def chunk(self, source: str, file_path: str, language: str) -> list[Chunk]:
+        parser = self._get_parser(language)
+        if parser is None:
+            return [self._fallback_chunk(source, file_path, language)]
+        tree = parser.parse(source.encode("utf-8"))
+        chunks = []
+        self._walk(tree.root_node, source, file_path, language, chunks)
+        if not chunks:
+            return [self._fallback_chunk(source, file_path, language)]
+        return chunks
+
+    def _walk(self, node, source, file_path, language, chunks):
+        if node.type in _FUNCTION_TYPES:
+            chunks.append(self._node_to_chunk(node, source, file_path, language, ChunkType.FUNCTION))
+        elif node.type in _CLASS_TYPES:
+            chunks.append(self._node_to_chunk(node, source, file_path, language, ChunkType.CLASS))
+        for child in node.children:
+            self._walk(child, source, file_path, language, chunks)
+
+    def _node_to_chunk(self, node, source, file_path, language, chunk_type):
+        content = source[node.start_byte:node.end_byte]
+        start_line = node.start_point.row + 1
+        end_line = node.end_point.row + 1
+        chunk_id = hashlib.sha256(
+            f"{file_path}:{start_line}:{end_line}:{content[:100]}".encode()
+        ).hexdigest()[:16]
+        return Chunk(
+            id=chunk_id, content=content, chunk_type=chunk_type,
+            file_path=file_path, start_line=start_line, end_line=end_line, language=language,
+        )
+
+    def chunk_with_imports(
+        self, source: str, file_path: str, language: str
+    ) -> tuple[list[Chunk], list[str]]:
+        chunks = self.chunk(source, file_path, language)
+        imports = self._extract_imports(source, language)
+        return chunks, imports
+
+    def _extract_imports(self, source: str, language: str) -> list[str]:
+        parser = self._get_parser(language)
+        if parser is None:
+            return []
+        tree = parser.parse(source.encode("utf-8"))
+        imports: list[str] = []
+        self._walk_imports(tree.root_node, source, language, imports)
+        return list(dict.fromkeys(imports))  # deduplicate while preserving order
+
+    def _walk_imports(self, node, source, language, imports):
+        if node.type in _IMPORT_TYPES:
+            module = self._parse_import_module(node, source, language)
+            if module:
+                imports.append(module)
+        for child in node.children:
+            self._walk_imports(child, source, language, imports)
+
+    def _parse_import_module(self, node, source, language) -> str | None:
+        if node.type == "import_statement":
+            # Python: "import os" or "import os.path"
+            # Also handles JS/TS: "import React from 'react'" (string child present)
+            for child in node.children:
+                if child.type == "string":
+                    # JavaScript/TypeScript import with string module specifier
+                    raw = source[child.start_byte:child.end_byte].strip("'\"")
+                    return raw.split("/")[0] if not raw.startswith("@") else "/".join(raw.split("/")[:2])
+                if child.type in ("dotted_name", "aliased_import"):
+                    # Python bare import
+                    name = source[child.start_byte:child.end_byte]
+                    name = name.split(" as ")[0].strip()
+                    return name.split(".")[0]
+        elif node.type == "import_from_statement":
+            # Python: "from pathlib import Path"
+            for child in node.children:
+                if child.type in ("dotted_name", "relative_import"):
+                    name = source[child.start_byte:child.end_byte].strip()
+                    name = name.lstrip(".")
+                    if name:
+                        return name.split(".")[0]
+        elif node.type == "import_declaration":
+            # TypeScript (tree-sitter-typescript): "import React from 'react'"
+            for child in node.children:
+                if child.type == "string":
+                    raw = source[child.start_byte:child.end_byte].strip("'\"")
+                    return raw.split("/")[0] if not raw.startswith("@") else "/".join(raw.split("/")[:2])
+        return None
+
+    def _fallback_chunk(self, source, file_path, language):
+        chunk_id = hashlib.sha256(f"{file_path}:module".encode()).hexdigest()[:16]
+        lines = source.strip().split("\n")
+        return Chunk(
+            id=chunk_id, content=source, chunk_type=ChunkType.MODULE,
+            file_path=file_path, start_line=1, end_line=len(lines), language=language,
+        )

context_engine/indexer/embedder.py

@@ -0,0 +1,154 @@
+"""Embedding generation using fastembed (lightweight ONNX-based embeddings).
+
+Uses BAAI/bge-small-en-v1.5 by default — 33% smaller and better quality
+than all-MiniLM-L6-v2. Parallel embedding for 3-4x faster indexing.
+
+Supports an optional EmbeddingCache so unchanged code chunks are never
+re-embedded across index runs (inspired by Cursor's content-hash cache).
+"""
+import logging
+import os
+import sys
+from functools import lru_cache
+
+from fastembed import TextEmbedding
+
+from context_engine.indexer.embedding_cache import EmbeddingCache
+from context_engine.models import Chunk
+
+log = logging.getLogger(__name__)
+
+_DEFAULT_MODEL = "BAAI/bge-small-en-v1.5"
+
+# Passed straight to fastembed's `parallel` argument:
+#   None → no data-parallel mp; use onnxruntime's own threading
+#   N>0  → spawn N forkserver workers around onnxruntime
+#
+# Even parallel=1 takes the multiprocessing path — and that path deadlocks on
+# macOS (workers idle on SimpleQueue.get while the main thread sits in
+# asyncio.poll, leaving `cce init` stuck after the file-scan progress bar
+# hits 100%). Default to None on darwin; allow override via CCE_EMBED_PARALLEL.
+def _resolve_parallel() -> int | None:
+    override = os.environ.get("CCE_EMBED_PARALLEL")
+    if override:
+        try:
+            return max(1, int(override))
+        except ValueError:
+            pass
+    if sys.platform == "darwin":
+        return None
+    return min(os.cpu_count() or 2, 4)
+
+
+_PARALLEL: int | None = _resolve_parallel()
+
+
+class Embedder:
+    def __init__(
+        self,
+        model_name: str = "BAAI/bge-small-en-v1.5",
+        cache: EmbeddingCache | None = None,
+    ) -> None:
+        self._cache = cache
+        # Resolve short names: "all-MiniLM-L6-v2" → "sentence-transformers/all-MiniLM-L6-v2"
+        # but leave fully qualified names like "BAAI/bge-small-en-v1.5" alone.
+        if "/" not in model_name:
+            resolved = f"sentence-transformers/{model_name}"
+        else:
+            resolved = model_name
+        try:
+            self._model = TextEmbedding(resolved)
+        except Exception as exc:
+            raise RuntimeError(
+                f"Failed to load embedding model '{model_name}'. "
+                f"Ensure fastembed is installed and the model name is valid. "
+                f"Supported models: TextEmbedding.list_supported_models(). "
+                f"Original error: {exc}"
+            ) from exc
+
+    def embed(
+        self,
+        chunks: list[Chunk],
+        batch_size: int = 64,
+        progress_fn=None,
+    ) -> None:
+        """Embed chunks in-place. With a cache attached, only chunks whose
+        content hash is not already in the cache go through the model.
+
+        `progress_fn(current, total)` is called as embedding proceeds, where
+        `total` is the count of chunks that actually needed embedding (cache
+        misses). Cache hits return instantly and don't trigger callbacks.
+        """
+        if not chunks:
+            return
+
+        if self._cache is None:
+            self._embed_all(chunks, batch_size, progress_fn=progress_fn)
+            return
+
+        # Hash + batched lookup: one SQL roundtrip for the whole batch
+        # instead of N roundtrips through the per-chunk get() path.
+        hashes = [self._cache.content_hash(c.content) for c in chunks]
+        cached = self._cache.get_batch(hashes)
+
+        miss_indices: list[int] = []
+        for i, h in enumerate(hashes):
+            if h in cached:
+                chunks[i].embedding = cached[h]
+            else:
+                miss_indices.append(i)
+
+        if miss_indices:
+            miss_chunks = [chunks[i] for i in miss_indices]
+            self._embed_all(miss_chunks, batch_size, progress_fn=progress_fn)
+            # Persist newly-computed embeddings back to the cache.
+            new_entries = [
+                (hashes[i], chunks[i].embedding)
+                for i in miss_indices
+                if chunks[i].embedding is not None
+            ]
+            if new_entries:
+                self._cache.put_batch(new_entries)
+
+        cache_total = len(chunks)
+        cache_hits = cache_total - len(miss_indices)
+        if cache_hits > 0:
+            log.info(
+                "Embedding cache: %d/%d hits (%.0f%% reused)",
+                cache_hits, cache_total, cache_hits / cache_total * 100,
+            )
+
+    def _embed_all(
+        self,
+        chunks: list[Chunk],
+        batch_size: int = 64,
+        progress_fn=None,
+    ) -> None:
+        """Embed all chunks via the model (no cache).
+
+        Iterates fastembed's generator one item at a time so we can tick a
+        progress callback. The model still embeds in batches internally; we
+        just observe one yielded vector at a time.
+        """
+        texts = [c.content for c in chunks]
+        total = len(texts)
+        if progress_fn:
+            progress_fn(0, total)
+        for i, emb in enumerate(self._model.embed(
+            texts,
+            batch_size=batch_size,
+            parallel=_PARALLEL,
+        )):
+            chunks[i].embedding = emb.tolist()
+            if progress_fn and ((i + 1) % batch_size == 0 or i + 1 == total):
+                progress_fn(i + 1, total)
+
+    @lru_cache(maxsize=256)
+    def embed_query(self, query: str) -> tuple:
+        """Embed a single query string. Returns tuple for LRU cache hashability.
+
+        Callers that need a list (e.g. LanceDB) should use list(result)
+        or the _to_list() helper in vector_store.
+        """
+        results = list(self._model.query_embed(query))
+        return tuple(results[0].tolist())

context_engine/indexer/embedding_cache.py

@@ -0,0 +1,168 @@
+"""SQLite-backed embedding cache keyed by content hash.
+
+Avoids recomputing embeddings for unchanged code chunks across re-index runs.
+Inspired by Cursor's approach of caching embeddings by chunk content hash so
+identical code is never re-embedded.
+
+Vectors are stored via `struct.pack` (binary float32) rather than JSON — same
+encoding the sqlite-vec store uses elsewhere in the codebase. JSON would be
+~4× larger on disk for typical 384-dim embeddings.
+"""
+import hashlib
+import logging
+import sqlite3
+import struct
+from pathlib import Path
+from threading import RLock
+
+log = logging.getLogger(__name__)
+
+_SCHEMA = """
+CREATE TABLE IF NOT EXISTS embedding_cache (
+    content_hash TEXT PRIMARY KEY,
+    dim          INTEGER NOT NULL,
+    embedding    BLOB NOT NULL
+);
+"""
+
+
+class EmbeddingCache:
+    """Maps content SHA-256 → embedding vector, persisted in SQLite.
+
+    When *model_name* is provided the content hash is salted with the model
+    identifier so that switching embedding models automatically invalidates
+    stale cache entries rather than silently returning vectors with the wrong
+    dimensionality or semantics.
+
+    All SQLite access is serialised with an RLock (same pattern as VectorStore
+    and FTSStore). ``check_same_thread=False`` only disables Python's ownership
+    check; concurrent calls still need explicit locking.
+    """
+
+    def __init__(self, cache_path: Path, *, model_name: str = "") -> None:
+        self._path = cache_path
+        self._model_name = model_name
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+        self._lock = RLock()
+        self._conn = sqlite3.connect(str(self._path), check_same_thread=False)
+        self._conn.execute("PRAGMA journal_mode=WAL")
+        self._conn.execute("PRAGMA synchronous=NORMAL")
+        self._conn.execute(_SCHEMA)
+        self._conn.commit()
+        self._hits = 0
+        self._misses = 0
+
+    def content_hash(self, text: str) -> str:
+        """SHA-256 of *text*, salted with model name when set."""
+        key = f"{self._model_name}:{text}" if self._model_name else text
+        return hashlib.sha256(key.encode("utf-8")).hexdigest()
+
+    @staticmethod
+    def _pack(vec) -> bytes:
+        v = list(vec) if not isinstance(vec, list) else vec
+        return struct.pack(f"{len(v)}f", *v)
+
+    @staticmethod
+    def _unpack(blob: bytes, dim: int) -> list[float]:
+        return list(struct.unpack(f"{dim}f", blob))
+
+    def get(self, content_hash: str) -> list[float] | None:
+        with self._lock:
+            row = self._conn.execute(
+                "SELECT dim, embedding FROM embedding_cache WHERE content_hash = ?",
+                (content_hash,),
+            ).fetchone()
+        if row is None:
+            self._misses += 1
+            return None
+        self._hits += 1
+        return self._unpack(row[1], row[0])
+
+    def put(self, content_hash: str, embedding) -> None:
+        v = list(embedding) if not isinstance(embedding, list) else embedding
+        with self._lock:
+            self._conn.execute(
+                "INSERT OR REPLACE INTO embedding_cache (content_hash, dim, embedding) VALUES (?, ?, ?)",
+                (content_hash, len(v), self._pack(v)),
+            )
+            self._conn.commit()
+
+    def put_batch(self, items: list[tuple[str, list[float]]]) -> None:
+        rows = [(h, len(e), self._pack(e)) for h, e in items]
+        with self._lock:
+            self._conn.executemany(
+                "INSERT OR REPLACE INTO embedding_cache (content_hash, dim, embedding) VALUES (?, ?, ?)",
+                rows,
+            )
+            self._conn.commit()
+
+    def get_batch(self, content_hashes: list[str]) -> dict[str, list[float]]:
+        """Retrieve multiple embeddings at once. Returns hash → embedding for hits."""
+        if not content_hashes:
+            return {}
+        results: dict[str, list[float]] = {}
+        with self._lock:
+            for i in range(0, len(content_hashes), 500):
+                batch = content_hashes[i : i + 500]
+                placeholders = ",".join("?" * len(batch))
+                rows = self._conn.execute(
+                    f"SELECT content_hash, dim, embedding FROM embedding_cache "
+                    f"WHERE content_hash IN ({placeholders})",
+                    batch,
+                ).fetchall()
+                for h, dim, blob in rows:
+                    results[h] = self._unpack(blob, dim)
+        self._hits += len(results)
+        self._misses += len(content_hashes) - len(results)
+        return results
+
+    def prune_orphans(self, known_hashes: set[str]) -> int:
+        """Drop cached entries whose content_hash is not in `known_hashes`.
+
+        Cache grows monotonically without this — every chunk content variant
+        ever seen accumulates forever even after the source files change or
+        get deleted. Call this after a `cce index --full` with the set of
+        hashes still present in the live index. Returns the count removed.
+        """
+        if not known_hashes:
+            return 0
+        with self._lock:
+            cur = self._conn.execute("SELECT content_hash FROM embedding_cache")
+            current = {row[0] for row in cur.fetchall()}
+            orphans = current - known_hashes
+            if not orphans:
+                return 0
+            removed = 0
+            orphan_list = list(orphans)
+            for i in range(0, len(orphan_list), 500):
+                batch = orphan_list[i : i + 500]
+                placeholders = ",".join("?" * len(batch))
+                self._conn.execute(
+                    f"DELETE FROM embedding_cache WHERE content_hash IN ({placeholders})",
+                    batch,
+                )
+                removed += len(batch)
+            self._conn.commit()
+        return removed
+
+    @property
+    def hits(self) -> int:
+        return self._hits
+
+    @property
+    def misses(self) -> int:
+        return self._misses
+
+    @property
+    def hit_rate(self) -> float:
+        total = self._hits + self._misses
+        return self._hits / total if total > 0 else 0.0
+
+    def size(self) -> int:
+        with self._lock:
+            row = self._conn.execute("SELECT COUNT(*) FROM embedding_cache").fetchone()
+        return row[0] if row else 0
+
+    def close(self) -> None:
+        with self._lock:
+            self._conn.close()

context_engine/indexer/git_hooks.py

@@ -0,0 +1,73 @@
+"""Git hook installer and handler for triggering re-indexing."""
+import os
+import shutil
+import stat
+import sys
+from pathlib import Path
+
+HOOK_MARKER = "# cce hook"
+HOOK_NAMES = ["post-commit", "post-checkout", "post-merge"]
+
+
+def _resolve_cce_binary() -> str:
+    """Find an absolute path to the `cce` launcher.
+
+    Preferring an absolute path means the git hook keeps working when the user
+    runs `git commit` from a shell that doesn't pick up the same PATH as the one
+    used to install the engine (e.g. different login shell, GUI git client).
+    """
+    candidate = Path(sys.executable).parent / "cce"
+    if candidate.exists():
+        return str(candidate)
+    which = shutil.which("cce") or shutil.which("code-context-engine")
+    if which:
+        return which
+    # Last-resort: rely on PATH at hook-run time.
+    return "cce"
+
+
+def _hook_script() -> str:
+    bin_path = _resolve_cce_binary()
+    return f"""{HOOK_MARKER}
+{bin_path} index --changed-only >/dev/null 2>&1 &
+"""
+
+
+def install_hooks(project_dir: str) -> list[str]:
+    """Install CCE git hooks. Returns [] gracefully if not a git repo."""
+    hooks_dir = Path(project_dir) / ".git" / "hooks"
+    if not hooks_dir.exists():
+        return []
+    installed = []
+    for hook_name in HOOK_NAMES:
+        hook_path = hooks_dir / hook_name
+        _install_single_hook(hook_path)
+        installed.append(str(hook_path))
+    return installed
+
+
+def _install_single_hook(hook_path: Path) -> None:
+    script = _hook_script()
+    if hook_path.exists():
+        existing = hook_path.read_text()
+        if HOOK_MARKER in existing:
+            return
+        new_content = existing.rstrip() + "\n\n" + script
+    else:
+        new_content = "#!/bin/sh\n\n" + script
+    hook_path.write_text(new_content)
+    hook_path.chmod(hook_path.stat().st_mode | stat.S_IEXEC)
+
+
+def get_changed_files_from_hook() -> list[str]:
+    import subprocess
+    try:
+        result = subprocess.run(
+            ["git", "diff", "--name-only", "HEAD~1", "HEAD"],
+            capture_output=True, text=True, timeout=10,
+        )
+        if result.returncode == 0:
+            return [f for f in result.stdout.strip().split("\n") if f]
+    except (subprocess.TimeoutExpired, FileNotFoundError):
+        pass
+    return []