contextual-engine 0.1.0__py3-none-any.whl

contextual/docs/pipeline.py

@@ -0,0 +1,513 @@
+"""Docs indexing pipeline for Contextual Phase 2.
+
+Mirrors ``contextual/indexing/pipeline.py`` but for markdown/MDX/RST/TXT files.
+Orchestrates:
+  1. File discovery   — respects .gitignore + .contextualignore
+  2. Docs chunking    — heading-aware via DocsChunker
+  3. Embedding        — fastembed (jina-embeddings-v2-base-code, 768-dim)
+  4. Storage          — writes to chunks / vec_chunks / fts_chunks with source_type='docs'
+  5. File watching    — integrates with the existing watchdog FileWatcher
+
+Design constraints:
+- Single-writer to SQLite (same pool as code indexer).
+- Incremental: skip files whose content hash has not changed.
+- source_type='docs' column used for filter pushdown in retrieval.
+- All logging to stderr via structlog (NEVER stdout — breaks stdio MCP transport).
+"""
+from __future__ import annotations
+
+import hashlib
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from sqlite3 import Connection
+from typing import Callable
+
+import pathspec
+import structlog
+
+from contextual.docs.chunker import DocChunk, DocsChunker
+from contextual.embedding import CodeEmbedder, DocsEmbedder
+from contextual.storage import SQLitePool
+
+logger = structlog.get_logger(__name__)
+
+# ---------------------------------------------------------------------------
+# Supported document extensions
+# ---------------------------------------------------------------------------
+DOC_EXTENSIONS: frozenset[str] = frozenset({".md", ".mdx", ".rst", ".txt"})
+
+# Directories always excluded regardless of .gitignore
+_ALWAYS_EXCLUDE: frozenset[str] = frozenset({
+    ".git", ".hg", ".svn",
+    "node_modules", "__pycache__", ".venv", "venv", ".env",
+    "target", "dist", "build", ".contextual",
+    ".mypy_cache", ".ruff_cache", ".pytest_cache",
+})
+
+# ---------------------------------------------------------------------------
+# Stats dataclass
+# ---------------------------------------------------------------------------
+
+@dataclass
+class DocsIndexStats:
+    """Counters returned after an indexing run."""
+    files_discovered: int = 0
+    files_skipped_unchanged: int = 0
+    files_indexed: int = 0
+    files_failed: int = 0
+    chunks_created: int = 0
+    embeddings_generated: int = 0
+    duration_seconds: float = 0.0
+
+    def __str__(self) -> str:  # noqa: D105
+        return (
+            f"files={self.files_indexed}/{self.files_discovered} "
+            f"chunks={self.chunks_created} "
+            f"embeddings={self.embeddings_generated} "
+            f"skipped={self.files_skipped_unchanged} "
+            f"failed={self.files_failed} "
+            f"duration={self.duration_seconds:.2f}s"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Schema helpers
+# ---------------------------------------------------------------------------
+
+def _ensure_docs_schema(conn: Connection) -> None:
+    """Ensure Phase 2 columns exist on the chunks table.
+
+    Idempotent — safe to call on every startup.  We use ALTER TABLE … ADD
+    COLUMN which is a no-op if the column already exists (SQLite ≥ 3.37).
+    """
+    existing_cols = {
+        row[1]
+        for row in conn.execute("PRAGMA table_info(chunks)").fetchall()
+    }
+
+    migrations: list[str] = []
+
+    if "source_type" not in existing_cols:
+        migrations.append(
+            "ALTER TABLE chunks ADD COLUMN source_type TEXT NOT NULL DEFAULT 'code'"
+        )
+    if "heading_path" not in existing_cols:
+        migrations.append(
+            "ALTER TABLE chunks ADD COLUMN heading_path TEXT NOT NULL DEFAULT ''"
+        )
+    if "heading_level" not in existing_cols:
+        migrations.append(
+            "ALTER TABLE chunks ADD COLUMN heading_level INTEGER NOT NULL DEFAULT 0"
+        )
+    if "chunk_type" not in existing_cols:
+        migrations.append(
+            "ALTER TABLE chunks ADD COLUMN chunk_type TEXT NOT NULL DEFAULT 'code'"
+        )
+    if "frontmatter_json" not in existing_cols:
+        migrations.append(
+            "ALTER TABLE chunks ADD COLUMN frontmatter_json TEXT NOT NULL DEFAULT '{}'"
+        )
+
+    if migrations:
+        for sql in migrations:
+            conn.execute(sql)
+        conn.commit()
+        logger.info("Applied Phase 2 schema columns", count=len(migrations))
+
+
+def _file_content_hash(path: Path) -> str:
+    """SHA-256 hex digest of file contents (for change detection)."""
+    h = hashlib.sha256()
+    h.update(path.read_bytes())
+    return h.hexdigest()
+
+
+# ---------------------------------------------------------------------------
+# Main pipeline class
+# ---------------------------------------------------------------------------
+
+class DocsPipeline:
+    """Full docs indexing pipeline.
+
+    Args:
+        repo_root: Absolute path to the repository root.
+        pool:      Shared SQLitePool (single-writer).
+        embedder:  Embedder instance (fastembed ONNX, 768-dim).
+        progress:  Optional callback ``(current, total, message)`` for CLI progress.
+    """
+
+    def __init__(
+        self,
+        repo_root: Path,
+        pool: SQLitePool,
+        embedder: DocsEmbedder | CodeEmbedder,
+        progress: Callable[[int, int, str], None] | None = None,
+    ) -> None:
+        self._root = repo_root.resolve()
+        self._pool = pool
+        self._embedder = embedder
+        self._chunker = DocsChunker()
+        self._progress = progress or (lambda *_: None)
+
+    # ------------------------------------------------------------------
+    # Public entry points
+    # ------------------------------------------------------------------
+
+    def index_all(self, force: bool = False) -> DocsIndexStats:
+        """Full docs index of the repository.
+
+        Args:
+            force: Re-index every file even if hash unchanged.
+
+        Returns:
+            ``DocsIndexStats`` with counters for this run.
+        """
+        t0 = time.perf_counter()
+        stats = DocsIndexStats()
+
+        with self._pool.writer() as conn:
+            _ensure_docs_schema(conn)
+
+            file_paths = self._discover_files()
+            stats.files_discovered = len(file_paths)
+            logger.info("Docs discovery complete", discovered=len(file_paths), root=str(self._root))
+
+            for i, fp in enumerate(file_paths, start=1):
+                self._progress(i, len(file_paths), fp.name)
+                result = self._index_file(conn, fp, force=force)
+                if result is None:
+                    stats.files_skipped_unchanged += 1
+                elif result == -1:
+                    stats.files_failed += 1
+                else:
+                    stats.files_indexed += 1
+                    stats.chunks_created += result
+
+        stats.duration_seconds = time.perf_counter() - t0
+        logger.info("Docs index_all complete", stats=str(stats))
+        return stats
+
+    def index_file(self, file_path: Path, force: bool = False) -> DocsIndexStats:
+        """Incrementally index (or re-index) a single document file.
+
+        Intended for file-watcher integration.
+
+        Args:
+            file_path: Absolute path to the changed file.
+            force:     Skip hash check and always re-index.
+        """
+        t0 = time.perf_counter()
+        stats = DocsIndexStats(files_discovered=1)
+
+        if file_path.suffix.lower() not in DOC_EXTENSIONS:
+            logger.debug("Skipped: not a doc extension", path=str(file_path))
+            return stats
+
+        with self._pool.writer() as conn:
+            result = self._index_file(conn, file_path, force=force)
+            if result is None:
+                stats.files_skipped_unchanged = 1
+            elif result == -1:
+                stats.files_failed = 1
+            else:
+                stats.files_indexed = 1
+                stats.chunks_created = result
+
+        stats.duration_seconds = time.perf_counter() - t0
+        return stats
+
+    def delete_file(self, file_path: Path) -> int:
+        """Remove all chunks for a deleted file from the index.
+
+        Returns:
+            Number of chunks deleted.
+        """
+        rel = self._rel_path(file_path)
+        with self._pool.writer() as conn:
+            cur = conn.execute(
+                "SELECT id FROM chunks WHERE file_path = ? AND source_type = 'docs'",
+                (rel,),
+            )
+            chunk_ids = [row[0] for row in cur.fetchall()]
+            if not chunk_ids:
+                return 0
+            placeholders = ",".join("?" * len(chunk_ids))
+            conn.execute(f"DELETE FROM vec_chunks WHERE chunk_id IN ({placeholders})", chunk_ids)
+            conn.execute(
+                f"DELETE FROM chunks WHERE id IN ({placeholders})",
+                chunk_ids,
+            )
+            conn.commit()
+        logger.info("Deleted docs chunks", path=rel, count=len(chunk_ids))
+        return len(chunk_ids)
+
+    # ------------------------------------------------------------------
+    # File discovery
+    # ------------------------------------------------------------------
+
+    def _discover_files(self) -> list[Path]:
+        """Walk repo_root and return all indexable doc files.
+
+        Respects:
+        - .gitignore (via pathspec)
+        - .contextualignore (project-specific additions)
+        - _ALWAYS_EXCLUDE hardcoded dir set
+        """
+        ignore_spec = self._load_ignore_spec()
+        results: list[Path] = []
+
+        for path in self._root.rglob("*"):
+            if not path.is_file():
+                continue
+            if path.suffix.lower() not in DOC_EXTENSIONS:
+                continue
+            # Exclude always-ignored dir components
+            parts = set(path.relative_to(self._root).parts)
+            if parts & _ALWAYS_EXCLUDE:
+                continue
+            # Exclude gitignore patterns
+            rel = str(path.relative_to(self._root))
+            if ignore_spec and ignore_spec.match_file(rel):
+                continue
+            results.append(path)
+
+        results.sort()
+        return results
+
+    def _load_ignore_spec(self) -> pathspec.PathSpec | None:
+        """Load .gitignore + .contextualignore patterns."""
+        patterns: list[str] = []
+        for ignore_file in (".gitignore", ".contextualignore"):
+            ig_path = self._root / ignore_file
+            if ig_path.exists():
+                try:
+                    patterns.extend(ig_path.read_text(encoding="utf-8").splitlines())
+                except OSError:
+                    pass
+        if not patterns:
+            return None
+        return pathspec.PathSpec.from_lines("gitwildmatch", patterns)
+
+    # ------------------------------------------------------------------
+    # Single-file indexing
+    # ------------------------------------------------------------------
+
+    def _index_file(
+        self, conn: Connection, file_path: Path, force: bool
+    ) -> int | None:
+        """Index a single file within an active writer connection.
+
+        Returns:
+            int:  number of chunks written (0 is valid for empty files).
+            None: file skipped (hash unchanged and not force).
+            -1:   file failed (logged internally).
+        """
+        rel = self._rel_path(file_path)
+
+        # --- Change detection ---
+        try:
+            content_hash = _file_content_hash(file_path)
+        except OSError as exc:
+            logger.warning("Cannot hash file — skipping", path=rel, error=str(exc))
+            return -1
+
+        if not force:
+            row = conn.execute(
+                "SELECT content_hash FROM files WHERE path = ?", (rel,)
+            ).fetchone()
+            if row and row[0] == content_hash:
+                logger.debug("File unchanged — skipping", path=rel)
+                return None
+
+        # --- Remove stale chunks ---
+        self._purge_file_chunks(conn, rel)
+
+        # --- Chunk ---
+        try:
+            chunks: list[DocChunk] = self._chunker.chunk_file(file_path, repo_root=self._root)
+        except Exception as exc:  # noqa: BLE001
+            logger.error("Chunking failed", path=rel, error=str(exc))
+            return -1
+
+        if not chunks:
+            logger.debug("No chunks produced", path=rel)
+            self._upsert_file_record(conn, rel, content_hash, chunk_count=0)
+            conn.commit()
+            return 0
+
+        # --- Embed ---
+        chunk_tuples = []
+        for c in chunks:
+            emb_hash = hashlib.sha256(c.embedding_text.encode("utf-8")).hexdigest()
+            chunk_tuples.append((emb_hash, c.embedding_text))
+
+        try:
+            from contextual.embedding.helpers import embed_chunks
+            cached_results = embed_chunks(conn, chunk_tuples)
+            embeddings: list[list[float]] = [emb for _, emb in cached_results]
+        except Exception as exc:  # noqa: BLE001
+            logger.error("Embedding failed", path=rel, chunks=len(chunks), error=str(exc))
+            return -1
+
+        if len(embeddings) != len(chunks):
+            logger.error(
+                "Embedding count mismatch",
+                expected=len(chunks),
+                got=len(embeddings),
+                path=rel,
+            )
+            return -1
+
+        # --- Write chunks + vectors ---
+        import json
+
+        file_id = self._upsert_file_record(conn, rel, content_hash, chunk_count=len(chunks))
+
+        chunk_ids: list[int] = []
+        chunk_data = []
+        for chunk in chunks:
+            chunk_data.append((
+                file_id,
+                rel,
+                chunk.content,
+                chunk.chunk_index,
+                chunk.start_line,
+                chunk.end_line,
+                "markdown",
+                "docs",
+                chunk.heading_path,
+                chunk.heading_level,
+                chunk.chunk_type,
+                json.dumps(chunk.frontmatter),
+            ))
+
+        # Batch insert chunks and collect IDs
+        # Note: SQLite executemany doesn't support RETURNING in older versions,
+        # so we insert one by one to get lastrowid, OR we use a high enough version.
+        # Given the "N+1" fix requirement, we should batch.
+        # However, to get IDs for FTS and Vec0, we'll insert one by one and THEN batch FTS/Vec.
+        for data in chunk_data:
+            cur = conn.execute(
+                """
+                INSERT INTO chunks (
+                    file_id, file_path, content, chunk_index,
+                    start_line, end_line, language,
+                    source_type, heading_path, heading_level,
+                    chunk_type, frontmatter_json, indexed_at
+                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, unixepoch('now'))
+                """,
+                data
+            )
+            if cur.lastrowid is None:
+                logger.error("Failed to retrieve inserted chunk ID", path=rel)
+                return -1
+            chunk_ids.append(cur.lastrowid)
+
+        # --- Batch FTS5 insert ---
+        if chunk_ids:
+            fts_rows = [
+                (cid, chunks[i].content, "", "", chunks[i].heading_path)
+                for i, cid in enumerate(chunk_ids)
+            ]
+            conn.executemany(
+                """
+                INSERT INTO fts_chunks (rowid, content, file_path, symbol_name, heading_path)
+                VALUES (?, ?, ?, ?, ?)
+                """,
+                fts_rows
+            )
+
+        # --- Vec0 batch insert ---
+        self._insert_vectors(conn, list(zip(chunk_ids, embeddings, strict=True)))
+
+        conn.commit()
+        logger.info(
+            "Docs file indexed",
+            path=rel,
+            chunks=len(chunks),
+            source_type="docs",
+        )
+        return len(chunks)
+
+    # ------------------------------------------------------------------
+    # Storage helpers
+    # ------------------------------------------------------------------
+
+    def _purge_file_chunks(self, conn: Connection, rel_path: str) -> None:
+        """Delete all existing chunks + vectors for a doc file."""
+        cur = conn.execute(
+            "SELECT id FROM chunks WHERE file_path = ? AND source_type IN ('docs','docs_code')",
+            (rel_path,),
+        )
+        ids = [row[0] for row in cur.fetchall()]
+        if ids:
+            ph = ",".join("?" * len(ids))
+            conn.execute(f"DELETE FROM vec_chunks WHERE chunk_id IN ({ph})", ids)
+            conn.execute(f"DELETE FROM fts_chunks WHERE rowid IN ({ph})", ids)
+            conn.execute(f"DELETE FROM chunks WHERE id IN ({ph})", ids)
+
+    def _upsert_file_record(
+        self, conn: Connection, rel_path: str, content_hash: str, chunk_count: int
+    ) -> int:
+        """Insert or update the files table row, returning the file_id."""
+        conn.execute(
+            """
+            INSERT INTO files (path, language, mtime, content_hash, indexed_at)
+            VALUES (?, ?, unixepoch('now'), ?, unixepoch('now'))
+            ON CONFLICT(path) DO UPDATE SET
+                content_hash = excluded.content_hash,
+                indexed_at   = excluded.indexed_at,
+                mtime        = excluded.mtime
+            """,
+            (rel_path, "markdown", content_hash),
+        )
+        row = conn.execute("SELECT id FROM files WHERE path = ?", (rel_path,)).fetchone()
+        return row[0]
+
+    def _insert_vectors(
+        self, conn: Connection, pairs: list[tuple[int, list[float]]]
+    ) -> None:
+        """Batch-insert (chunk_id, embedding) into vec_chunks (vec0)."""
+        import json as _json
+
+        if pairs:
+            vec_rows = [(chunk_id, _json.dumps(vec)) for chunk_id, vec in pairs]
+            conn.executemany(
+                "INSERT OR REPLACE INTO vec_chunks (chunk_id, embedding) VALUES (?, ?)",
+                vec_rows,
+            )
+
+    def _rel_path(self, abs_path: Path) -> str:
+        """Return repo-relative string path."""
+        try:
+            return str(abs_path.relative_to(self._root))
+        except ValueError:
+            return str(abs_path)
+
+
+# ---------------------------------------------------------------------------
+# Module-level convenience
+# ---------------------------------------------------------------------------
+
+def index_docs(
+    repo_root: Path,
+    pool: SQLitePool,
+    embedder: CodeEmbedder | DocsEmbedder,
+    force: bool = False,
+    progress: Callable[[int, int, str], None] | None = None,
+) -> DocsIndexStats:
+    """One-shot convenience: index all docs in a repository.
+
+    Args:
+        repo_root: Absolute repository root.
+        pool:      Active SQLitePool.
+        embedder:  Embedder instance.
+        force:     Re-index regardless of content hash.
+        progress:  Optional ``(current, total, filename)`` progress callback.
+
+    Returns:
+        ``DocsIndexStats`` summary.
+    """
+    pipeline = DocsPipeline(repo_root, pool, embedder, progress=progress)
+    return pipeline.index_all(force=force)