schift-cli 0.1.0__py3-none-any.whl

schift_cli/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

schift_cli/client.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+from typing import Any
+
+import click
+import httpx
+
+from schift_cli.config import get_api_key, get_api_url
+
+# Timeout: 30s connect, 120s read (migrations can be slow)
+DEFAULT_TIMEOUT = httpx.Timeout(30.0, read=120.0)
+
+
+class SchiftAPIError(Exception):
+    """Raised when the Schift API returns a non-2xx response."""
+
+    def __init__(self, status_code: int, detail: str):
+        self.status_code = status_code
+        self.detail = detail
+        super().__init__(f"HTTP {status_code}: {detail}")
+
+
+class SchiftClient:
+    """HTTP client for the Schift API.
+
+    Handles authentication headers, base URL resolution, and consistent
+    error handling so command modules can stay focused on CLI logic.
+    """
+
+    def __init__(self, api_key: str | None = None, base_url: str | None = None):
+        self.api_key = api_key or get_api_key()
+        self.base_url = (base_url or get_api_url()).rstrip("/")
+        self._http = httpx.Client(
+            base_url=self.base_url,
+            timeout=DEFAULT_TIMEOUT,
+            headers=self._build_headers(),
+        )
+
+    def _build_headers(self) -> dict[str, str]:
+        headers: dict[str, str] = {
+            "User-Agent": "schift-cli/0.1.0",
+            "Accept": "application/json",
+        }
+        if self.api_key:
+            headers["Authorization"] = f"Bearer {self.api_key}"
+        return headers
+
+    # -- HTTP verbs ----------------------------------------------------------
+
+    def get(self, path: str, **kwargs: Any) -> Any:
+        return self._request("GET", path, **kwargs)
+
+    def post(self, path: str, **kwargs: Any) -> Any:
+        return self._request("POST", path, **kwargs)
+
+    def put(self, path: str, **kwargs: Any) -> Any:
+        return self._request("PUT", path, **kwargs)
+
+    def delete(self, path: str, **kwargs: Any) -> Any:
+        return self._request("DELETE", path, **kwargs)
+
+    # -- Internal -------------------------------------------------------------
+
+    def _request(self, method: str, path: str, **kwargs: Any) -> Any:
+        try:
+            resp = self._http.request(method, path, **kwargs)
+        except httpx.ConnectError:
+            raise click.ClickException(
+                f"Could not connect to Schift API at {self.base_url}\n"
+                "  The server may be unavailable. Check your network or set "
+                "SCHIFT_API_URL if using a custom endpoint."
+            )
+        except httpx.TimeoutException:
+            raise click.ClickException(
+                "Request timed out. The server may be under heavy load — try again."
+            )
+
+        if resp.status_code == 401:
+            raise click.ClickException(
+                "Authentication failed. Run `schift auth login` to set your API key."
+            )
+
+        if resp.status_code >= 400:
+            try:
+                body = resp.json()
+                detail = body.get("detail") or body.get("message") or resp.text
+            except Exception:
+                detail = resp.text
+            raise SchiftAPIError(resp.status_code, str(detail))
+
+        if resp.status_code == 204:
+            return None
+        return resp.json()
+
+    def close(self) -> None:
+        self._http.close()
+
+    def __enter__(self) -> SchiftClient:
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        self.close()
+
+
+def require_api_key() -> str:
+    """Return the API key or abort with a helpful message."""
+    key = get_api_key()
+    if not key:
+        raise click.ClickException(
+            "No API key configured.\n"
+            "  Run `schift auth login` or set the SCHIFT_API_KEY environment variable."
+        )
+    return key
+
+
+def get_client() -> SchiftClient:
+    """Create a client, ensuring an API key is present."""
+    api_key = require_api_key()
+    return SchiftClient(api_key=api_key)

