vecshift 0.1.0__py3-none-any.whl

vecshift/cli.py

@@ -0,0 +1,236 @@
+"""Command-line interface."""
+
+from __future__ import annotations
+
+import json
+from enum import StrEnum
+from pathlib import Path
+from typing import Annotated
+
+import typer
+
+from vecshift import __version__
+from vecshift.cli_style import STYLE as _STYLE
+from vecshift.cli_style import warn_if_password_on_command_line
+from vecshift.cli_style import wrap as _wrap
+from vecshift.core.fingerprint import EmbeddingFingerprint
+from vecshift.doctor import Report, Severity, run_checks
+
+app = typer.Typer(
+    name="vecshift",
+    # Tracebacks must never print local variables: they can hold connection strings and keys.
+    pretty_exceptions_show_locals=False,
+    help="Safe, observable embedding migrations for any vector store.",
+    no_args_is_help=True,
+    add_completion=False,
+)
+
+
+def _print_version(value: bool) -> None:
+    if value:
+        typer.echo(f"vecshift {__version__}")
+        raise typer.Exit()
+
+
+@app.callback()
+def main(
+    version: Annotated[
+        bool,
+        typer.Option(
+            "--version",
+            callback=_print_version,
+            is_eager=True,
+            help="Show the version and exit.",
+        ),
+    ] = False,
+) -> None:
+    """Safe, observable embedding migrations for any vector store."""
+
+
+@app.command()
+def fingerprint(
+    provider: Annotated[str, typer.Option(help="Embedding provider, e.g. openai.")],
+    model: Annotated[str, typer.Option(help="Model name, e.g. text-embedding-3-small.")],
+    dimensions: Annotated[int, typer.Option(help="Output dimensions.")],
+    version: Annotated[str | None, typer.Option(help="Model version, if pinned.")] = None,
+    task: Annotated[str | None, typer.Option(help="Task type, e.g. retrieval_document.")] = None,
+    prefix: Annotated[str | None, typer.Option(help="Text prefix, e.g. 'passage: '.")] = None,
+    normalized: Annotated[
+        bool,
+        typer.Option("--normalized/--not-normalized", help="Whether vectors are L2-normalized."),
+    ] = True,
+) -> None:
+    """Print the model tag that identifies an embedding configuration's vector space."""
+    try:
+        fp = EmbeddingFingerprint(
+            provider=provider,
+            model=model,
+            dimensions=dimensions,
+            version=version,
+            task=task,
+            prefix=prefix,
+            normalized=normalized,
+        )
+    except ValueError as exc:
+        raise typer.BadParameter(str(exc)) from exc
+    typer.echo(fp.model_tag)
+
+
+class FailOn(StrEnum):
+    NEVER = "never"
+    WARNING = "warning"
+    ERROR = "error"
+
+
+def _render(report: Report, connection: str) -> None:
+    dims = f"({report.declared_dimensions})" if report.declared_dimensions else ""
+    rows = f"~{report.estimated_rows:,}" if report.estimated_rows is not None else "unknown"
+    typer.secho(f"vecshift doctor · {report.store} via {connection}", bold=True)
+    typer.echo(f"Target  {report.target}  {report.vector_type}{dims}")
+    typer.echo(f"Rows    {rows} (inspected {report.sample_rows:,}, {report.sample_method})")
+    typer.echo()
+    for finding in report.sorted_findings():
+        label, color, _ = _STYLE[finding.severity]
+        typer.secho(f"{label}  ", fg=color, bold=True, nl=False)
+        typer.secho(finding.title, bold=True)
+        typer.echo(_wrap(finding.detail))
+        if finding.hint:
+            typer.secho(_wrap(f"→ {finding.hint}"), dim=True)
+    typer.echo()
+    counts = [(report.count(sev), name) for sev, (_, _, name) in _STYLE.items()]
+    summary = ", ".join(
+        f"{n} {name}{'s' if n != 1 and name in ('error', 'warning') else ''}" for n, name in counts
+    )
+    typer.echo(f"Summary: {summary}")
+
+
+@app.command()
+def doctor(
+    ctx: typer.Context,
+    dsn: Annotated[
+        str,
+        typer.Option(
+            envvar=["VECSHIFT_DSN", "DATABASE_URL"],
+            help="PostgreSQL connection string. Prefer the VECSHIFT_DSN environment variable "
+            "so the password stays out of your shell history.",
+            show_default=False,
+        ),
+    ],
+    table: Annotated[
+        str | None, typer.Option(help="Table to inspect, as table or schema.table.")
+    ] = None,
+    column: Annotated[
+        str | None, typer.Option(help="Vector column, if there's more than one.")
+    ] = None,
+    text_column: Annotated[
+        str | None, typer.Option(help="Column holding the source text, if not detected.")
+    ] = None,
+    sample_size: Annotated[
+        int, typer.Option(min=1, max=100_000, help="Maximum rows to inspect.")
+    ] = 2000,
+    timeout: Annotated[int, typer.Option(min=1, help="Statement timeout in seconds.")] = 60,
+    output_json: Annotated[bool, typer.Option("--json", help="Print the report as JSON.")] = False,
+    html: Annotated[
+        Path | None,
+        typer.Option(
+            "--html",
+            dir_okay=False,
+            help="Also write a self-contained HTML report to this file.",
+            show_default=False,
+        ),
+    ] = None,
+    fail_on: Annotated[
+        FailOn, typer.Option(help="Exit with status 1 if any finding is this severe or worse.")
+    ] = FailOn.NEVER,
+) -> None:
+    """Inspect a pgvector index and report problems. Read-only."""
+    warn_if_password_on_command_line(ctx, dsn)
+    try:
+        from vecshift.connectors import pgvector
+    except ImportError as exc:  # pragma: no cover - a broken install
+        typer.secho(
+            "The PostgreSQL driver is missing. Reinstall: pip install --force-reinstall vecshift",
+            err=True,
+            fg=typer.colors.RED,
+        )
+        raise typer.Exit(2) from exc
+
+    try:
+        settings = pgvector.prepare(dsn)
+        conn = pgvector.connect(settings, statement_timeout_s=timeout)
+    except pgvector.ConnectError as exc:
+        typer.secho(f"Couldn't connect: {exc}", err=True, fg=typer.colors.RED)
+        if exc.hint:
+            typer.echo(f"→ {exc.hint}", err=True)
+        raise typer.Exit(2) from exc
+
+    try:
+        profile = pgvector.inspect(
+            conn, table=table, column=column, text_column=text_column, sample_size=sample_size
+        )
+    except pgvector.TargetSelectionError as exc:
+        typer.secho(str(exc), err=True, fg=typer.colors.RED)
+        if exc.candidates:
+            typer.echo("Vector columns found:", err=True)
+            for c in exc.candidates:
+                dims = f"({c.dimensions})" if c.dimensions else ""
+                typer.echo(f"  {c.qualified}  {c.type}{dims}", err=True)
+        raise typer.Exit(2) from exc
+    except Exception as exc:
+        import psycopg
+
+        if isinstance(exc, psycopg.Error):
+            typer.secho(f"Inspection failed: {exc}", err=True, fg=typer.colors.RED)
+            raise typer.Exit(2) from exc
+        raise
+    finally:
+        conn.rollback()
+        conn.close()
+
+    report = run_checks(profile)
+    if html is not None:
+        from vecshift.doctor.html import render_html
+
+        page = render_html(
+            report,
+            connection=settings.display,
+            connection_kind=settings.description,
+            version=__version__,
+        )
+        try:
+            html.write_text(page, encoding="utf-8")
+        except OSError as exc:
+            typer.secho(f"Couldn't write {html}: {exc.strerror}", err=True, fg=typer.colors.RED)
+            raise typer.Exit(2) from exc
+
+    if output_json:
+        typer.echo(json.dumps({"connection": settings.display, **report.to_dict()}, indent=2))
+    else:
+        _render(report, settings.description)
+    if html is not None:
+        typer.echo(f"HTML report written to {html}", err=output_json)
+
+    if fail_on is not FailOn.NEVER and report.worst.rank >= Severity(fail_on.value).rank:
+        raise typer.Exit(1)
+
+
+from vecshift.cli_bench import bench  # noqa: E402
+
+app.command()(bench)
+
+from vecshift.cli_plan import init, plan  # noqa: E402
+
+app.command()(init)
+app.command()(plan)
+
+from vecshift.cli_apply import apply  # noqa: E402
+from vecshift.cli_cutover import cleanup, cutover, rollback  # noqa: E402
+
+app.command()(apply)
+app.command()(cutover)
+app.command()(rollback)
+app.command()(cleanup)
+
+from vecshift.cli_eval import eval_  # noqa: E402
+
+app.command(name="eval")(eval_)

