coderay 1.0.0__py3-none-any.whl

coderay/state/machine.py

@@ -0,0 +1,242 @@
+from __future__ import annotations
+
+import json
+import time
+from dataclasses import asdict, dataclass, field
+from enum import Enum
+from pathlib import Path
+
+META_FILENAME = "meta.json"
+FILE_HASHES_FILENAME = "file_hashes.json"
+
+
+class MetaState(str, Enum):
+    """State of the indexer run."""
+
+    IN_PROGRESS = "in_progress"
+    DONE = "done"
+    ERRORED = "errored"
+    INCOMPLETE = "incomplete"
+
+
+@dataclass
+class CurrentRun:
+    """Resume info for the current indexing run."""
+
+    paths_to_process: list[str] = field(default_factory=list)
+    processed_count: int = 0
+    error: str | None = None
+
+
+@dataclass(frozen=True)
+class IndexMeta:
+    """Immutable snapshot of meta.json contents."""
+
+    state: MetaState
+    started_at: float
+    last_commit: str
+    branch: str
+    indexed_at: float
+    current_run: CurrentRun
+    error: str | None = None
+
+    def is_in_progress(self):
+        """Return True if state is IN_PROGRESS."""
+        return self.state == MetaState.IN_PROGRESS
+
+    def is_incomplete(self):
+        """Return True if state is INCOMPLETE (previous run did not finish)."""
+        return self.state == MetaState.INCOMPLETE
+
+
+class StateMachine:
+    """Manages index metadata and file hashes on disk."""
+
+    def __init__(self, index_dir: Path | str) -> None:
+        """Initialize with index dir (contains meta.json and file_hashes.json)."""
+        self._index_dir = Path(index_dir)
+        self._meta_path = self._index_dir / META_FILENAME
+        self._file_hashes_path = self._index_dir / FILE_HASHES_FILENAME
+        self._current_state: IndexMeta | None = self._load_meta()
+        self._file_hashes: dict[str, str] = self._load_file_hashes()
+
+    @property
+    def index_dir(self) -> Path:
+        """Path to the index directory."""
+        return self._index_dir
+
+    @property
+    def meta_path(self) -> Path:
+        """Path to meta.json."""
+        return self._meta_path
+
+    @property
+    def current_state(self) -> IndexMeta | None:
+        """Current meta state; None if meta.json is missing or invalid."""
+        return self._current_state
+
+    @current_state.setter
+    def current_state(self, value: IndexMeta) -> None:
+        self._current_state = value
+
+    @property
+    def file_hashes(self) -> dict[str, str]:
+        """Mapping of relative path -> content hash."""
+        return self._file_hashes
+
+    @file_hashes.setter
+    def file_hashes(self, value: dict[str, str]) -> None:
+        self._file_hashes = value
+
+    def _load_meta(self) -> IndexMeta | None:
+        """Load meta from disk. Returns None if file missing or invalid."""
+        if not self._meta_path.exists():
+            return None
+        try:
+            data = json.loads(self._meta_path.read_text())
+            data["state"] = MetaState(data["state"])
+            cr = data.get("current_run") or {}
+            data["current_run"] = CurrentRun(
+                paths_to_process=cr.get("paths_to_process") or [],
+                processed_count=int(cr.get("processed_count") or 0),
+                error=cr.get("error"),
+            )
+            return IndexMeta(**data)
+        except Exception:
+            return None
+
+    def _load_file_hashes(self) -> dict[str, str]:
+        """Load file hashes from disk. Returns {} if missing or invalid."""
+        if not self._file_hashes_path.exists():
+            return {}
+        try:
+            data = json.loads(self._file_hashes_path.read_text())
+            return dict(data) if isinstance(data, dict) else {}
+        except Exception:
+            return {}
+
+    @property
+    def is_in_progress(self) -> bool:
+        """True if state is in_progress or incomplete."""
+        if self._current_state is None:
+            return False
+        return (
+            self._current_state.is_in_progress() or self._current_state.is_incomplete()
+        )
+
+    @property
+    def has_partial_progress(self) -> bool:
+        """True if we have a run with paths and non-zero processed count."""
+        if self._current_state is None:
+            return False
+        run = self._current_state.current_run
+        return bool(run.paths_to_process and run.processed_count > 0)
+
+    def set_incomplete(self) -> None:
+        """Mark the current run as incomplete."""
+        if self._current_state is None or not self.is_in_progress:
+            return
+        data = asdict(self._current_state)
+        data["state"] = MetaState.INCOMPLETE
+        data["current_run"] = asdict(self._current_state.current_run)
+        self._current_state = IndexMeta(
+            state=data["state"],
+            started_at=data["started_at"],
+            last_commit=data["last_commit"],
+            branch=data["branch"],
+            indexed_at=data["indexed_at"],
+            current_run=CurrentRun(**data["current_run"]),
+        )
+        self._save_meta()
+
+    def set_errored(self, exc: str) -> None:
+        if self._current_state is None:
+            return
+        data = asdict(self._current_state)
+        current_run = asdict(self._current_state.current_run)
+        self._current_state = IndexMeta(
+            state=MetaState.ERRORED,
+            error=exc,
+            started_at=data["started_at"],
+            last_commit=data["last_commit"],
+            branch=data["branch"],
+            indexed_at=data["indexed_at"],
+            current_run=CurrentRun(
+                paths_to_process=current_run["paths_to_process"],
+                processed_count=current_run["processed_count"],
+                error=exc,
+            ),
+        )
+        self._save_meta()
+
+    def start(self, branch: str | None, last_commit: str | None) -> None:
+        """Start a new indexing run."""
+        now = time.time()
+        self._current_state = IndexMeta(
+            state=MetaState.IN_PROGRESS,
+            started_at=now,
+            last_commit=last_commit or "",
+            branch=branch or "",
+            indexed_at=now,
+            current_run=CurrentRun(),
+        )
+        self._save_meta()
+
+    def save_progress(
+        self,
+        full_rel_paths: list[str],
+        processed_count: int,
+    ) -> None:
+        """Save resume checkpoint with processed count."""
+        if self._current_state is None or not self.is_in_progress:
+            return
+        run = CurrentRun(
+            paths_to_process=full_rel_paths,
+            processed_count=processed_count,
+        )
+        self._current_state = IndexMeta(
+            state=self._current_state.state,
+            started_at=self._current_state.started_at,
+            last_commit=self._current_state.last_commit,
+            branch=self._current_state.branch,
+            indexed_at=self._current_state.indexed_at,
+            current_run=run,
+        )
+        self._save_meta()
+
+    def finish(
+        self,
+        last_commit: str | None = None,
+        branch: str | None = None,
+    ) -> None:
+        """Mark the current run as done and persist state."""
+        self._save_file_hashes()
+        if self._current_state is None:
+            return
+        commit = (
+            last_commit if last_commit is not None else self._current_state.last_commit
+        )
+        br = branch if branch is not None else self._current_state.branch
+        self._current_state = IndexMeta(
+            state=MetaState.DONE,
+            started_at=self._current_state.started_at,
+            last_commit=commit,
+            branch=br,
+            indexed_at=time.time(),
+            current_run=CurrentRun(),
+        )
+        self._save_meta()
+
+    def _save_meta(self) -> None:
+        """Write current state to meta.json."""
+        self._meta_path.parent.mkdir(parents=True, exist_ok=True)
+        data = asdict(self._current_state)
+        data["state"] = self._current_state.state.value
+        self._meta_path.write_text(json.dumps(data, indent=2))
+
+    def _save_file_hashes(self) -> None:
+        """Write file_hashes to file_hashes.json."""
+        self._file_hashes_path.parent.mkdir(parents=True, exist_ok=True)
+        self._file_hashes_path.write_text(
+            json.dumps(self._file_hashes, indent=0, sort_keys=True)
+        )