schift_cli/commands/auth.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import click
+
+from schift_cli.config import clear_api_key, get_api_key, set_api_key, CONFIG_FILE
+from schift_cli.display import success, info, error
+
+
+@click.group("auth")
+def auth() -> None:
+    """Manage authentication with the Schift platform."""
+
+
+@auth.command()
+def login() -> None:
+    """Set your Schift API key."""
+    existing = get_api_key()
+    if existing:
+        overwrite = click.confirm(
+            "An API key is already configured. Overwrite?", default=False
+        )
+        if not overwrite:
+            info("Keeping existing API key.")
+            return
+
+    api_key = click.prompt("Enter your Schift API key", hide_input=True)
+    api_key = api_key.strip()
+
+    if not api_key:
+        raise click.ClickException("API key cannot be empty.")
+
+    set_api_key(api_key)
+    success(f"API key saved to {CONFIG_FILE}")
+
+
+@auth.command()
+def logout() -> None:
+    """Remove the stored API key."""
+    if not get_api_key():
+        info("No API key is currently stored.")
+        return
+
+    clear_api_key()
+    success("API key removed.")
+
+
+@auth.command()
+def status() -> None:
+    """Show current authentication status."""
+    import os
+    from schift_cli.config import ENV_API_KEY
+
+    env_key = os.environ.get(ENV_API_KEY)
+    file_key = None
+    try:
+        from schift_cli.config import load_config
+        file_key = load_config().get("api_key")
+    except Exception:
+        pass
+
+    if env_key:
+        masked = env_key[:8] + "..." + env_key[-4:] if len(env_key) > 12 else "***"
+        success(f"Authenticated via {ENV_API_KEY} env var (key: {masked})")
+    elif file_key:
+        masked = file_key[:8] + "..." + file_key[-4:] if len(file_key) > 12 else "***"
+        success(f"Authenticated via config file (key: {masked})")
+    else:
+        error("Not authenticated. Run `schift auth login` to get started.")

schift_cli/commands/bench.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import click
+
+from schift_cli.client import get_client, SchiftAPIError
+from schift_cli.display import console, error, info, print_kv, spinner, success
+
+
+@click.command("bench")
+@click.option("--source", "-s", required=True, help="Source model ID (e.g. openai/text-embedding-3-large)")
+@click.option("--target", "-t", required=True, help="Target model ID (e.g. google/gemini-embedding-004)")
+@click.option("--data", "-d", type=click.Path(exists=True, path_type=Path), required=True,
+              help="JSONL file with benchmark queries")
+@click.option("--top-k", "-k", type=int, default=10, show_default=True,
+              help="Number of results to compare per query")
+def bench(source: str, target: str, data: Path, top_k: int) -> None:
+    """Benchmark embedding quality between two models.
+
+    Measures how well a Schift projection preserves retrieval quality
+    when switching from SOURCE to TARGET model.
+    """
+    info(f"Benchmarking projection: {source} -> {target}")
+    info(f"Data: {data}  |  top-k: {top_k}")
+
+    try:
+        with get_client() as client:
+            with spinner("Running benchmark...") as progress:
+                progress.add_task("Running benchmark...", total=None)
+                result = client.post(
+                    "/bench",
+                    json={
+                        "source_model": source,
+                        "target_model": target,
+                        "data_path": str(data),
+                        "top_k": top_k,
+                    },
+                )
+    except SchiftAPIError as e:
+        error(f"Benchmark failed: {e.detail}")
+        raise SystemExit(1)
+    except click.ClickException:
+        raise
+
+    report = result.get("report", result)
+    print_kv("Benchmark Report", {
+        "Source Model": source,
+        "Target Model": target,
+        "Queries": report.get("num_queries", "-"),
+        "Recall@k": report.get("recall_at_k", "-"),
+        "MRR": report.get("mrr", "-"),
+        "Cosine Similarity (avg)": report.get("avg_cosine_similarity", "-"),
+        "Latency (p50)": report.get("latency_p50_ms", "-"),
+        "Latency (p99)": report.get("latency_p99_ms", "-"),
+    })
+
+    quality = report.get("recall_at_k")
+    if quality is not None:
+        if float(quality) >= 0.95:
+            success("Projection quality is excellent.")
+        elif float(quality) >= 0.85:
+            console.print("[yellow]Projection quality is acceptable but may degrade edge cases.[/]")
+        else:
+            console.print("[red]Projection quality is low. Consider increasing sample size in `migrate fit`.[/]")