vecshift/cli_apply.py

@@ -0,0 +1,283 @@
+"""The ``vecshift apply`` command."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import signal
+import sys
+import time
+from pathlib import Path
+from types import FrameType
+from typing import TYPE_CHECKING, Annotated, Any
+
+import typer
+
+from vecshift.cli_plan import DEFAULT_JOB, _fail, _money, _size, prepare
+from vecshift.cli_style import finding_lines
+
+if TYPE_CHECKING:
+    from vecshift.migrate import ApplyResult, Event
+
+EXIT_STOPPED = 3
+"""Stopped on purpose (Ctrl-C or budget): run apply again to continue."""
+
+
+class _Stop:
+    """First Ctrl-C finishes the current batch and stops cleanly; a second one aborts."""
+
+    def __init__(self) -> None:
+        self.requested = False
+
+    def __call__(self, signum: int, frame: FrameType | None) -> None:
+        if self.requested:
+            raise KeyboardInterrupt
+        self.requested = True
+        typer.secho(
+            "\nStopping after the current batch. Press Ctrl-C again to abort now.",
+            err=True,
+            fg=typer.colors.YELLOW,
+        )
+
+
+class _Progress:
+    def __init__(self, live: bool) -> None:
+        self.live = live
+        self.started = time.monotonic()
+        self.pending = 0
+        self.last_print = 0.0
+
+    def __call__(self, event: Event) -> None:
+        d = event.data
+        if event.kind == "start":
+            self.pending = d["pending"]
+            typer.echo(f"{self.pending:,} rows need a new vector.")
+        elif event.kind == "column" and d["added"]:
+            typer.echo("✔ Added the new column.")
+        elif event.kind == "trigger" and d["added"]:
+            typer.echo("✔ Added the sync trigger.")
+        elif event.kind == "batch":
+            self._batch(d)
+        elif event.kind == "pass" and d["number"] > 1 and self.live:
+            typer.echo()
+        elif event.kind == "index":
+            if self.live:
+                typer.echo()
+            if d["state"] == "building":
+                typer.echo(f"Building the {d['method']} index concurrently. Writes continue…")
+            else:
+                typer.echo(f"✔ Index {d['state']}.")
+
+    def _batch(self, d: dict[str, Any]) -> None:
+        done = d["rows_written"]
+        elapsed = max(time.monotonic() - self.started, 1e-6)
+        rate = done / elapsed
+        left = max(self.pending - done, 0)
+        eta = f"~{_eta(left / rate)} left" if rate > 0 and left else ""
+        cost = f" · {_money(d['spent_usd'])}" if d["spent_usd"] is not None else ""
+        line = f"  {done:,} / {self.pending:,} rows · {rate:,.0f} rows/s{cost} {eta}"
+        if self.live:
+            typer.echo("\r" + line.ljust(78), nl=False)
+        elif time.monotonic() - self.last_print > 15 or not left:
+            typer.echo(line)
+            self.last_print = time.monotonic()
+
+
+def _eta(seconds: float) -> str:
+    if seconds < 90:
+        return f"{max(1, round(seconds))} s"
+    if seconds < 5400:
+        return f"{round(seconds / 60)} min"
+    return f"{seconds / 3600:.1f} h"
+
+
+def _summary(result: ApplyResult) -> None:
+    color = {"complete": typer.colors.GREEN, "failed": typer.colors.RED}.get(
+        result.status, typer.colors.YELLOW
+    )
+    typer.echo()
+    typer.secho(f"Apply {result.status}.", fg=color, bold=True)
+    if result.message:
+        typer.echo(result.message)
+    spent = _money(result.spent_usd) if result.spent_usd is not None else "unknown"
+    total = (
+        f" ({_money(result.total_spent_usd)} across runs)"
+        if result.total_spent_usd is not None
+        else ""
+    )
+    typer.echo(
+        f"  Rows written {result.rows_written:,} · tokens {result.tokens:,} · spent {spent}"
+        f"{total} · {result.seconds:,.0f} s"
+    )
+    if result.rows_failed:
+        typer.echo(f"  Rejected by the provider: {result.rows_failed:,} rows (retried next run)")
+    if result.status == "complete":
+        typer.echo(
+            "\nNext: check it's safe to switch with `vecshift cutover --check`, then run "
+            "`vecshift cutover`.\nIt embeds any rows added or edited since this run first."
+        )
+
+
+def apply(
+    job_file: Annotated[Path, typer.Argument(help="The job file.")] = DEFAULT_JOB,
+    yes: Annotated[bool, typer.Option("--yes", "-y", help="Don't ask before starting.")] = False,
+    no_index: Annotated[
+        bool, typer.Option("--no-index", help="Backfill only; build the index on a later run.")
+    ] = False,
+    until: Annotated[
+        int | None,
+        typer.Option(
+            "--until",
+            min=1,
+            max=99,
+            metavar="PERCENT",
+            help="Stop once this share of rows has a new vector, to check before continuing.",
+        ),
+    ] = None,
+    output_json: Annotated[
+        bool, typer.Option("--json", help="Print progress as JSON lines, for scripts and UIs.")
+    ] = False,
+) -> None:
+    """Run (or resume) a migration: add the column, embed every row, and build the index."""
+    import psycopg
+
+    from vecshift.connectors import pgvector
+    from vecshift.connectors.pgvector.writer import AlreadyRunning, Busy, Layout, PgWriter
+    from vecshift.embeddings import EmbeddingError, create_embedder
+    from vecshift.embeddings.providers import CONCURRENCY
+    from vecshift.migrate import JobState
+
+    ready = prepare(job_file)
+    plan, job, spec = ready.plan, ready.job, ready.spec
+    if not plan.ok:
+        typer.secho("The plan has errors, so nothing was changed.", fg=typer.colors.RED, bold=True)
+        finding_lines([f for f in plan.findings if f.severity.rank >= 2])
+        raise typer.Exit(1)
+    if plan.dimensions is None:
+        raise _fail(
+            "The new vector size is unknown.",
+            "Add dims= to the model spec, or check it with: vecshift plan --probe",
+        )
+    if job.limits.budget_usd is not None and spec.price is None:
+        raise _fail(
+            "A budget is set, but the model has no price, so it can't be enforced.",
+            "Add price= (USD per million tokens) to the model spec.",
+        )
+    if len(ready.target.primary_key) != 1 or ready.profile.text_field is None:
+        raise _fail("Apply needs a single-column primary key and a text column.")
+
+    state = JobState.for_job(job_file, job.name)
+    est = plan.estimates
+    warnings = [f for f in plan.findings if f.severity.rank == 2]
+    if not output_json:
+        typer.secho(f"vecshift apply · {job.name}", bold=True)
+        for change in plan.changes:
+            if change.kind != "cutover":
+                typer.echo(f"  · {change.summary}")
+        cost = _money(est.cost_usd)
+        budget = f", budget ${job.limits.budget_usd:,.2f}" if job.limits.budget_usd else ""
+        typer.echo(f"  Estimated cost {cost}{budget}; new column {_size(est.new_bytes)}.")
+        if state.runs:
+            typer.echo(
+                f"  Resuming: {state.rows_written:,} rows written in earlier runs, "
+                f"{_money(state.spent_usd)} spent."
+            )
+        if warnings:
+            finding_lines(warnings)
+    if not yes:
+        where = "your machine" if spec.is_local else spec.url
+        typer.echo(
+            f"\nThis changes {plan.source.split()[0]} and sends row text to {where}.",
+            err=True,
+        )
+        if not sys.stdin.isatty():
+            raise _fail("Not changing the database without confirmation.", "Pass --yes to proceed.")
+        if not typer.confirm("Apply?", err=True):
+            raise typer.Exit(1)
+
+    layout = Layout(
+        schema=ready.target.source.schema,
+        table=ready.target.source.table,
+        pk=ready.target.primary_key[0],
+        text=ready.profile.text_field,
+        target=job.target.column,
+        vector_type=job.target.vector_type.value,
+        dims=plan.dimensions,
+        extension_schema=ready.target.extension_schema,
+    )
+    index = None
+    if not no_index and job.target.index.value != "none":
+        index = (job.target.index.value, plan.metric)
+
+    try:
+        conn = pgvector.connect_writer(ready.settings)
+    except pgvector.ConnectError as exc:
+        raise _fail(f"Couldn't connect: {exc}", exc.hint) from exc
+
+    stop = _Stop()
+    previous = signal.signal(signal.SIGINT, stop)
+    if output_json:
+
+        def on_event(event: Event) -> None:
+            typer.echo(json.dumps({"event": event.kind, **event.data}))
+
+    else:
+        on_event = _Progress(live=sys.stdout.isatty())
+    try:
+        embedder = create_embedder(spec)
+        result = asyncio.run(
+            _run(
+                PgWriter(conn, layout),
+                embedder,
+                state,
+                dims=plan.dimensions,
+                price_per_million=spec.price,
+                budget_usd=job.limits.budget_usd,
+                chunk_rows=spec.batch_size * CONCURRENCY,
+                index=index,
+                on_event=on_event,
+                should_stop=lambda: stop.requested,
+                until=until / 100 if until else None,
+            )
+        )
+    except AlreadyRunning as exc:
+        raise _fail(str(exc), "Wait for it to finish, or stop it first.") from exc
+    except (Busy, EmbeddingError) as exc:
+        raise _fail(f"{exc}", "Progress is saved; run apply again to continue.") from exc
+    except psycopg.Error as exc:
+        detail = str(exc).strip().splitlines()[0] if str(exc).strip() else type(exc).__name__
+        hint = "Every finished batch is saved; run apply again to continue."
+        if "shared memory segment" in detail:
+            hint = (
+                "The parallel index build ran out of shared memory, sized by "
+                "maintenance_work_mem. In Docker, /dev/shm is 64 MB unless the container "
+                "runs with --shm-size (e.g. 2g); or lower maintenance_work_mem. " + hint
+            )
+        raise _fail(f"The database reported an error: {detail}", hint) from exc
+    finally:
+        signal.signal(signal.SIGINT, previous)
+        conn.close()
+
+    if output_json:
+        typer.echo(json.dumps({"event": "result", **_result_dict(result)}))
+    else:
+        _summary(result)
+    if result.status == "failed":
+        raise typer.Exit(1)
+    if result.status in {"stopped", "budget"}:
+        raise typer.Exit(EXIT_STOPPED)
+
+
+async def _run(writer: Any, embedder: Any, state: Any, **kwargs: Any) -> ApplyResult:
+    from vecshift.migrate import apply as run_apply
+
+    try:
+        return await run_apply(writer, embedder, state, **kwargs)
+    finally:
+        await embedder.aclose()
+
+
+def _result_dict(result: ApplyResult) -> dict[str, Any]:
+    from dataclasses import asdict
+
+    return asdict(result)