coderay/state/version.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+INDEX_SCHEMA_VERSION = 2
+VERSION_FILENAME = "version.json"
+
+
+class IndexVersionError(Exception):
+    """Raised when the index schema version doesn't match the current code."""
+
+
+def write_index_version(index_dir: str | Path) -> None:
+    """Write the current schema version to index_dir/version.json."""
+    path = Path(index_dir) / VERSION_FILENAME
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps({"schema_version": INDEX_SCHEMA_VERSION}))
+
+
+def read_index_version(index_dir: str | Path) -> int | None:
+    """Read the schema version from index_dir/version.json. Returns None if missing."""
+    path = Path(index_dir) / VERSION_FILENAME
+    if not path.is_file():
+        return None
+    try:
+        data = json.loads(path.read_text())
+        return int(data.get("schema_version", 0))
+    except Exception:
+        return None
+
+
+def check_index_version(index_dir: str | Path) -> None:
+    """Warn if the index version doesn't match the current schema."""
+    version = read_index_version(index_dir)
+    if version is None:
+        return
+    if version != INDEX_SCHEMA_VERSION:
+        logger.warning(
+            "Index schema version mismatch: index=%s, code=%s. "
+            "Consider rebuilding with 'coderay build --full'.",
+            version,
+            INDEX_SCHEMA_VERSION,
+        )