schift_cli/commands/catalog.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import click
+
+from schift_cli.client import get_client, SchiftAPIError
+from schift_cli.display import print_table, print_kv, error
+
+
+@click.group("catalog")
+def catalog() -> None:
+    """Browse available embedding models."""
+
+
+@catalog.command("list")
+def list_models() -> None:
+    """List all supported embedding models."""
+    try:
+        with get_client() as client:
+            data = client.get("/catalog/models")
+    except SchiftAPIError as e:
+        error(f"Failed to fetch model catalog: {e.detail}")
+        raise SystemExit(1)
+    except click.ClickException:
+        raise
+
+    models = data.get("models", [])
+    if not models:
+        click.echo("No models found in the catalog.")
+        return
+
+    rows = [
+        (
+            m.get("id", ""),
+            m.get("provider", ""),
+            str(m.get("dimensions", "")),
+            m.get("max_tokens", ""),
+            m.get("status", ""),
+        )
+        for m in models
+    ]
+    print_table(
+        "Embedding Model Catalog",
+        ["Model ID", "Provider", "Dimensions", "Max Tokens", "Status"],
+        rows,
+    )
+
+
+@catalog.command("get")
+@click.argument("model_id")
+def get_model(model_id: str) -> None:
+    """Show details for a specific model.
+
+    MODEL_ID is the fully qualified model name, e.g. openai/text-embedding-3-large
+    """
+    try:
+        with get_client() as client:
+            data = client.get(f"/catalog/models/{model_id}")
+    except SchiftAPIError as e:
+        if e.status_code == 404:
+            error(f"Model not found: {model_id}")
+        else:
+            error(f"Failed to fetch model: {e.detail}")
+        raise SystemExit(1)
+    except click.ClickException:
+        raise
+
+    model = data.get("model", data)
+    print_kv(f"Model: {model_id}", {
+        "Provider": model.get("provider", "-"),
+        "Dimensions": model.get("dimensions", "-"),
+        "Max Tokens": model.get("max_tokens", "-"),
+        "Status": model.get("status", "-"),
+        "Description": model.get("description", "-"),
+    })

schift_cli/commands/db.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import click
+
+from schift_cli.client import get_client, SchiftAPIError
+from schift_cli.display import error, info, print_kv, print_table, success
+
+
+@click.group("db")
+def db() -> None:
+    """Manage vector collections."""
+
+
+@db.command("create")
+@click.argument("name")
+@click.option("--dim", "-d", type=int, required=True, help="Vector dimensions (e.g. 3072)")
+@click.option("--metric", type=click.Choice(["cosine", "euclidean", "dot"]),
+              default="cosine", show_default=True, help="Distance metric")
+def create(name: str, dim: int, metric: str) -> None:
+    """Create a new vector collection."""
+    try:
+        with get_client() as client:
+            data = client.post(
+                "/collections",
+                json={"name": name, "dimensions": dim, "metric": metric},
+            )
+    except SchiftAPIError as e:
+        error(f"Failed to create collection: {e.detail}")
+        raise SystemExit(1)
+    except click.ClickException:
+        raise
+
+    collection = data.get("collection", data)
+    success(f"Collection '{name}' created (id: {collection.get('id', '-')})")
+
+
+@db.command("list")
+def list_collections() -> None:
+    """List all vector collections."""
+    try:
+        with get_client() as client:
+            data = client.get("/collections")
+    except SchiftAPIError as e:
+        error(f"Failed to list collections: {e.detail}")
+        raise SystemExit(1)
+    except click.ClickException:
+        raise
+
+    collections = data.get("collections", [])
+    if not collections:
+        info("No collections found. Create one with `schift db create`.")
+        return
+
+    rows = [
+        (
+            c.get("name", ""),
+            str(c.get("dimensions", "")),
+            c.get("metric", ""),
+            str(c.get("vector_count", "")),
+            c.get("created_at", ""),
+        )
+        for c in collections
+    ]
+    print_table(
+        "Vector Collections",
+        ["Name", "Dimensions", "Metric", "Vectors", "Created"],
+        rows,
+    )
+
+
+@db.command("stats")
+@click.argument("name")
+def stats(name: str) -> None:
+    """Show statistics for a collection."""
+    try:
+        with get_client() as client:
+            data = client.get(f"/collections/{name}/stats")
+    except SchiftAPIError as e:
+        if e.status_code == 404:
+            error(f"Collection not found: {name}")
+        else:
+            error(f"Failed to get stats: {e.detail}")
+        raise SystemExit(1)
+    except click.ClickException:
+        raise
+
+    s = data.get("stats", data)
+    print_kv(f"Collection: {name}", {
+        "Vectors": s.get("vector_count", "-"),
+        "Dimensions": s.get("dimensions", "-"),
+        "Metric": s.get("metric", "-"),
+        "Index Type": s.get("index_type", "-"),
+        "Storage Size": s.get("storage_size", "-"),
+        "Created": s.get("created_at", "-"),
+        "Last Updated": s.get("updated_at", "-"),
+    })

