embedkombinat-annotator 0.3.0__py3-none-any.whl

annotator/__init__.py

@@ -0,0 +1,6 @@
+"""Distributed annotation worker for embedkombinat."""
+
+__version__ = "0.3.0"
+
+TEAL = "#00E5B0"
+AMBER = "#c05d3b"

annotator/auth.py

@@ -0,0 +1,236 @@
+"""GitHub OAuth Device Flow and token management.
+
+Why device flow: the previous web-flow design ran a localhost callback server
+inside the CLI process, then handed GitHub a `redirect_uri=http://localhost:PORT/callback`
+URI. That works on a single machine where browser and server are colocated
+(your laptop), but breaks the moment the CLI runs on a remote host the user
+SSH'd into (Runpod, Lambda, EC2, etc.) — the user's browser hits localhost
+on *their machine*, not the remote one. The device flow has no callback at
+all: the CLI prints a short user code, the user enters it in any browser on
+any device, the CLI polls GitHub directly. Same auth, no machine-boundary
+assumption.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import os
+import stat
+import time
+import webbrowser
+from datetime import UTC, datetime, timedelta
+from typing import TYPE_CHECKING, Any
+
+import httpx
+from pydantic import BaseModel
+
+from annotator import TEAL
+from annotator.errors import AuthError
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from rich.console import Console
+
+    from annotator.config import Settings
+
+AUTH_FILE = "auth.json"
+DEVICE_CODE_URL = "https://github.com/login/device/code"
+TOKEN_URL = "https://github.com/login/oauth/access_token"
+DEFAULT_SCOPE = "read:user"
+EXPIRY_BUFFER = timedelta(minutes=5)
+MIN_POLL_INTERVAL = 5.0
+
+
+class ContributorInfo(BaseModel):
+    id: str
+    github_username: str
+    github_avatar_url: str | None = None
+
+
+class AuthToken(BaseModel):
+    kombinat_url: str
+    access_token: str
+    expires_at: datetime
+    contributor: ContributorInfo
+
+    def is_expired(self) -> bool:
+        return datetime.now(tz=UTC) >= (self.expires_at - EXPIRY_BUFFER)
+
+
+def save_token(token: AuthToken, home: Path) -> None:
+    """Save auth token to disk with restricted permissions."""
+    home.mkdir(parents=True, exist_ok=True)
+    path = home / AUTH_FILE
+    path.write_text(token.model_dump_json(indent=2))
+    os.chmod(path, stat.S_IRUSR | stat.S_IWUSR)
+
+
+def load_token(home: Path) -> AuthToken | None:
+    """Load auth token from disk. Returns None if missing or expired."""
+    path = home / AUTH_FILE
+    if not path.exists():
+        return None
+    try:
+        token = AuthToken.model_validate_json(path.read_text())
+    except (json.JSONDecodeError, ValueError):
+        return None
+    if token.is_expired():
+        return None
+    return token
+
+
+def delete_token(home: Path) -> None:
+    """Remove auth token from disk."""
+    path = home / AUTH_FILE
+    if path.exists():
+        path.unlink()
+
+
+def fetch_client_id(kombinat_url: str) -> str:
+    """Fetch the public GitHub OAuth client_id from kombinat."""
+    try:
+        with httpx.Client() as client:
+            resp = client.get(f"{kombinat_url}/v1/auth/config", timeout=10.0)
+    except httpx.HTTPError as exc:
+        raise AuthError(f"could not reach kombinat at {kombinat_url}: {exc}") from exc
+    if resp.status_code != 200:
+        raise AuthError(f"kombinat auth config fetch failed: {resp.status_code} {resp.text}")
+    client_id = resp.json().get("client_id")
+    if not isinstance(client_id, str) or not client_id:
+        raise AuthError("kombinat returned empty client_id — server is misconfigured")
+    return client_id
+
+
+def request_device_code(client_id: str) -> dict[str, Any]:
+    """Ask GitHub for a device code + user code. Returns the full response payload."""
+    try:
+        with httpx.Client() as client:
+            resp = client.post(
+                DEVICE_CODE_URL,
+                data={"client_id": client_id, "scope": DEFAULT_SCOPE},
+                headers={"Accept": "application/json"},
+                timeout=10.0,
+            )
+    except httpx.HTTPError as exc:
+        raise AuthError(f"could not reach GitHub: {exc}") from exc
+    if resp.status_code != 200:
+        raise AuthError(f"device code request failed: {resp.status_code} {resp.text}")
+    data: dict[str, Any] = resp.json()
+    required = {"device_code", "user_code", "verification_uri", "expires_in", "interval"}
+    if not required.issubset(data):
+        raise AuthError(
+            "device code response is missing required fields. "
+            "This usually means the OAuth app does not have Device Flow enabled — "
+            "check 'Enable Device Flow' in the OAuth app settings on GitHub."
+        )
+    return data
+
+
+def poll_for_access_token(
+    client_id: str,
+    device_code: str,
+    interval: float,
+    expires_in: float,
+) -> str:
+    """Poll GitHub's token endpoint until the user authorizes. Returns the access token."""
+    deadline = time.monotonic() + expires_in
+    poll_interval = max(interval, MIN_POLL_INTERVAL)
+
+    while time.monotonic() < deadline:
+        time.sleep(poll_interval)
+        try:
+            with httpx.Client() as client:
+                resp = client.post(
+                    TOKEN_URL,
+                    data={
+                        "client_id": client_id,
+                        "device_code": device_code,
+                        "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+                    },
+                    headers={"Accept": "application/json"},
+                    timeout=10.0,
+                )
+        except httpx.HTTPError as exc:
+            raise AuthError(f"polling GitHub failed: {exc}") from exc
+
+        data = resp.json()
+        if "access_token" in data:
+            return str(data["access_token"])
+
+        error = data.get("error")
+        if error == "authorization_pending":
+            continue
+        if error == "slow_down":
+            poll_interval += 5.0
+            continue
+        if error == "expired_token":
+            raise AuthError("device code expired before authorization completed")
+        if error == "access_denied":
+            raise AuthError("authorization was denied")
+        raise AuthError(f"unexpected error during device flow: {data}")
+
+    raise AuthError(f"device flow timed out after {expires_in:.0f}s")
+
+
+def exchange_github_token(github_access_token: str, kombinat_url: str) -> AuthToken:
+    """Exchange a GitHub access token for a kombinat JWT."""
+    try:
+        with httpx.Client() as client:
+            resp = client.post(
+                f"{kombinat_url}/v1/auth/github-device",
+                json={"access_token": github_access_token},
+                timeout=30.0,
+            )
+    except httpx.HTTPError as exc:
+        raise AuthError(f"could not reach kombinat at {kombinat_url}: {exc}") from exc
+    if resp.status_code == 401:
+        raise AuthError("kombinat rejected the GitHub access token")
+    if resp.status_code != 200:
+        raise AuthError(f"kombinat auth failed: {resp.status_code} {resp.text}")
+    data = resp.json()
+    expires_at = datetime.now(tz=UTC) + timedelta(seconds=data["expires_in"])
+    contributor_data = data["contributor"]
+    return AuthToken(
+        kombinat_url=kombinat_url,
+        access_token=data["access_token"],
+        expires_at=expires_at,
+        contributor=ContributorInfo(
+            id=contributor_data["id"],
+            github_username=contributor_data["github_username"],
+            github_avatar_url=contributor_data.get("github_avatar_url"),
+        ),
+    )
+
+
+def login(settings: Settings, console: Console) -> AuthToken:
+    """Run the GitHub Device Flow and exchange the access token for a kombinat JWT."""
+    console.print("  No credentials found. Starting login...\n")
+
+    client_id = fetch_client_id(settings.kombinat_url)
+    device_data = request_device_code(client_id)
+
+    user_code = device_data["user_code"]
+    verification_uri = device_data["verification_uri"]
+    device_code = device_data["device_code"]
+    interval = float(device_data["interval"])
+    expires_in = float(device_data["expires_in"])
+
+    console.print(f"  -> Open in any browser: [bold]{verification_uri}[/bold]")
+    console.print(f"  -> Enter code: [bold {TEAL}]{user_code}[/bold {TEAL}]\n")
+
+    # Best-effort browser open as a convenience on machines that have one;
+    # silently a no-op on headless hosts (Runpod, etc.) which is the whole
+    # point of the device flow.
+    with contextlib.suppress(Exception):
+        webbrowser.open(verification_uri)
+
+    console.print("  -> Waiting for authorization...\n")
+
+    github_token = poll_for_access_token(client_id, device_code, interval, expires_in)
+    token = exchange_github_token(github_token, settings.kombinat_url)
+    save_token(token, settings.annotator_home)
+
+    console.print(f"  [{TEAL}]✓[/{TEAL}] Authenticated as {token.contributor.github_username}")
+    return token

annotator/cli.py

@@ -0,0 +1,134 @@
+"""CLI interface for annotator."""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+import typer
+from rich.console import Console
+from rich.panel import Panel
+from rich.text import Text
+
+from annotator import AMBER, TEAL, __version__
+from annotator.config import ExitCode, Settings
+
+app = typer.Typer(
+    name="annotator",
+    no_args_is_help=False,
+    invoke_without_command=True,
+    add_completion=False,
+)
+
+
+def _print_banner(console: Console) -> None:
+    """Print the EmbedKombinat branded header."""
+    logo = Text()
+    logo.append(" EEEEE  K   K\n", style=f"bold {TEAL}")
+    logo.append(" E      K  K\n", style=f"bold {TEAL}")
+    logo.append(" EEEE   KKK\n", style=f"bold {TEAL}")
+    logo.append(" E      K  K\n", style=f"bold {TEAL}")
+    logo.append(" EEEEE  K   K\n", style=f"bold {TEAL}")
+
+    title = f"embed kombinat \u00b7 annotator v{__version__}"
+    console.print()
+    console.print(logo, end="")
+    console.print(Panel(title, style=f"bold {TEAL}", width=len(title) + 6))
+    console.print()
+
+
+@app.callback(invoke_without_command=True)
+def main(
+    ctx: typer.Context,
+    batch_size: Annotated[
+        int | None, typer.Option("--batch-size", help="Pairs per batch claimed from kombinat")
+    ] = None,
+    model: Annotated[
+        str | None, typer.Option("--model", help="Override auto-selected model (HuggingFace ID)")
+    ] = None,
+    quantization: Annotated[
+        str | None, typer.Option("--quantization", help="Override quantization (awq, fp16, etc.)")
+    ] = None,
+    backend: Annotated[
+        str | None, typer.Option("--backend", help="Override backend (vllm, mlx, llama_cpp)")
+    ] = None,
+    gpu_memory_utilization: Annotated[
+        float, typer.Option("--gpu-memory-utilization", help="GPU memory fraction (vLLM only)")
+    ] = 0.9,
+    dry_run: Annotated[
+        bool, typer.Option("--dry-run", help="Process one pair without submitting")
+    ] = False,
+) -> None:
+    """Start the labeling loop. Default command."""
+    if ctx.invoked_subcommand is not None:
+        return
+
+    console = Console()
+    _print_banner(console)
+
+    from annotator.runner import AnnotatorRunner
+
+    settings = Settings()
+    runner = AnnotatorRunner(settings, console)
+    exit_code = runner.run(
+        batch_size=batch_size if batch_size is not None else settings.batch_size,
+        model_override=model,
+        quantization_override=quantization,
+        backend_override=backend,
+        gpu_memory_utilization=gpu_memory_utilization,
+        dry_run=dry_run,
+    )
+    raise typer.Exit(code=exit_code)
+
+
+@app.command()
+def login() -> None:
+    """Authenticate with GitHub."""
+    from annotator import auth
+
+    console = Console()
+    settings = Settings()
+    try:
+        token = auth.login(settings, console)
+        console.print(
+            f"  [{TEAL}]\u2713[/{TEAL}] Authenticated as {token.contributor.github_username}"
+        )
+    except Exception as e:
+        console.print(f"  [{AMBER}]\u2717[/{AMBER}] Login failed: {e}")
+        raise typer.Exit(code=ExitCode.AUTH_FAILURE) from e
+
+
+@app.command()
+def status() -> None:
+    """Show contributor profile and stats."""
+    from annotator import auth
+    from annotator.client import KombinatClient
+
+    console = Console()
+    settings = Settings()
+    token = auth.load_token(settings.annotator_home)
+    if token is None:
+        console.print(f"  [{AMBER}]Not logged in.[/{AMBER}] Run 'annotator login'.")
+        raise typer.Exit(code=ExitCode.AUTH_FAILURE)
+
+    client = KombinatClient(token.kombinat_url, token.access_token)
+    try:
+        profile = client.get_profile()
+        console.print(f"  [{TEAL}]\u2713[/{TEAL}] Logged in as {profile.github_username}")
+        console.print(f"    Total annotations: {profile.total_annotations}")
+        console.print(f"    Reputation score:  {profile.reputation_score:.2f}")
+    except Exception as e:
+        console.print(f"  [{AMBER}]\u2717[/{AMBER}] Failed to fetch status: {e}")
+        raise typer.Exit(code=ExitCode.KOMBINAT_UNREACHABLE) from e
+    finally:
+        client.close()
+
+
+@app.command()
+def logout() -> None:
+    """Remove stored credentials."""
+    from annotator import auth
+
+    console = Console()
+    settings = Settings()
+    auth.delete_token(settings.annotator_home)
+    console.print(f"  [{TEAL}]\u2713[/{TEAL}] Logged out.")

annotator/client.py

@@ -0,0 +1,174 @@
+"""HTTP client for kombinat API."""
+
+from __future__ import annotations
+
+import logging
+import random
+import time
+from datetime import datetime  # noqa: TC003 - needed at runtime by Pydantic
+
+import httpx
+from pydantic import BaseModel
+
+from annotator.errors import AuthError, KombinatError
+
+logger = logging.getLogger(__name__)
+
+MAX_RETRIES = 3
+MIN_RETRY_DELAY = 1.0
+MAX_RETRY_DELAY = 30.0
+NO_PAIRS_INITIAL_WAIT = 30.0
+NO_PAIRS_MAX_WAIT = 600.0
+
+
+class PairData(BaseModel):
+    pair_id: str
+    query_text: str
+    doc_text: str
+
+
+class BatchResponse(BaseModel):
+    batch_id: str
+    expires_at: datetime
+    pairs: list[PairData]
+
+
+class AnnotationPayload(BaseModel):
+    pair_id: str
+    label: int
+    input_tokens: int
+    output_tokens: int
+    raw_response_hash: str
+
+
+class AnnotationSubmission(BaseModel):
+    batch_id: str
+    model_id: str
+    quantization: str
+    annotations: list[AnnotationPayload]
+
+
+class AnnotationResult(BaseModel):
+    accepted: int
+    rejected: int
+    honeypot_accuracy: float | None = None
+    pairs_verified: int = 0
+    contributor_tokens: dict[str, int] = {}
+
+
+class ContributorProfile(BaseModel):
+    id: str
+    github_username: str
+    github_avatar_url: str | None = None
+    total_annotations: int = 0
+    reputation_score: float = 0.0
+    total_input_tokens: int = 0
+    total_output_tokens: int = 0
+    created_at: datetime | None = None
+    last_seen_at: datetime | None = None
+
+
+class NoPairsBackoff:
+    """Tracks consecutive 204s and computes wait duration."""
+
+    def __init__(self) -> None:
+        self._consecutive_empty = 0
+
+    def wait_duration(self) -> float:
+        """Get the next wait duration in seconds."""
+        duration = NO_PAIRS_INITIAL_WAIT * (2**self._consecutive_empty)
+        result: float = min(duration, NO_PAIRS_MAX_WAIT)
+        return result
+
+    def record_empty(self) -> None:
+        self._consecutive_empty += 1
+
+    def reset(self) -> None:
+        self._consecutive_empty = 0
+
+    @property
+    def consecutive_empty(self) -> int:
+        return self._consecutive_empty
+
+
+SUBMIT_TIMEOUT = httpx.Timeout(120.0, connect=10.0)
+
+
+class KombinatClient:
+    def __init__(self, base_url: str, access_token: str) -> None:
+        self.http = httpx.Client(
+            base_url=base_url,
+            headers={"Authorization": f"Bearer {access_token}"},
+            timeout=httpx.Timeout(30.0, connect=10.0),
+        )
+
+    def claim_batch(self, size: int = 100) -> BatchResponse | None:
+        """Claim a batch of pairs. Returns None if no pairs available (204)."""
+        resp = self._request_with_retry("POST", "/v1/batches/claim", json={"size": size})
+        if resp.status_code == 204:
+            return None
+        return BatchResponse.model_validate(resp.json())
+
+    def submit_annotations(self, submission: AnnotationSubmission) -> AnnotationResult:
+        """Submit a batch of annotations. Uses a longer timeout since processing
+        100 annotations with honeypot checks can be slow over Railway's proxy."""
+        resp = self._request_with_retry(
+            "POST",
+            "/v1/annotations",
+            json=submission.model_dump(),
+            timeout=SUBMIT_TIMEOUT,
+        )
+        return AnnotationResult.model_validate(resp.json())
+
+    def release_batch(self, batch_id: str) -> None:
+        """Release an unfinished batch back to the pool."""
+        self._request_with_retry("DELETE", f"/v1/batches/{batch_id}")
+
+    def get_profile(self) -> ContributorProfile:
+        """Get the contributor's profile and stats."""
+        resp = self._request_with_retry("GET", "/v1/contributors/me")
+        return ContributorProfile.model_validate(resp.json())
+
+    def close(self) -> None:
+        self.http.close()
+
+    def _request_with_retry(self, method: str, url: str, **kwargs: object) -> httpx.Response:
+        """Make an HTTP request with exponential backoff retry on 5xx/network errors."""
+        last_exception: Exception | None = None
+
+        for attempt in range(MAX_RETRIES + 1):
+            try:
+                resp = self.http.request(method, url, **kwargs)  # type: ignore[arg-type]
+
+                if resp.status_code == 401:
+                    raise AuthError("Authentication failed (401). Run 'annotator login'.")
+                if resp.status_code == 204:
+                    return resp
+                if 400 <= resp.status_code < 500:
+                    raise KombinatError(f"kombinat error {resp.status_code}: {resp.text}")
+                if resp.status_code >= 500:
+                    last_exception = KombinatError(
+                        f"kombinat server error {resp.status_code}: {resp.text}"
+                    )
+                    if attempt < MAX_RETRIES:
+                        self._backoff_sleep(attempt)
+                        continue
+                    raise last_exception
+
+                return resp
+
+            except (httpx.ConnectError, httpx.TimeoutException) as e:
+                last_exception = KombinatError(f"kombinat unreachable: {e}")
+                last_exception.__cause__ = e
+                if attempt < MAX_RETRIES:
+                    self._backoff_sleep(attempt)
+                    continue
+                raise KombinatError(f"kombinat unreachable after {MAX_RETRIES} retries: {e}") from e
+
+        msg = "Request failed after all retries"
+        raise KombinatError(msg) if last_exception is None else last_exception
+
+    def _backoff_sleep(self, attempt: int) -> None:
+        delay = min(MIN_RETRY_DELAY * (2**attempt) + random.uniform(0, 1), MAX_RETRY_DELAY)
+        logger.debug("Retrying in %.1fs (attempt %d/%d)", delay, attempt + 1, MAX_RETRIES)
+        time.sleep(delay)

annotator/config.py

@@ -0,0 +1,38 @@
+"""Annotator configuration via environment variables."""
+
+from __future__ import annotations
+
+from enum import IntEnum
+from pathlib import Path
+
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class ExitCode(IntEnum):
+    """CLI exit codes."""
+
+    SUCCESS = 0
+    AUTH_FAILURE = 1
+    NO_COMPATIBLE_HARDWARE = 2
+    MODEL_LOADING_FAILED = 3
+    KOMBINAT_UNREACHABLE = 4
+    UNRECOVERABLE = 5
+
+
+class Settings(BaseSettings):
+    model_config = SettingsConfigDict(
+        env_prefix="ANNOTATOR_",
+        env_file=Path.cwd() / ".env",
+        extra="ignore",
+    )
+
+    # Override with ANNOTATOR_KOMBINAT_URL for local dev against a non-production hub.
+    kombinat_url: str = "https://kombinat-production.up.railway.app"
+
+    batch_size: int = 100
+    chunk_size: int = 50
+    gpu_memory_utilization: float = 0.9
+    max_model_len: int = 4096
+    max_output_tokens: int = 256
+
+    annotator_home: Path = Path.home() / ".annotator"

annotator/engine/__init__.py

@@ -0,0 +1,34 @@
+"""Inference engine backends."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from annotator.engine.base import BaseEngine
+    from annotator.resolver import ResolvedRuntime
+
+
+def create_engine(
+    runtime: ResolvedRuntime,
+    gpu_memory_utilization: float = 0.9,
+    max_model_len: int = 4096,
+    max_output_tokens: int = 256,
+) -> BaseEngine:
+    """Create the appropriate engine for the resolved runtime."""
+    if runtime.backend == "vllm":
+        from annotator.engine.vllm import VLLMEngine
+
+        return VLLMEngine(
+            runtime.model_spec, gpu_memory_utilization, max_model_len, max_output_tokens
+        )
+    elif runtime.backend == "mlx":
+        from annotator.engine.mlx import MLXEngine
+
+        return MLXEngine(runtime.model_spec, max_output_tokens=max_output_tokens)
+    elif runtime.backend == "llama_cpp":
+        from annotator.engine.llama_cpp import LlamaCppEngine
+
+        return LlamaCppEngine(runtime.model_spec)
+    else:
+        raise ValueError(f"Unknown backend: {runtime.backend}")

annotator/engine/base.py

@@ -0,0 +1,54 @@
+"""Base engine interface and data structures."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from pydantic import BaseModel
+
+
+class LabelingInput(BaseModel):
+    """A single (query, document) pair to label."""
+
+    pair_id: str
+    query_text: str
+    doc_text: str
+
+
+class LabelingOutput(BaseModel):
+    """LLM response + engine metadata. Ready for submission to kombinat."""
+
+    pair_id: str
+    label: int
+    reasoning: str
+    input_tokens: int
+    output_tokens: int
+    raw_response_hash: str
+
+
+class EngineInfo(BaseModel):
+    """Model metadata — submitted to kombinat with every annotation."""
+
+    model_id: str
+    quantization: str
+    backend: str
+
+
+class BaseEngine(ABC):
+    @abstractmethod
+    def load(self) -> None:
+        """Download model (if not cached) and load into memory."""
+        ...
+
+    @abstractmethod
+    def label_batch(self, pairs: list[LabelingInput]) -> list[LabelingOutput]:
+        """Run inference on a batch. Returns results for successfully labeled pairs only.
+
+        Pairs that fail parsing/validation after retry are silently dropped.
+        """
+        ...
+
+    @abstractmethod
+    def info(self) -> EngineInfo:
+        """Return model metadata for submission to kombinat."""
+        ...