vecshift 0.1.0__py3-none-any.whl

vecshift/jobs/spec.py

@@ -0,0 +1,202 @@
+"""The job spec, usually ``vecshift.yaml``.
+
+One file describes a migration: where the vectors are, where the new ones go, which model
+makes them, and the limits to respect. The CLI reads it today; the API and UI will read and
+write the same file. It never holds credentials: the database connection string comes
+from an environment variable the file names.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from enum import StrEnum
+from pathlib import Path
+from typing import Literal
+
+import yaml
+from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator
+
+from vecshift.embeddings.spec import SpecError, parse_spec
+
+_IDENT = re.compile(r"^[A-Za-z_][A-Za-z0-9_$]*$")
+
+
+class JobError(ValueError):
+    """The job file couldn't be read or isn't valid."""
+
+
+class _Strict(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+
+def _identifier(value: str | None, what: str) -> str | None:
+    if value is not None and not _IDENT.match(value):
+        raise ValueError(f"{what} must be a plain identifier, got {value!r}")
+    return value
+
+
+class Source(_Strict):
+    type: Literal["pgvector"] = "pgvector"
+    dsn_env: str = "VECSHIFT_DSN"
+    """Environment variable holding the connection string. The file never stores it."""
+    table: str
+    """``table`` or ``schema.table``."""
+    vector_column: str | None = None
+    text_column: str | None = None
+    model: str | None = None
+    """The model that made the current vectors, so ``eval`` can embed queries for them."""
+
+    @field_validator("table")
+    @classmethod
+    def _table(cls, value: str) -> str:
+        for part in value.split("."):
+            _identifier(part, "table")
+        if value.count(".") > 1:
+            raise ValueError("table must be table or schema.table")
+        return value
+
+    @field_validator("vector_column", "text_column")
+    @classmethod
+    def _columns(cls, value: str | None) -> str | None:
+        return _identifier(value, "column")
+
+    @field_validator("model")
+    @classmethod
+    def _old_model(cls, value: str | None) -> str | None:
+        if value is not None:
+            try:
+                parse_spec(value)
+            except SpecError as exc:
+                raise ValueError(str(exc)) from None
+        return value
+
+
+class VectorType(StrEnum):
+    VECTOR = "vector"
+    HALFVEC = "halfvec"
+
+
+class IndexMethod(StrEnum):
+    HNSW = "hnsw"
+    IVFFLAT = "ivfflat"
+    NONE = "none"
+
+
+class Metric(StrEnum):
+    COSINE = "cosine"
+    INNER_PRODUCT = "inner_product"
+    L2 = "l2"
+
+
+class Target(_Strict):
+    column: str = "embedding_v2"
+    """New column on the source table. Cutover swaps it with the old one by renaming."""
+    vector_type: VectorType = VectorType.VECTOR
+    index: IndexMethod = IndexMethod.HNSW
+    metric: Metric | None = None
+    """Distance for the index. Defaults to the existing index's metric, or cosine."""
+
+    @field_validator("column")
+    @classmethod
+    def _column(cls, value: str) -> str:
+        _identifier(value, "column")
+        return value
+
+
+class Limits(_Strict):
+    budget_usd: float | None = Field(default=None, gt=0)
+    """Stop before spending more than this. Plan fails when the estimate exceeds it."""
+    tokens_per_minute: int | None = Field(default=None, gt=0)
+    """The provider's rate limit, used to estimate duration."""
+    requests_per_minute: int | None = Field(default=None, gt=0)
+
+
+class JobSpec(_Strict):
+    version: Literal[1] = 1
+    name: str = "migration"
+    source: Source
+    target: Target = Target()
+    model: str
+    """A model spec, as for ``vecshift bench``: ``provider/model[,option=value]``."""
+    limits: Limits = Limits()
+
+    @field_validator("target", "limits", mode="before")
+    @classmethod
+    def _empty_section(cls, value: object) -> object:
+        # A section whose settings are all commented out reads as null: use the defaults.
+        return {} if value is None else value
+
+    @field_validator("model")
+    @classmethod
+    def _model(cls, value: str) -> str:
+        try:
+            parse_spec(value)
+        except SpecError as exc:
+            raise ValueError(str(exc)) from None
+        return value
+
+
+def format_errors(error: ValidationError) -> str:
+    lines = []
+    for item in error.errors():
+        where = ".".join(str(p) for p in item["loc"]) or "file"
+        message = item["msg"].removeprefix("Value error, ")
+        if item["type"] == "extra_forbidden":
+            message = "unknown setting"
+        elif item["type"] == "missing":
+            message = "required"
+        lines.append(f"  {where}: {message}")
+    return "\n".join(lines)
+
+
+def load_job(path: Path) -> JobSpec:
+    try:
+        raw = yaml.safe_load(path.read_text(encoding="utf-8"))
+    except OSError as exc:
+        raise JobError(f"Couldn't read {path}: {exc.strerror}") from exc
+    except yaml.YAMLError as exc:
+        mark = getattr(exc, "problem_mark", None)
+        where = f" at line {mark.line + 1}" if mark else ""
+        raise JobError(f"{path} isn't valid YAML{where}") from exc
+    if not isinstance(raw, dict):
+        raise JobError(f"{path} should contain settings, starting with source: and model:")
+    try:
+        return JobSpec.model_validate(raw)
+    except ValidationError as exc:
+        raise JobError(f"{path} has problems:\n{format_errors(exc)}") from exc
+
+
+def template(table: str, model: str, column: str = "embedding_v2") -> str:
+    """A commented job file to start from. Values are quoted, so any input stays valid YAML."""
+    name = json.dumps(table.split(".")[-1] + "-reembed")
+    return f"""# vecshift job: re-embed a pgvector column with a new model, side by side.
+# Check it with `vecshift plan`. Nothing is changed until `vecshift apply`.
+version: 1
+name: {name}
+
+source:
+  type: pgvector
+  # The connection string is read from this environment variable, never from this file.
+  dsn_env: VECSHIFT_DSN
+  table: {json.dumps(table)}
+  # vector_column: embedding     # needed only if the table has several
+  # text_column: content         # detected automatically when it's a common name
+  # model: openai/text-embedding-3-small  # the model that made the current vectors, for eval
+
+target:
+  # New vectors go in this column next to the old ones. Cutover renames the two columns
+  # in one transaction, and rollback renames them back.
+  column: {json.dumps(column)}
+  vector_type: vector            # or halfvec: half the storage, indexable up to 4,000 dims
+  index: hnsw                    # hnsw, ivfflat, or none
+  # metric: cosine               # defaults to the current index's metric, or cosine
+
+# The new model, written as for `vecshift bench`, e.g. openai/text-embedding-3-large,dims=1024
+model: {json.dumps(model)}
+
+limits:
+  # budget_usd: 50               # plan fails if the estimate is higher; apply stops there
+  # tokens_per_minute: 1000000   # your provider rate limit, for the duration estimate
+  # requests_per_minute: 3000
+"""