schift_cli/commands/embed.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import click
+
+from schift_cli.client import get_client, SchiftAPIError
+from schift_cli.display import console, error, info, success
+
+
+@click.group("embed", invoke_without_command=True)
+@click.argument("text", required=False, default=None)
+@click.option("--model", "-m", default=None,
+              help="Embedding model ID (e.g. openai/text-embedding-3-large)")
+@click.pass_context
+def embed(ctx: click.Context, text: str | None, model: str | None) -> None:
+    """Generate embeddings for text.
+
+    \b
+    Usage:
+        schift embed "hello world" --model openai/text-embedding-3-large
+        schift embed batch --file texts.jsonl --model google/gemini-embedding-004
+    """
+    if ctx.invoked_subcommand is not None:
+        return
+
+    if not text:
+        raise click.UsageError("Provide TEXT to embed, or use `schift embed batch`.")
+    if not model:
+        raise click.UsageError("--model is required.")
+
+    try:
+        with get_client() as client:
+            data = client.post("/embed", json={"text": text, "model": model})
+    except SchiftAPIError as e:
+        error(f"Embedding failed: {e.detail}")
+        raise SystemExit(1)
+    except click.ClickException:
+        raise
+
+    embedding = data.get("embedding", [])
+    dims = len(embedding)
+    preview = embedding[:5]
+    preview_str = ", ".join(f"{v:.6f}" for v in preview)
+
+    success(f"Generated {dims}-dimensional embedding")
+    console.print(f"  [{preview_str}, ...]")
+
+
+@embed.command("batch")
+@click.option("--file", "-f", "file_path", required=True,
+              type=click.Path(exists=True, path_type=Path),
+              help="JSONL file with one text per line (field: \"text\")")
+@click.option("--model", "-m", required=True,
+              help="Embedding model ID (e.g. google/gemini-embedding-004)")
+@click.option("--output", "-o", type=click.Path(path_type=Path), default=None,
+              help="Output JSONL file for results (default: stdout summary)")
+def embed_batch(file_path: Path, model: str, output: Path | None) -> None:
+    """Embed multiple texts from a JSONL file.
+
+    Each line in the input file must be a JSON object with a "text" field.
+    """
+    texts: list[str] = []
+    with open(file_path) as f:
+        for i, line in enumerate(f, 1):
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                obj = json.loads(line)
+                texts.append(obj["text"])
+            except (json.JSONDecodeError, KeyError):
+                error(f"Line {i}: expected JSON with a \"text\" field")
+                raise SystemExit(1)
+
+    if not texts:
+        error("No texts found in input file.")
+        raise SystemExit(1)
+
+    info(f"Embedding {len(texts)} texts with model {model}")
+
+    try:
+        with get_client() as client:
+            data = client.post(
+                "/embed/batch",
+                json={"texts": texts, "model": model},
+            )
+    except SchiftAPIError as e:
+        error(f"Batch embedding failed: {e.detail}")
+        raise SystemExit(1)
+    except click.ClickException:
+        raise
+
+    embeddings = data.get("embeddings", [])
+
+    if output:
+        with open(output, "w") as f:
+            for text_val, emb in zip(texts, embeddings):
+                f.write(json.dumps({"text": text_val, "embedding": emb}) + "\n")
+        success(f"Wrote {len(embeddings)} embeddings to {output}")
+    else:
+        dims = len(embeddings[0]) if embeddings else 0
+        success(f"Generated {len(embeddings)} embeddings ({dims} dimensions each)")

