attocode-code-intel 0.1.0__py3-none-any.whl

attocode_intel/_internal/integrations/context/embedding_rotation.py

@@ -0,0 +1,597 @@
+"""Embedding rotation state machine for the local ``vectors`` store.
+
+Migrates every stored embedding from one model to another without ever
+touching (let alone wiping) the live rows — the opposite of the
+destructive-NULL path Phase 1 killed.
+
+States::
+
+        none
+         │    embeddings_rotate_start(new_model, new_dim)
+         ▼
+      pending ─── staging table created, total row count captured
+         │
+         │    embeddings_rotate_step(batch_size)   (call repeatedly)
+         ▼
+    backfilling ── each batch: read N old rows, embed with new provider,
+         │                     insert into vectors_rotating
+         │
+         │    last batch processed (cursor exhausted)
+         ▼
+  ready_to_cutover
+         │
+         │    embeddings_rotate_cutover(confirm=True)
+         ▼                             ┌── rename vectors → vectors_archive
+    cutover_done                       └── rename vectors_rotating → vectors
+         │                                 update store_metadata.dimension
+         │    embeddings_rotate_gc_old(confirm=True)
+         ▼
+      gc_done ──── vectors_archive dropped; rotation metadata cleared
+         │
+         ▼
+        none
+
+At any pre-cutover state, ``embeddings_rotate_abort(confirm=True)`` drops
+``vectors_rotating`` and resets state to ``none`` — the primary ``vectors``
+table is untouched.
+
+**Known limitation:** new writes to ``vectors`` that happen DURING rotation
+(e.g. ``notify_file_changed`` triggering a reindex mid-rotation) land in
+the old table and never reach ``vectors_rotating``. After cutover they're
+effectively missing. For Phase 2b the recommendation is to avoid concurrent
+indexing during a rotation, or run a normal ``semantic_search`` reindex
+afterwards to catch up.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import sqlite3
+import struct
+import time
+from dataclasses import dataclass
+from enum import StrEnum
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from attocode_intel._internal.integrations.context.embeddings import EmbeddingProvider
+
+logger = logging.getLogger(__name__)
+
+# Staging table name used for the rotation target.
+STAGING_TABLE = "vectors_rotating"
+# Archive table name used to hold the pre-rotation rows until GC.
+ARCHIVE_TABLE = "vectors_archive"
+
+# All store_metadata keys the rotator owns. Clearing these drops all
+# rotation state back to a clean slate.
+_ROTATION_KEYS: tuple[str, ...] = (
+    "rotation_state",
+    "rotation_from_model",
+    "rotation_from_version",
+    "rotation_from_dim",
+    "rotation_to_model",
+    "rotation_to_version",
+    "rotation_to_dim",
+    "rotation_started_at",
+    "rotation_cursor",
+    "rotation_processed_rows",
+    "rotation_total_rows",
+    "rotation_last_error",
+)
+
+
+class RotationState(StrEnum):
+    """Finite state machine for one embedding rotation."""
+
+    NONE = "none"
+    PENDING = "pending"
+    BACKFILLING = "backfilling"
+    READY_TO_CUTOVER = "ready_to_cutover"
+    CUTOVER_DONE = "cutover_done"
+    GC_DONE = "gc_done"
+    ABORTED = "aborted"
+    FAILED = "failed"
+
+
+@dataclass(slots=True)
+class RotationStatus:
+    """Public snapshot of the current rotation (or lack thereof)."""
+
+    state: RotationState
+    from_model: str
+    from_version: str
+    from_dim: int
+    to_model: str
+    to_version: str
+    to_dim: int
+    started_at: float
+    processed_rows: int
+    total_rows: int
+    last_error: str
+
+    def is_active(self) -> bool:
+        return self.state not in (
+            RotationState.NONE,
+            RotationState.GC_DONE,
+            RotationState.ABORTED,
+            RotationState.FAILED,
+        )
+
+    def progress_pct(self) -> float:
+        if self.total_rows <= 0:
+            return 0.0
+        return 100.0 * self.processed_rows / self.total_rows
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "state": self.state.value,
+            "from": {
+                "model": self.from_model,
+                "version": self.from_version,
+                "dim": self.from_dim,
+            },
+            "to": {
+                "model": self.to_model,
+                "version": self.to_version,
+                "dim": self.to_dim,
+            },
+            "started_at": self.started_at,
+            "processed_rows": self.processed_rows,
+            "total_rows": self.total_rows,
+            "progress_pct": round(self.progress_pct(), 2),
+            "last_error": self.last_error,
+        }
+
+
+# ---------------------------------------------------------------------------
+# EmbeddingRotator
+# ---------------------------------------------------------------------------
+
+
+class EmbeddingRotator:
+    """Rotation driver operating directly on a ``vectors.db`` file.
+
+    Deliberately bypasses ``VectorStore`` so that a dimension mismatch
+    between the old and new models doesn't trip ``_validate_dimension``.
+    All reads / writes go through raw ``sqlite3``; the rotator is
+    responsible for its own locking.
+    """
+
+    def __init__(self, db_path: str) -> None:
+        self.db_path = db_path
+        self._conn: sqlite3.Connection | None = None
+        if os.path.exists(db_path):
+            self._conn = sqlite3.connect(db_path, check_same_thread=False)
+            self._conn.execute("PRAGMA journal_mode=WAL")
+
+    def _get_conn(self) -> sqlite3.Connection:
+        if self._conn is None:
+            raise RuntimeError(
+                f"EmbeddingRotator: no vectors.db at {self.db_path}"
+            )
+        return self._conn
+
+    def close(self) -> None:
+        if self._conn:
+            self._conn.close()
+            self._conn = None
+
+    # ------------------------------------------------------------------
+    # Metadata helpers
+    # ------------------------------------------------------------------
+
+    def _set_meta(self, key: str, value: str) -> None:
+        conn = self._get_conn()
+        conn.execute(
+            "INSERT OR REPLACE INTO store_metadata (key, value) VALUES (?, ?)",
+            (key, str(value)),
+        )
+
+    def _get_meta(self, key: str, default: str = "") -> str:
+        conn = self._get_conn()
+        row = conn.execute(
+            "SELECT value FROM store_metadata WHERE key = ?", (key,),
+        ).fetchone()
+        return row[0] if row else default
+
+    def _clear_rotation_meta(self) -> None:
+        conn = self._get_conn()
+        for key in _ROTATION_KEYS:
+            conn.execute("DELETE FROM store_metadata WHERE key = ?", (key,))
+
+    # ------------------------------------------------------------------
+    # Status
+    # ------------------------------------------------------------------
+
+    def status(self) -> RotationStatus:
+        """Read the current rotation state from ``store_metadata``."""
+        if self._conn is None:
+            return RotationStatus(
+                state=RotationState.NONE,
+                from_model="", from_version="", from_dim=0,
+                to_model="", to_version="", to_dim=0,
+                started_at=0.0,
+                processed_rows=0, total_rows=0,
+                last_error="vectors.db not found",
+            )
+        state_raw = self._get_meta("rotation_state", RotationState.NONE.value)
+        try:
+            state = RotationState(state_raw)
+        except ValueError:
+            state = RotationState.NONE
+        return RotationStatus(
+            state=state,
+            from_model=self._get_meta("rotation_from_model"),
+            from_version=self._get_meta("rotation_from_version"),
+            from_dim=int(self._get_meta("rotation_from_dim") or "0"),
+            to_model=self._get_meta("rotation_to_model"),
+            to_version=self._get_meta("rotation_to_version"),
+            to_dim=int(self._get_meta("rotation_to_dim") or "0"),
+            started_at=float(self._get_meta("rotation_started_at") or "0"),
+            processed_rows=int(self._get_meta("rotation_processed_rows") or "0"),
+            total_rows=int(self._get_meta("rotation_total_rows") or "0"),
+            last_error=self._get_meta("rotation_last_error"),
+        )
+
+    # ------------------------------------------------------------------
+    # start
+    # ------------------------------------------------------------------
+
+    def start(
+        self,
+        *,
+        to_model: str,
+        to_version: str,
+        to_dim: int,
+    ) -> RotationStatus:
+        """Begin a new rotation.
+
+        Raises RuntimeError if a rotation is already active or if the
+        main ``vectors`` table is empty.
+        """
+        conn = self._get_conn()
+
+        current = self.status()
+        if current.is_active():
+            raise RuntimeError(
+                f"rotation already active (state={current.state.value}); "
+                "abort it before starting a new one"
+            )
+
+        # Capture the current "from" metadata from the active store_metadata.
+        from_dim = int(self._get_meta("dimension") or "0")
+        # "from" model is approximated from the most common model_name on
+        # existing rows. If the column was never populated (legacy rows
+        # from Phase 0), we fall back to "legacy".
+        from_model = self._dominant_model_name(conn) or "legacy"
+        from_version = self._dominant_model_version(conn)
+
+        row_count = conn.execute("SELECT COUNT(*) FROM vectors").fetchone()[0]
+        if row_count == 0:
+            raise RuntimeError("no vectors to rotate — table is empty")
+
+        # Clean any prior staging table from a failed rotation.
+        conn.execute(f"DROP TABLE IF EXISTS {STAGING_TABLE}")
+        # Mirror the vectors table schema. Using a literal CREATE keeps
+        # the PK + NOT NULL / DEFAULT constraints that a
+        # ``CREATE TABLE x AS SELECT`` would strip.
+        conn.execute(
+            f"""
+            CREATE TABLE {STAGING_TABLE} (
+                id TEXT PRIMARY KEY,
+                file_path TEXT NOT NULL,
+                chunk_type TEXT NOT NULL,
+                name TEXT NOT NULL,
+                text TEXT NOT NULL,
+                vector BLOB NOT NULL,
+                model_name TEXT NOT NULL DEFAULT '',
+                model_version TEXT NOT NULL DEFAULT '',
+                dimension INTEGER NOT NULL DEFAULT 0,
+                produced_at REAL NOT NULL DEFAULT 0,
+                blob_oid TEXT NOT NULL DEFAULT '',
+                action_hash TEXT NOT NULL DEFAULT ''
+            )
+            """
+        )
+        # Matching indexes on the staging table so they're in place for
+        # search as soon as the cutover rename swap completes.
+        conn.execute(
+            f"CREATE INDEX IF NOT EXISTS idx_rot_vectors_file "
+            f"ON {STAGING_TABLE}(file_path)"
+        )
+        conn.execute(
+            f"CREATE INDEX IF NOT EXISTS idx_rot_vectors_model "
+            f"ON {STAGING_TABLE}(model_name, model_version)"
+        )
+        conn.execute(
+            f"CREATE INDEX IF NOT EXISTS idx_rot_vectors_blob "
+            f"ON {STAGING_TABLE}(blob_oid)"
+        )
+
+        self._clear_rotation_meta()
+        self._set_meta("rotation_state", RotationState.PENDING.value)
+        self._set_meta("rotation_from_model", from_model)
+        self._set_meta("rotation_from_version", from_version)
+        self._set_meta("rotation_from_dim", str(from_dim))
+        self._set_meta("rotation_to_model", to_model)
+        self._set_meta("rotation_to_version", to_version)
+        self._set_meta("rotation_to_dim", str(to_dim))
+        self._set_meta("rotation_started_at", str(time.time()))
+        self._set_meta("rotation_cursor", "")
+        self._set_meta("rotation_processed_rows", "0")
+        self._set_meta("rotation_total_rows", str(row_count))
+        self._set_meta("rotation_last_error", "")
+        conn.commit()
+
+        return self.status()
+
+    def _dominant_model_name(self, conn: sqlite3.Connection) -> str:
+        """Return the most common ``model_name`` in the vectors table.
+
+        Returns an empty string if no rows have a non-empty name (e.g.
+        a fresh v2 store where rows were inserted before provenance was
+        wired through).
+        """
+        try:
+            row = conn.execute(
+                "SELECT model_name, COUNT(*) as c FROM vectors "
+                "WHERE model_name != '' "
+                "GROUP BY model_name ORDER BY c DESC LIMIT 1"
+            ).fetchone()
+        except sqlite3.OperationalError:
+            return ""
+        return row[0] if row else ""
+
+    def _dominant_model_version(self, conn: sqlite3.Connection) -> str:
+        try:
+            row = conn.execute(
+                "SELECT model_version, COUNT(*) as c FROM vectors "
+                "WHERE model_version != '' "
+                "GROUP BY model_version ORDER BY c DESC LIMIT 1"
+            ).fetchone()
+        except sqlite3.OperationalError:
+            return ""
+        return row[0] if row else ""
+
+    # ------------------------------------------------------------------
+    # step — advance the backfill
+    # ------------------------------------------------------------------
+
+    def step(
+        self,
+        provider: EmbeddingProvider,
+        batch_size: int = 32,
+    ) -> int:
+        """Embed one batch of rows with ``provider`` and write them to
+        the staging table.
+
+        Returns the number of rows processed. Returns 0 when the
+        rotation is done (state transitions to ``ready_to_cutover``).
+
+        Raises RuntimeError if there's no active rotation or if
+        ``provider.dimension()`` doesn't match the target dim recorded
+        at ``start``.
+        """
+        conn = self._get_conn()
+        status = self.status()
+
+        # Idempotent no-op once the backfill is already done — callers
+        # that drive step() in a while-loop would otherwise have to
+        # duplicate state-check logic.
+        if status.state == RotationState.READY_TO_CUTOVER:
+            return 0
+
+        if status.state not in (RotationState.PENDING, RotationState.BACKFILLING):
+            raise RuntimeError(
+                f"step() requires state pending|backfilling, got {status.state.value}"
+            )
+
+        provider_dim = provider.dimension()
+        if provider_dim != status.to_dim:
+            self._set_meta(
+                "rotation_last_error",
+                f"provider dim {provider_dim} != target dim {status.to_dim}",
+            )
+            self._set_meta("rotation_state", RotationState.FAILED.value)
+            conn.commit()
+            raise RuntimeError(
+                f"provider {provider.name} has dim {provider_dim}, "
+                f"but rotation target dim is {status.to_dim}"
+            )
+
+        last_cursor = self._get_meta("rotation_cursor", "")
+        rows = conn.execute(
+            "SELECT id, file_path, chunk_type, name, text, blob_oid, action_hash "
+            "FROM vectors WHERE id > ? ORDER BY id LIMIT ?",
+            (last_cursor, batch_size),
+        ).fetchall()
+
+        if not rows:
+            # Nothing left to process.
+            self._set_meta("rotation_state", RotationState.READY_TO_CUTOVER.value)
+            conn.commit()
+            return 0
+
+        # Promote to BACKFILLING on the first processed batch.
+        if status.state == RotationState.PENDING:
+            self._set_meta("rotation_state", RotationState.BACKFILLING.value)
+
+        texts = [r[4] for r in rows]
+        try:
+            new_vectors = provider.embed(texts)
+        except Exception as exc:  # noqa: BLE001 — surface provider errors loudly
+            self._set_meta("rotation_last_error", f"embed failed: {exc}")
+            self._set_meta("rotation_state", RotationState.FAILED.value)
+            conn.commit()
+            raise
+
+        if len(new_vectors) != len(rows):
+            self._set_meta(
+                "rotation_last_error",
+                f"embed returned {len(new_vectors)} vectors for {len(rows)} inputs",
+            )
+            self._set_meta("rotation_state", RotationState.FAILED.value)
+            conn.commit()
+            raise RuntimeError("embed() length mismatch")
+
+        produced_at = time.time()
+        provider_name = provider.name  # property, not callable
+        to_version = status.to_version
+        to_dim = status.to_dim
+
+        insert_sql = (
+            f"INSERT OR REPLACE INTO {STAGING_TABLE} "
+            "(id, file_path, chunk_type, name, text, vector, "
+            "model_name, model_version, dimension, produced_at, "
+            "blob_oid, action_hash) "
+            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
+        )
+        for (row, vec) in zip(rows, new_vectors, strict=True):
+            if len(vec) != to_dim:
+                self._set_meta(
+                    "rotation_last_error",
+                    f"provider returned vec of len {len(vec)}, expected {to_dim}",
+                )
+                self._set_meta("rotation_state", RotationState.FAILED.value)
+                conn.commit()
+                raise RuntimeError(
+                    f"provider returned vec of wrong length: {len(vec)} != {to_dim}"
+                )
+            packed = struct.pack(f"{to_dim}f", *vec)
+            conn.execute(
+                insert_sql,
+                (
+                    row[0],        # id
+                    row[1],        # file_path
+                    row[2],        # chunk_type
+                    row[3],        # name
+                    row[4],        # text
+                    packed,        # new vector
+                    provider_name, # model_name
+                    to_version,    # model_version
+                    to_dim,        # dimension
+                    produced_at,   # produced_at
+                    row[5],        # blob_oid (carried over)
+                    row[6],        # action_hash (carried over)
+                ),
+            )
+
+        # Update cursor + processed count.
+        last_id = rows[-1][0]
+        processed = status.processed_rows + len(rows)
+        self._set_meta("rotation_cursor", last_id)
+        self._set_meta("rotation_processed_rows", str(processed))
+        conn.commit()
+
+        # If we drained the table, transition to ready_to_cutover.
+        if len(rows) < batch_size:
+            # We got fewer rows than requested, so the table is exhausted.
+            self._set_meta("rotation_state", RotationState.READY_TO_CUTOVER.value)
+            conn.commit()
+
+        return len(rows)
+
+    # ------------------------------------------------------------------
+    # cutover
+    # ------------------------------------------------------------------
+
+    def cutover(self) -> RotationStatus:
+        """Atomically swap ``vectors`` and ``vectors_rotating``.
+
+        The existing ``vectors`` becomes ``vectors_archive`` (still
+        present until ``gc_old``). Subsequent reads / writes hit the
+        new table. ``store_metadata.dimension`` is updated to the new
+        target dim.
+
+        Requires state ``ready_to_cutover``.
+        """
+        conn = self._get_conn()
+        status = self.status()
+        if status.state != RotationState.READY_TO_CUTOVER:
+            raise RuntimeError(
+                f"cutover() requires state ready_to_cutover, got {status.state.value}"
+            )
+
+        try:
+            conn.execute("BEGIN IMMEDIATE")
+            conn.execute(f"DROP TABLE IF EXISTS {ARCHIVE_TABLE}")
+            conn.execute(f"ALTER TABLE vectors RENAME TO {ARCHIVE_TABLE}")
+            conn.execute(f"ALTER TABLE {STAGING_TABLE} RENAME TO vectors")
+            # Update stored dimension pointer so next VectorStore open
+            # doesn't trip the mismatch check.
+            conn.execute(
+                "INSERT OR REPLACE INTO store_metadata (key, value) VALUES ('dimension', ?)",
+                (str(status.to_dim),),
+            )
+            conn.execute(
+                "INSERT OR REPLACE INTO store_metadata (key, value) VALUES (?, ?)",
+                ("rotation_state", RotationState.CUTOVER_DONE.value),
+            )
+            # Invalidate the vector cache version so open VectorStores
+            # reload. Uses a monotonic bump stored in store_metadata.
+            prev_ver = int(self._get_meta("_rotator_cache_ver") or "0")
+            conn.execute(
+                "INSERT OR REPLACE INTO store_metadata (key, value) VALUES (?, ?)",
+                ("_rotator_cache_ver", str(prev_ver + 1)),
+            )
+            conn.commit()
+        except sqlite3.Error:
+            conn.execute("ROLLBACK")
+            self._set_meta("rotation_state", RotationState.FAILED.value)
+            self._set_meta("rotation_last_error", "cutover SQL failed")
+            conn.commit()
+            raise
+
+        return self.status()
+
+    # ------------------------------------------------------------------
+    # gc_old
+    # ------------------------------------------------------------------
+
+    def gc_old(self) -> RotationStatus:
+        """Drop the archive table and clear rotation metadata.
+
+        Requires state ``cutover_done``. Transitions through ``gc_done``
+        back to ``none``.
+        """
+        conn = self._get_conn()
+        status = self.status()
+        if status.state != RotationState.CUTOVER_DONE:
+            raise RuntimeError(
+                f"gc_old() requires state cutover_done, got {status.state.value}"
+            )
+        conn.execute(f"DROP TABLE IF EXISTS {ARCHIVE_TABLE}")
+        self._set_meta("rotation_state", RotationState.GC_DONE.value)
+        conn.commit()
+        # Then clean the rotation metadata entirely so state goes back to NONE.
+        self._clear_rotation_meta()
+        conn.commit()
+        return self.status()
+
+    # ------------------------------------------------------------------
+    # abort
+    # ------------------------------------------------------------------
+
+    def abort(self) -> RotationStatus:
+        """Cancel a pre-cutover rotation.
+
+        Drops ``vectors_rotating`` and clears rotation metadata. The
+        primary ``vectors`` table is untouched. Legal from any state
+        other than ``cutover_done`` / ``gc_done``.
+        """
+        conn = self._get_conn()
+        status = self.status()
+        if status.state in (RotationState.CUTOVER_DONE, RotationState.GC_DONE):
+            raise RuntimeError(
+                f"abort() would corrupt state {status.state.value}; "
+                "use gc_old instead or restore from a snapshot"
+            )
+        conn.execute(f"DROP TABLE IF EXISTS {STAGING_TABLE}")
+        self._set_meta("rotation_state", RotationState.ABORTED.value)
+        conn.commit()
+        # Clear metadata back to NONE so a fresh rotation can begin.
+        self._clear_rotation_meta()
+        conn.commit()
+        return self.status()