coderay/storage/lancedb.py

@@ -0,0 +1,268 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any
+
+import lancedb
+
+from coderay.core.models import Chunk
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_DIMENSIONS = 384
+TABLE_NAME = "chunks"
+DEFAULT_DISTANCE_METRIC = "cosine"
+
+
+def index_exists(index_dir: str | Path) -> bool:
+    """True if a LanceDB index (chunks table) exists at index_dir."""
+    path = Path(index_dir)
+    return (path / f"{TABLE_NAME}.lance").is_dir()
+
+
+class Store:
+    """LanceDB-backed vector store for code chunks."""
+
+    def __init__(self, db_path: str | Path, dimensions: int = DEFAULT_DIMENSIONS):
+        """Initialize the LanceDB store."""
+        self.db_path = Path(db_path)
+        self.dimensions = dimensions
+        self._ensure_dir()
+        self._db = lancedb.connect(str(self.db_path))
+        self._table_known = False
+        self._fts_stale = True
+
+    def _ensure_dir(self) -> None:
+        self.db_path.mkdir(parents=True, exist_ok=True)
+
+    def _table_exists(self) -> bool:
+        if self._table_known:
+            return True
+        resp = self._db.list_tables()
+        tables = resp.tables if hasattr(resp, "tables") else list(resp)
+        exists = TABLE_NAME in tables
+        if exists:
+            self._table_known = True
+        return exists
+
+    def _rows_from_chunks_embeddings(
+        self,
+        chunks: list[Chunk],
+        embeddings: list[list[float]],
+    ) -> list[dict[str, Any]]:
+        rows = []
+        for chunk, emb in zip(chunks, embeddings):
+            if len(emb) != self.dimensions:
+                raise ValueError(
+                    f"Embedding dimension {len(emb)} "
+                    f"!= store dimension {self.dimensions}"
+                )
+            rows.append(
+                {
+                    "path": chunk.path,
+                    "start_line": chunk.start_line,
+                    "end_line": chunk.end_line,
+                    "symbol": chunk.symbol,
+                    "language": chunk.language,
+                    "content": chunk.content,
+                    "vector": emb,
+                }
+            )
+        return rows
+
+    def _get_table(self):
+        return self._db.open_table(TABLE_NAME)
+
+    def _ensure_fts_index(self, table) -> None:
+        """Create or replace the full-text search index on the content column."""
+        try:
+            table.create_fts_index("content", replace=True)
+        except Exception as exc:
+            logger.debug("FTS index creation skipped: %s", exc)
+
+    def insert_chunks(
+        self,
+        chunks: list[Chunk],
+        embeddings: list[list[float]],
+    ) -> None:
+        """Insert chunks and their embeddings. Lengths must match."""
+        if len(chunks) != len(embeddings):
+            raise ValueError("chunks and embeddings length mismatch")
+        if not chunks:
+            return
+        rows = self._rows_from_chunks_embeddings(chunks, embeddings)
+        if not self._table_exists():
+            self._db.create_table(TABLE_NAME, rows)
+            self._table_known = True
+        else:
+            self._get_table().add(rows)
+        self._fts_stale = True
+
+    def delete_by_paths(self, paths: list[str]) -> None:
+        """Remove all chunks whose path is in the given list."""
+        if not paths:
+            return
+        if not self._table_exists():
+            return
+        table = self._get_table()
+        safe = [p.replace("'", "''") for p in paths]
+        quoted = ", ".join(f"'{p}'" for p in safe)
+        table.delete(f"path IN ({quoted})")
+
+    def search(
+        self,
+        query_embedding: list[float],
+        top_k: int = 10,
+        path_prefix: str | None = None,
+        language: str | None = None,
+        query_text: str | None = None,
+    ) -> list[dict[str, Any]]:
+        """Nearest-neighbor search with optional hybrid scoring."""
+        if not self._table_exists():
+            return []
+        table = self._get_table()
+
+        use_hybrid = bool(query_text)
+        if use_hybrid:
+            if self._fts_stale:
+                self._ensure_fts_index(table)
+                self._fts_stale = False
+            try:
+                query = (
+                    table.search(query_type="hybrid")
+                    .vector(query_embedding)
+                    .distance_type(DEFAULT_DISTANCE_METRIC)
+                    .text(query_text)
+                )
+            except Exception:
+                query = table.search(query_embedding).distance_type(
+                    DEFAULT_DISTANCE_METRIC
+                )
+                use_hybrid = False
+        else:
+            query = table.search(query_embedding).distance_type(DEFAULT_DISTANCE_METRIC)
+
+        if path_prefix:
+            prefix = (path_prefix.rstrip("/") + "/").replace("'", "''")
+            query = query.where(f"path LIKE '{prefix}%'")
+        if language:
+            lang = (language or "").replace("'", "''")
+            query = query.where(f"language = '{lang}'")
+
+        query = query.limit(top_k)
+        rows = query.to_list()
+
+        results = []
+        for r in rows:
+            row = dict(r)
+            if "_relevance_score" in row:
+                score = row.pop("_relevance_score")
+                row.pop("_distance", None)
+            elif "_distance" in row:
+                score = 1.0 - row.pop("_distance")
+            else:
+                score = row.pop("distance", 0.0)
+            row["score"] = round(float(score), 4)
+            row["score_type"] = "rrf" if use_hybrid else "cosine"
+            row.pop("vector", None)
+            results.append(row)
+
+        return results
+
+    def chunk_count(self) -> int:
+        """Total number of chunks in the store."""
+        if not self._table_exists():
+            return 0
+        return self._get_table().count_rows()
+
+    def list_chunks(
+        self,
+        limit: int = 500,
+        path_prefix: str | None = None,
+    ) -> list[dict[str, Any]]:
+        """List indexed chunks (no vectors). For visualization / debugging."""
+        if not self._table_exists():
+            return []
+        table = self._get_table()
+        n = table.count_rows()
+        if n == 0:
+            return []
+
+        col_names = ["path", "start_line", "end_line", "symbol", "language"]
+
+        if path_prefix:
+            prefix = (path_prefix.rstrip("/") + "/").replace("'", "''")
+            try:
+                arrow = (
+                    table.search()
+                    .where(f"path LIKE '{prefix}%'")
+                    .select(col_names)
+                    .limit(limit)
+                    .to_arrow()
+                )
+            except Exception:
+                arrow = table.head(min(n, limit * 2))
+                arrow = arrow.select(col_names)
+                rows = arrow.to_pylist()
+                pfix = path_prefix.rstrip("/") + "/"
+                return [r for r in rows if str(r.get("path", "")).startswith(pfix)][
+                    :limit
+                ]
+        else:
+            to_read = min(n, limit)
+            arrow = table.head(to_read)
+            arrow = arrow.select(col_names)
+
+        return arrow.to_pylist()[:limit]
+
+    def chunks_by_path(self) -> dict[str, int]:
+        """Return mapping of file path -> chunk count for the whole index."""
+        if not self._table_exists():
+            return {}
+        table = self._get_table()
+        n = table.count_rows()
+        if n == 0:
+            return {}
+        arrow = table.head(n).select(["path"])
+        paths = arrow.column("path").to_pylist()
+        counts: dict[str, int] = {}
+        for p in paths:
+            key = str(p) if p is not None else "?"
+            counts[key] = counts.get(key, 0) + 1
+        return counts
+
+    def maintain(self) -> dict[str, Any]:
+        """Run maintenance on the chunks table to reclaim space."""
+        result: dict[str, Any] = {"cleanup_done": False, "compact_done": False}
+        if not self._table_exists():
+            return result
+        table = self._get_table()
+        try:
+            dataset = table.to_lance()
+        except Exception as e:
+            logger.warning("to_lance failed (install pylance?): %s", e)
+            result["error_cleanup"] = str(e)
+            return result
+        try:
+            dataset.cleanup_old_versions(retain_versions=1)
+            result["cleanup_done"] = True
+            logger.info("Cleaned up old table versions")
+        except Exception as e:
+            logger.warning("cleanup_old_versions failed: %s", e)
+            result["error_cleanup"] = str(e)
+        try:
+            dataset.optimize.compact_files()
+            result["compact_done"] = True
+            logger.info("Compacted table fragments")
+        except Exception as e:
+            logger.warning("compact_files failed: %s", e)
+            result["error_compact"] = str(e)
+        return result
+
+    def clear(self) -> None:
+        """Drop table so next insert_chunks creates a fresh one (full rebuild)."""
+        if self._table_exists():
+            self._db.drop_table(TABLE_NAME)
+            self._table_known = False
+            self._fts_stale = True