vecshift/migrate/__init__.py

@@ -0,0 +1,6 @@
+"""Carrying out a migration: backfill, keep in sync, index, and verify."""
+
+from vecshift.migrate.engine import ApplyResult, Event, apply
+from vecshift.migrate.state import JobState
+
+__all__ = ["ApplyResult", "Event", "JobState", "apply"]

vecshift/migrate/engine.py

@@ -0,0 +1,272 @@
+"""The apply loop: fill the new column, keep up with writes, then index and verify.
+
+It talks to the database through a small writer interface and to the model through the
+embedding contract, so it can be tested with fakes and reused for other stores.
+"""
+
+from __future__ import annotations
+
+import math
+import time
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, field
+from typing import Any, Literal, Protocol
+
+from vecshift.embeddings.providers import EmbeddingError
+from vecshift.migrate.state import JobState
+
+CHARS_PER_TOKEN = 4
+MAX_PASSES = 5
+SETTLED_ROWS = 100
+SETTLED_SHARE = 0.001
+"""On a busy table some rows always change while a pass runs. Once no more than this many
+remain (at least ``SETTLED_ROWS``), apply builds the index and finishes; the stragglers get
+their vector on the next run, and cutover embeds them before it switches."""
+"""Catch-up passes after the first, for rows that changed or arrived during the run."""
+MAX_CHARS = 24_000
+"""Longer texts are cut to this length, inside common model input limits (about 6K tokens)."""
+
+Status = Literal["complete", "stopped", "budget", "failed"]
+
+
+class Row(Protocol):
+    @property
+    def key(self) -> Any: ...
+    @property
+    def text(self) -> str: ...
+    @property
+    def id(self) -> str: ...
+
+
+class Writer(Protocol):
+    def acquire(self) -> None: ...
+    def release(self) -> None: ...
+    def ensure_column(self) -> bool: ...
+    def ensure_trigger(self) -> bool: ...
+    def total_count(self) -> int: ...
+    def pending_count(self, exclude: Sequence[str] = ()) -> int: ...
+    def fetch(self, after: Any, limit: int, exclude: Sequence[str] = ()) -> Sequence[Row]: ...
+    def write(self, items: Sequence[tuple[Any, Sequence[float]]]) -> int: ...
+    def build_index(self, method: str, metric: str, rows: int) -> str: ...
+    def dimensions_in_use(self) -> set[int]: ...
+
+
+class Embedder(Protocol):
+    tokens: int
+
+    async def embed(
+        self, texts: Sequence[str], mode: Literal["document", "query"] = ...
+    ) -> list[list[float]]: ...
+
+
+@dataclass(frozen=True, slots=True)
+class Event:
+    kind: str
+    """``start``, ``column``, ``trigger``, ``batch``, ``pass``, ``index``, or ``done``."""
+    data: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass(slots=True)
+class ApplyResult:
+    status: Status
+    rows_written: int = 0
+    rows_failed: int = 0
+    remaining: int = 0
+    tokens: int = 0
+    spent_usd: float | None = 0.0
+    """This run's spend, or ``None`` when the model has no price."""
+    total_spent_usd: float | None = 0.0
+    seconds: float = 0.0
+    index: str = "pending"
+    """``built``, ``rebuilt``, ``exists``, ``skipped``, or ``pending``."""
+    message: str | None = None
+
+
+class DimensionMismatch(Exception):
+    pass
+
+
+async def _embed_isolating(
+    embedder: Embedder, rows: Sequence[Row], dims: int
+) -> tuple[list[tuple[Row, list[float]]], dict[str, str]]:
+    """Embed rows; if the provider rejects a request, split it to find the bad rows."""
+    texts = [r.text[:MAX_CHARS] for r in rows]
+    try:
+        vectors = await embedder.embed(texts, "document")
+    except EmbeddingError as exc:
+        if len(rows) == 1:
+            return [], {rows[0].id: str(exc)}
+        mid = len(rows) // 2
+        left, lf = await _embed_isolating(embedder, rows[:mid], dims)
+        right, rf = await _embed_isolating(embedder, rows[mid:], dims)
+        return left + right, {**lf, **rf}
+    for vector in vectors:
+        if len(vector) != dims:
+            raise DimensionMismatch(
+                f"The model returned {len(vector)} dimensions, but the column holds {dims}."
+            )
+    return list(zip(rows, vectors, strict=True)), {}
+
+
+async def apply(
+    writer: Writer,
+    embedder: Embedder,
+    state: JobState,
+    *,
+    dims: int,
+    price_per_million: float | None,
+    budget_usd: float | None,
+    chunk_rows: int,
+    index: tuple[str, str] | None,
+    on_event: Callable[[Event], None] = lambda e: None,
+    should_stop: Callable[[], bool] = lambda: False,
+    until: float | None = None,
+) -> ApplyResult:
+    """Run (or resume) a migration. Safe to run again at any time to catch up.
+
+    ``until`` (a fraction, such as 0.5) stops once that share of the rows with text has a
+    new vector, so a migration can be checked part way and continued later.
+    """
+    started = time.monotonic()
+    result = ApplyResult(status="complete")
+    tokens_before = embedder.tokens
+    spent_before = state.spent_usd
+    state.runs += 1
+
+    def spent() -> float:
+        if price_per_million is None:
+            return 0.0
+        return (embedder.tokens - tokens_before) * price_per_million / 1_000_000
+
+    def finish(status: Status, message: str | None = None) -> ApplyResult:
+        result.status = status
+        result.message = message
+        result.tokens = embedder.tokens - tokens_before
+        run_spend = spent()
+        result.spent_usd = run_spend if price_per_million is not None else None
+        state.spent_usd = spent_before + run_spend
+        state.tokens += result.tokens
+        state.save()
+        result.total_spent_usd = state.spent_usd if price_per_million is not None else None
+        result.rows_failed = len(state.failed)
+        result.seconds = time.monotonic() - started
+        on_event(Event("done", {"status": status, "message": message}))
+        return result
+
+    writer.acquire()
+    try:
+        # The column has to exist before rows can be counted against it.
+        on_event(Event("column", {"added": writer.ensure_column()}))
+        on_event(Event("trigger", {"added": writer.ensure_trigger()}))
+        pending = writer.pending_count(tuple(state.failed))
+        quota: int | None = None
+        if until is not None:
+            total = writer.total_count()
+            done = total - writer.pending_count()
+            quota = max(0, math.ceil(total * until) - done)
+            pending = min(pending, quota)
+        on_event(Event("start", {"pending": pending, "spent_before": spent_before}))
+        settled = max(SETTLED_ROWS, math.ceil(SETTLED_SHARE * max(pending, state.rows_written)))
+
+        backlog = pending
+        for number in range(1, MAX_PASSES + 2):
+            after: Any = None
+            progress = 0
+            # A pass covers what was pending when it began; rows that arrive meanwhile wait
+            # for the next pass, so a pass ends even if inserts outpace embedding.
+            budget = backlog + chunk_rows
+            fetched = 0
+            while fetched < budget:
+                if should_stop():
+                    result.remaining = writer.pending_count(tuple(state.failed))
+                    return finish("stopped", "Stopped. Run apply again to continue.")
+                limit = chunk_rows
+                if quota is not None and result.rows_written >= quota:
+                    result.remaining = writer.pending_count(tuple(state.failed))
+                    if result.remaining:
+                        return finish(
+                            "stopped",
+                            f"Reached {until or 0:.0%} of rows. Check the new vectors, then "
+                            "run apply again to continue.",
+                        )
+                    break
+                if quota is not None:
+                    limit = min(chunk_rows, quota - result.rows_written)
+                rows = writer.fetch(after, limit, tuple(state.failed))
+                if not rows:
+                    break
+                fetched += len(rows)
+                after = rows[-1].key
+                if budget_usd is not None and price_per_million is not None:
+                    estimate = sum(len(r.text[:MAX_CHARS]) for r in rows) / CHARS_PER_TOKEN
+                    projected = spent_before + spent() + estimate * price_per_million / 1e6
+                    if projected > budget_usd:
+                        result.remaining = writer.pending_count(tuple(state.failed))
+                        return finish(
+                            "budget",
+                            f"Stopped before the next batch would pass the ${budget_usd:,.2f} "
+                            "budget. Raise limits.budget_usd and run apply again to continue.",
+                        )
+                try:
+                    pairs, failed = await _embed_isolating(embedder, rows, dims)
+                except DimensionMismatch as exc:
+                    return finish("failed", str(exc))
+                except EmbeddingError as exc:  # pragma: no cover - isolation catches these
+                    return finish("failed", str(exc))
+                state.failed.update(failed)
+                written = writer.write(pairs)
+                progress += written
+                result.rows_written += written
+                state.rows_written += written
+                state.spent_usd = spent_before + spent()
+                state.save()
+                on_event(
+                    Event(
+                        "batch",
+                        {
+                            "pass": number,
+                            "written": written,
+                            "skipped": len(pairs) - written,
+                            "failed": len(failed),
+                            "rows_written": result.rows_written,
+                            "spent_usd": state.spent_usd if price_per_million is not None else None,
+                        },
+                    )
+                )
+            remaining = writer.pending_count(tuple(state.failed))
+            on_event(Event("pass", {"number": number, "remaining": remaining}))
+            result.remaining = backlog = remaining
+            if remaining == 0 or progress == 0:
+                break
+
+        if result.remaining > settled:
+            return finish(
+                "stopped",
+                f"{result.remaining:,} rows kept changing during the run, faster than they "
+                "could be embedded. Run apply again to catch up.",
+            )
+        if index is None:
+            result.index = "skipped"
+        else:
+            on_event(Event("index", {"state": "building", "method": index[0]}))
+            result.index = writer.build_index(index[0], index[1], state.rows_written)
+            on_event(Event("index", {"state": result.index}))
+            # A big index takes minutes, and a busy table keeps changing meanwhile.
+            result.remaining = writer.pending_count(tuple(state.failed))
+        sizes = writer.dimensions_in_use()
+        if sizes and sizes != {dims}:
+            return finish("failed", f"Found vectors of sizes {sorted(sizes)} in the new column.")
+        failed_note = (
+            f" {len(state.failed):,} rows were rejected by the provider; see the state file."
+            if state.failed
+            else ""
+        )
+        if result.remaining:
+            return finish(
+                "complete",
+                f"Every row has a new vector except {result.remaining:,} that changed during "
+                f"the run; cutover (or the next apply) embeds them.{failed_note}",
+            )
+        return finish("complete", f"Every row with text has a new vector.{failed_note}")
+    finally:
+        writer.release()