schift_cli/commands/migrate.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import click
+
+from schift_cli.client import get_client, SchiftAPIError
+from schift_cli.display import console, error, info, print_kv, spinner, success
+
+
+@click.group("migrate")
+def migrate() -> None:
+    """Fit projection matrices and migrate vector databases."""
+
+
+@migrate.command("fit")
+@click.option("--source", "-s", required=True, help="Source embedding model ID")
+@click.option("--target", "-t", required=True, help="Target embedding model ID")
+@click.option("--sample", type=float, default=0.1, show_default=True,
+              help="Fraction of data to sample for fitting (0.0-1.0)")
+def fit(source: str, target: str, sample: float) -> None:
+    """Fit a projection matrix between two embedding models.
+
+    The projection is computed server-side. You never see the matrix --
+    Schift stores it and returns a projection ID for use in `migrate run`.
+    """
+    if not 0.0 < sample <= 1.0:
+        raise click.BadParameter("Sample must be between 0 and 1.", param_hint="--sample")
+
+    info(f"Fitting projection: {source} -> {target} (sample={sample})")
+
+    try:
+        with get_client() as client:
+            with spinner("Fitting projection matrix...") as progress:
+                progress.add_task("Fitting projection matrix...", total=None)
+                result = client.post(
+                    "/migrate/fit",
+                    json={
+                        "source_model": source,
+                        "target_model": target,
+                        "sample_fraction": sample,
+                    },
+                )
+    except SchiftAPIError as e:
+        error(f"Fit failed: {e.detail}")
+        raise SystemExit(1)
+    except click.ClickException:
+        raise
+
+    projection = result.get("projection", result)
+    proj_id = projection.get("id", "unknown")
+
+    print_kv("Projection Created", {
+        "Projection ID": proj_id,
+        "Source Model": source,
+        "Target Model": target,
+        "Sample Fraction": sample,
+        "Status": projection.get("status", "ready"),
+        "Quality (R2)": projection.get("r2_score", "-"),
+    })
+
+    success(f"Use this projection ID to migrate: schift migrate run --projection {proj_id}")
+
+
+@migrate.command("run")
+@click.option("--projection", "-p", required=True, help="Projection ID from `migrate fit`")
+@click.option("--db", required=True, help="Database connection string (e.g. pgvector://...)")
+@click.option("--dry-run", is_flag=True, default=False, help="Preview the migration without applying changes")
+@click.option("--batch-size", type=int, default=1000, show_default=True,
+              help="Number of vectors to migrate per batch")
+def run(projection: str, db: str, dry_run: bool, batch_size: int) -> None:
+    """Apply a projection to migrate vectors in a database.
+
+    Use --dry-run first to preview the migration plan.
+    """
+    mode = "DRY RUN" if dry_run else "LIVE"
+    info(f"Migration [{mode}]: projection={projection}")
+    info(f"Database: {_mask_connection_string(db)}")
+
+    if not dry_run:
+        click.confirm(
+            "This will modify vectors in your database. Continue?",
+            abort=True,
+        )
+
+    try:
+        with get_client() as client:
+            with spinner("Running migration...") as progress:
+                progress.add_task("Running migration...", total=None)
+                result = client.post(
+                    "/migrate/run",
+                    json={
+                        "projection_id": projection,
+                        "db_connection": db,
+                        "dry_run": dry_run,
+                        "batch_size": batch_size,
+                    },
+                )
+    except SchiftAPIError as e:
+        error(f"Migration failed: {e.detail}")
+        raise SystemExit(1)
+    except click.ClickException:
+        raise
+
+    migration = result.get("migration", result)
+
+    print_kv(f"Migration Result ({mode})", {
+        "Vectors Processed": migration.get("vectors_processed", "-"),
+        "Vectors Skipped": migration.get("vectors_skipped", "-"),
+        "Duration": migration.get("duration", "-"),
+        "Status": migration.get("status", "-"),
+    })
+
+    if dry_run:
+        info("No changes were applied. Remove --dry-run to execute.")
+    else:
+        success("Migration complete.")
+
+
+def _mask_connection_string(conn: str) -> str:
+    """Hide password in connection strings for display."""
+    if "@" in conn:
+        # pgvector://user:pass@host -> pgvector://user:***@host
+        before_at = conn.split("@")[0]
+        after_at = conn.split("@", 1)[1]
+        if ":" in before_at:
+            scheme_user = before_at.rsplit(":", 1)[0]
+            return f"{scheme_user}:***@{after_at}"
+    return conn