kiteml-cli 0.1.0__py3-none-any.whl

kite_cli/__init__.py

@@ -0,0 +1,3 @@
+"""Kite CLI and MCP server for the Kite robotics training platform."""
+
+__version__ = "0.1.0"

kite_cli/auth.py

@@ -0,0 +1,184 @@
+"""
+Authentication for the Kite CLI.
+
+Supports:
+  - API key: KITE_API_KEY env var (recommended for MCP / programmatic access)
+  - Browser OAuth: `kite auth login` opens the dashboard in a browser → an API key is stored
+  - Environment variable: KITE_TOKEN for CI/headless
+  - Token file: ~/.kite/credentials.json (written by login, read by client)
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import secrets
+import sys
+import threading
+import time
+import webbrowser
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from pathlib import Path
+from typing import Optional
+from urllib.parse import parse_qs, urlparse
+
+from kite_cli.config import KITE_DIR, get_app_url
+
+CREDENTIALS_PATH = KITE_DIR / "credentials.json"
+
+
+def get_token() -> Optional[str]:
+    """Get a valid auth token. Checks env vars first, then credentials file."""
+    # 1. API key (MCP / programmatic access)
+    api_key = os.environ.get("KITE_API_KEY")
+    if api_key:
+        return api_key
+
+    # 2. Token override (CI, headless)
+    env_token = os.environ.get("KITE_TOKEN")
+    if env_token:
+        return env_token
+
+    # 3. Stored credentials
+    creds = _load_credentials()
+    if creds and creds.get("id_token"):
+        return creds["id_token"]
+
+    return None
+
+
+def save_credentials(id_token: str, user_info: Optional[dict] = None) -> None:
+    """Save credentials to disk with restricted permissions."""
+    KITE_DIR.mkdir(parents=True, exist_ok=True)
+    creds = {
+        "id_token": id_token,
+        "user_info": user_info,
+        "saved_at": time.time(),
+    }
+    CREDENTIALS_PATH.write_text(json.dumps(creds, indent=2) + "\n")
+    # Restrict file permissions (owner read/write only)
+    try:
+        CREDENTIALS_PATH.chmod(0o600)
+    except OSError:
+        pass
+
+
+def clear_credentials() -> None:
+    """Remove stored credentials."""
+    if CREDENTIALS_PATH.exists():
+        CREDENTIALS_PATH.unlink()
+
+
+def _load_credentials() -> Optional[dict]:
+    """Load credentials from disk."""
+    if not CREDENTIALS_PATH.exists():
+        return None
+    try:
+        return json.loads(CREDENTIALS_PATH.read_text())
+    except (json.JSONDecodeError, OSError):
+        return None
+
+
+class _CallbackHandler(BaseHTTPRequestHandler):
+    """HTTP handler that receives the OAuth callback with the token."""
+
+    token: Optional[str] = None
+    user_info: Optional[dict] = None
+    expected_state: Optional[str] = None
+
+    def do_GET(self):
+        parsed = urlparse(self.path)
+        params = parse_qs(parsed.query)
+
+        if parsed.path != "/callback":
+            self.send_response(404)
+            self.end_headers()
+            return
+
+        state = params.get("state", [None])[0]
+        token = params.get("token", [None])[0]
+        email = params.get("email", [None])[0]
+
+        # Only accept the login this process started. Without this, any local
+        # page could hit the ephemeral port and plant a token. compare_digest
+        # keeps the nonce from being probed a byte at a time.
+        expected = _CallbackHandler.expected_state
+        if not expected or not state or not secrets.compare_digest(state, expected):
+            self._reply(403, "Login failed", "This callback did not come from the login you started.")
+            return
+
+        if not token:
+            self._reply(400, "Login failed", "No token received.")
+            return
+
+        _CallbackHandler.token = token
+        _CallbackHandler.user_info = {"email": email} if email else None
+        self._reply(200, "Login successful!", "You can close this tab.")
+
+    def _reply(self, status: int, heading: str, detail: str) -> None:
+        self.send_response(status)
+        self.send_header("Content-Type", "text/html")
+        self.end_headers()
+        self.wfile.write(f"<html><body><h2>{heading}</h2><p>{detail}</p></body></html>".encode())
+
+    def log_message(self, format, *args):
+        pass  # Suppress HTTP server logs
+
+
+def login_browser(timeout: int = 120) -> bool:
+    """Open browser for OAuth login. Returns True on success."""
+    import socket
+
+    # Find a free port
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        port = s.getsockname()[1]
+
+    # Single-use nonce the /cli-auth page echoes back, so we can tell our own
+    # callback apart from any other request that reaches this port.
+    state = secrets.token_urlsafe(24)
+
+    # Reset handler state
+    _CallbackHandler.token = None
+    _CallbackHandler.user_info = None
+    _CallbackHandler.expected_state = state
+
+    server = HTTPServer(("127.0.0.1", port), _CallbackHandler)
+    # Poll in short slices so a stray request can't hold us past the deadline.
+    server.timeout = 1
+
+    # The dashboard's /cli-auth page signs you in, then redirects to our
+    # loopback callback with the token. Derive nothing from the API URL: the
+    # app is a different host (app.kiteml.com vs api.kiteml.com).
+    login_url = f"{get_app_url().rstrip('/')}/cli-auth?redirect_port={port}&state={state}"
+
+    print(f"Opening browser for login...")
+    print(f"If the browser doesn't open, visit: {login_url}")
+    webbrowser.open(login_url)
+
+    # Wait for callback
+    deadline = time.time() + timeout
+    while time.time() < deadline and _CallbackHandler.token is None:
+        server.handle_request()
+
+    server.server_close()
+
+    if _CallbackHandler.token:
+        save_credentials(
+            id_token=_CallbackHandler.token,
+            user_info=_CallbackHandler.user_info,
+        )
+        return True
+
+    return False
+
+
+def login_headless() -> bool:
+    """Headless login: prompt user to paste a token."""
+    print("Visit your Kite dashboard to generate an API token.")
+    print("Then paste it here:")
+    token = input("> ").strip()
+    if token:
+        save_credentials(id_token=token)
+        return True
+    return False

kite_cli/cli/commands/augment.py

@@ -0,0 +1,164 @@
+"""`kite augment` — dataset augmentation jobs via the public /v1 API.
+
+create / status / download / cancel / list. Authenticates with KITE_API_KEY like the
+rest of the CLI; calls the bare `/v1/augmentations` surface.
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+from pathlib import Path
+from typing import List, Optional
+
+import typer
+
+from kite_cli.client import KiteClient, KiteAPIError
+from kite_cli.output import print_error, print_json
+from kite_cli.progress import wait_for_completion
+
+app = typer.Typer(no_args_is_help=True, help="Dataset augmentation jobs")
+
+_TERMINAL = {"succeeded", "failed", "canceled"}
+_MEDIA_TYPES = {".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".png": "image/png", ".webp": "image/webp"}
+
+
+def reference_image(path: str, camera: Optional[str] = None) -> dict:
+    """A local photo of the target look → a /v1 ``reference_images`` entry (base64 inline)."""
+    p = Path(path).expanduser()
+    media_type = _MEDIA_TYPES.get(p.suffix.lower())
+    if media_type is None:
+        raise ValueError(f"{p}: reference images must be .jpg, .png or .webp")
+    entry = {"data": base64.b64encode(p.read_bytes()).decode("ascii"), "media_type": media_type}
+    return {"camera": camera, **entry} if camera else entry
+
+
+@app.command()
+def create(
+    repo_id: str = typer.Option(..., "--repo-id", help="Source HF dataset id, e.g. lerobot/pusht"),
+    instructions: str = typer.Option(..., "-i", "--instructions", help="Augmentation instruction"),
+    episodes: int = typer.Option(3, "-n", "--episodes", help="Episodes to generate with video augmentation (1-50); relight covers every episode"),
+    model: Optional[str] = typer.Option(
+        None, "--model", help="auto (default: pick relight or video) | relight (colour/exposure, every episode) | video (generative)"
+    ),
+    reference_images: Optional[List[str]] = typer.Option(
+        None, "--reference-image", help="Photo of the target look, [camera=]PATH (repeatable)"
+    ),
+    relight: Optional[str] = typer.Option(
+        None, "--relight", help='JSON relight params per camera, e.g. \'{"*": {"exposure": -1.2, "temperature": 3200}}\''
+    ),
+    output: str = typer.Option("download", "--output", help="'download' or 'huggingface'"),
+    hf_repo: Optional[str] = typer.Option(None, "--hf-repo", help="Target HF repo (for --output huggingface)"),
+    private: bool = typer.Option(True, "--private/--public", help="HF output visibility"),
+    variation_prompts: Optional[str] = typer.Option(
+        None, "--variation-prompts", help="JSON array of per-episode prompts (overrides --episodes)"
+    ),
+    wait: bool = typer.Option(False, "--wait", help="Poll until the run finishes"),
+):
+    """Create an augmentation run."""
+    body: dict = {
+        "source": {"type": "huggingface", "repo_id": repo_id},
+        "instructions": instructions,
+        "config": {"episode_count": episodes},
+        "output": {"type": output, "private": private},
+    }
+    if hf_repo:
+        body["output"]["repo_id"] = hf_repo
+    if model:
+        body["model"] = model
+    try:
+        if relight:
+            body["config"]["relight"] = json.loads(relight)
+        if reference_images:
+            body["reference_images"] = [
+                reference_image(path, camera or None)
+                for camera, _, path in (
+                    spec.partition("=") if "=" in spec and not Path(spec).exists() else ("", "", spec)
+                    for spec in reference_images
+                )
+            ]
+    except (json.JSONDecodeError, OSError, ValueError) as e:
+        print_error(f"invalid --relight / --reference-image: {e}")
+        return
+    if variation_prompts:
+        try:
+            body["config"]["variation_prompts"] = json.loads(variation_prompts)
+        except json.JSONDecodeError:
+            print_error("--variation-prompts must be a JSON array of strings")
+            return
+    try:
+        client = KiteClient()
+        result = client.post("/v1/augmentations", json_data=body)
+        if wait and result.get("id"):
+            result = wait_for_completion(
+                client, f"/v1/augmentations/{result['id']}", terminal_states=_TERMINAL
+            )
+        print_json(result)
+    except KiteAPIError as e:
+        print_error(e.message, code=e.code, details=e.details)
+
+
+@app.command()
+def status(
+    aug_id: str = typer.Argument(..., help="Augmentation id (aug_...)"),
+    wait: bool = typer.Option(False, "--wait", help="Poll until the run finishes"),
+):
+    """Get an augmentation's status."""
+    try:
+        client = KiteClient()
+        path = f"/v1/augmentations/{aug_id}"
+        result = wait_for_completion(client, path, terminal_states=_TERMINAL) if wait else client.get(path)
+        print_json(result)
+    except KiteAPIError as e:
+        print_error(e.message, code=e.code, details=e.details)
+
+
+@app.command("list")
+def list_augmentations(limit: int = typer.Option(20, "--limit", help="Max rows (1-100)")):
+    """List your augmentation runs (newest first)."""
+    try:
+        client = KiteClient()
+        result = client.get("/v1/augmentations", params={"limit": limit})
+        print_json(result)
+    except KiteAPIError as e:
+        print_error(e.message, code=e.code, details=e.details)
+
+
+@app.command()
+def cancel(aug_id: str = typer.Argument(..., help="Augmentation id (aug_...)")):
+    """Cancel a running augmentation."""
+    try:
+        client = KiteClient()
+        result = client.post(f"/v1/augmentations/{aug_id}/cancel")
+        print_json(result)
+    except KiteAPIError as e:
+        print_error(e.message, code=e.code, details=e.details)
+
+
+@app.command()
+def download(
+    aug_id: str = typer.Argument(..., help="Augmentation id (aug_...)"),
+    out: Path = typer.Option(Path("."), "-o", "--out", help="Output directory"),
+):
+    """Download a finished download-output dataset, preserving its file layout."""
+    try:
+        client = KiteClient()
+        result = client.get(f"/v1/augmentations/{aug_id}")
+        if result.get("status") != "succeeded":
+            print_error(f"Augmentation is '{result.get('status')}', not succeeded yet.")
+            return
+        files = (result.get("output") or {}).get("files") or []
+        if not files:
+            print_error("No downloadable files (is this a 'download'-output run?).")
+            return
+        written = []
+        for f in files:
+            # Always fetch via the authenticated proxy route (never send the Bearer token
+            # to a GCS signed URL); the layout is reconstructed from each file's path.
+            dest = out / f["path"]
+            dest.parent.mkdir(parents=True, exist_ok=True)
+            client.download(f"/v1/augmentations/{aug_id}/files/{f['path']}", str(dest))
+            written.append(str(dest))
+        print_json({"downloaded": written, "dir": str(out)})
+    except KiteAPIError as e:
+        print_error(e.message, code=e.code, details=e.details)

kite_cli/cli/commands/auth_cmd.py

@@ -0,0 +1,57 @@
+"""kite auth — login, logout, whoami, token."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import typer
+
+from kite_cli.output import print_json, print_error
+
+app = typer.Typer(no_args_is_help=True)
+
+
+@app.command()
+def login(headless: bool = typer.Option(False, help="Paste token instead of browser flow")):
+    """Log in to Kite via browser OAuth or token paste."""
+    from kite_cli.auth import login_browser, login_headless
+
+    if headless:
+        ok = login_headless()
+    else:
+        ok = login_browser()
+
+    if ok:
+        print_json({"logged_in": True})
+    else:
+        print_error("Login failed. Try --headless to paste a token manually.")
+
+
+@app.command()
+def logout():
+    """Clear stored credentials."""
+    from kite_cli.auth import clear_credentials
+    clear_credentials()
+    print_json({"logged_out": True})
+
+
+@app.command()
+def whoami():
+    """Show the key or session you are signed in with."""
+    from kite_cli.client import KiteClient, KiteAPIError
+    try:
+        client = KiteClient()
+        print_json(client.get("/v1/keys/me"))
+    except KiteAPIError as e:
+        print_error(e.message, code=e.code)
+
+
+@app.command()
+def token():
+    """Print the current auth token (for piping)."""
+    from kite_cli.auth import get_token
+    t = get_token()
+    if t:
+        print(t)
+    else:
+        print_error("No token found. Run: kite auth login")

kite_cli/cli/commands/config_cmd.py

@@ -0,0 +1,39 @@
+"""kite config — get, set, show, set-project."""
+
+from __future__ import annotations
+
+import typer
+
+from kite_cli.config import get_value, set_value, load_config
+from kite_cli.output import print_json
+
+app = typer.Typer(no_args_is_help=True)
+
+
+@app.command()
+def show():
+    """Show all configuration values."""
+    print_json(load_config())
+
+
+@app.command()
+def get(key: str = typer.Argument(..., help="Config key to read")):
+    """Get a configuration value."""
+    print_json({"key": key, "value": get_value(key)})
+
+
+@app.command("set")
+def set_cmd(
+    key: str = typer.Argument(..., help="Config key"),
+    value: str = typer.Argument(..., help="Config value"),
+):
+    """Set a configuration value."""
+    set_value(key, value)
+    print_json({"key": key, "value": value, "saved": True})
+
+
+@app.command("set-project")
+def set_project(project_id: str = typer.Argument(..., help="Default project ID")):
+    """Set the default project ID."""
+    set_value("default_project_id", project_id)
+    print_json({"default_project_id": project_id, "saved": True})