vecshift/migrate/state.py

@@ -0,0 +1,50 @@
+"""Progress that has to survive between runs: spend so far and failed rows.
+
+The database itself records which rows are done (their new vector is filled in), so this
+file only holds what the database can't: money spent across runs, for the budget, and rows
+the provider rejected. It lives in ``.vecshift/`` next to the job file, readable by its owner
+only, and never holds document text.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import asdict, dataclass, field
+from datetime import UTC, datetime
+from pathlib import Path
+
+
+@dataclass(slots=True)
+class JobState:
+    path: Path
+    spent_usd: float = 0.0
+    tokens: int = 0
+    rows_written: int = 0
+    runs: int = 0
+    failed: dict[str, str] = field(default_factory=dict)
+    """Row ID → error, for rows the provider rejected. Text is never stored."""
+    history: list[dict[str, str]] = field(default_factory=list)
+    """Cutovers and rollbacks: what happened, when, and to which columns."""
+    updated_at: str | None = None
+
+    @classmethod
+    def for_job(cls, job_file: Path, name: str) -> JobState:
+        safe = "".join(c if c.isalnum() or c in "-_" else "_" for c in name) or "job"
+        path = job_file.resolve().parent / ".vecshift" / f"{safe}.state.json"
+        if not path.exists():
+            return cls(path=path)
+        data = json.loads(path.read_text(encoding="utf-8"))
+        data.pop("path", None)
+        return cls(path=path, **data)
+
+    def save(self) -> None:
+        self.updated_at = datetime.now(UTC).isoformat(timespec="seconds")
+        self.path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
+        data = {k: v for k, v in asdict(self).items() if k != "path"}
+        tmp = self.path.with_suffix(".tmp")
+        # Write then rename, so a crash never leaves a half-written state file.
+        fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+        with os.fdopen(fd, "w", encoding="utf-8") as fh:
+            json.dump(data, fh, indent=2)
+        os.replace(tmp, self.path)

vecshift/planning/__init__.py

@@ -0,0 +1,14 @@
+"""``vecshift plan``: validate a migration job and estimate it before anything runs."""
+
+from vecshift.planning.plan import Change, Estimates, Plan, ProbeResult
+from vecshift.planning.planner import CHARS_PER_TOKEN, SampleStats, build_plan
+
+__all__ = [
+    "CHARS_PER_TOKEN",
+    "Change",
+    "Estimates",
+    "Plan",
+    "ProbeResult",
+    "SampleStats",
+    "build_plan",
+]

vecshift/planning/plan.py

@@ -0,0 +1,87 @@
+"""The result of planning a migration."""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+from vecshift.doctor.findings import Finding, Severity
+
+
+@dataclass(frozen=True, slots=True)
+class Change:
+    kind: str
+    """``add_column``, ``embed``, ``index``, or ``cutover``."""
+    summary: str
+    sql: str | None = None
+    note: str | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class ProbeResult:
+    """Measured by embedding a few sample documents with the real model."""
+
+    documents: int
+    dimensions: int
+    tokens_per_char: float
+    docs_per_second: float
+
+
+@dataclass(slots=True)
+class Estimates:
+    rows: int | None = None
+    rows_exact: bool = False
+    tokens: int | None = None
+    tokens_method: str | None = None
+    cost_usd: float | None = None
+    requests: int | None = None
+    seconds: float | None = None
+    seconds_method: str | None = None
+    new_bytes: int | None = None
+    old_bytes: int | None = None
+    index_memory_bytes: int | None = None
+    maintenance_work_mem: int | None = None
+
+
+@dataclass(slots=True)
+class Plan:
+    job: str
+    source: str
+    target_column: str
+    model: str
+    dimensions: int | None
+    dimensions_source: str | None
+    """``spec``, ``known``, ``probe``, or ``None`` when unknown."""
+    vector_type: str
+    metric: str = "cosine"
+    """Distance for the new index: the job's, the current index's, or cosine."""
+    changes: list[Change] = field(default_factory=list)
+    estimates: Estimates = field(default_factory=Estimates)
+    findings: list[Finding] = field(default_factory=list)
+
+    @property
+    def errors(self) -> int:
+        return sum(1 for f in self.findings if f.severity is Severity.ERROR)
+
+    @property
+    def ok(self) -> bool:
+        return self.errors == 0
+
+    def sorted_findings(self) -> list[Finding]:
+        return sorted(self.findings, key=lambda f: -f.severity.rank)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "job": self.job,
+            "ok": self.ok,
+            "source": self.source,
+            "target_column": self.target_column,
+            "model": self.model,
+            "dimensions": self.dimensions,
+            "dimensions_source": self.dimensions_source,
+            "vector_type": self.vector_type,
+            "metric": self.metric,
+            "changes": [asdict(c) for c in self.changes],
+            "estimates": asdict(self.estimates),
+            "findings": [f.to_dict() for f in self.sorted_findings()],
+        }