logogram 0.1.0__py3-none-any.whl

logogram/server/models.py

@@ -0,0 +1,426 @@
+"""Response models for the JSON API. Request bodies are defined beside their routes in app.py."""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict
+
+from logogram.schema import RunListing, SiteSummary
+from logogram.spec import ModelRef, PredictionSettings
+
+
+class _Out(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+# -- models and jobs ---------------------------------------------------------------------------
+
+
+class ModelInfoOut(_Out):
+    id: str
+    revision: str | None
+    architecture: str
+    n_layers: int
+    n_heads: int
+    d_model: int
+    d_head: int
+    d_mlp: int | None
+    d_vocab: int
+    n_ctx: int
+    n_params: int | None
+    dtype: str
+    device: str
+    device_name: str
+    process_weights: bool
+    site_kinds: list[str]
+    backend: str
+    backend_version: str
+    extra: dict[str, Any]
+
+
+class ModelStatus(_Out):
+    state: Literal["none", "loading", "ready", "error"]
+    id: str | None = None
+    model_id: str | None = None
+    stage: str | None = None
+    done: int | None = None
+    total: int | None = None
+    file: str | None = None
+    revision: str | None = None
+    error: str | None = None
+    info: ModelInfoOut | None = None
+    memory: int | None = None
+    ref: ModelRef | None = None
+
+
+class JobInfo(_Out):
+    id: str
+    kind: Literal["load_model", "load_sae", "run"]
+    title: str
+    status: Literal["running", "finished", "failed", "cancelled"]
+    run_id: str | None
+    project_session: str | None = None
+    progress: dict[str, Any]
+    error: str | None
+    started: float
+    version: int
+    cancelling: bool | None = None
+
+
+class StartedRun(_Out):
+    run_id: str
+    job: JobInfo
+
+
+class CancelOut(_Out):
+    job: JobInfo | None
+
+
+class DraftSaved(_Out):
+    run_id: str
+    path: str
+
+
+# -- projects and files ------------------------------------------------------------------------
+
+
+class DatasetListing(_Out):
+    name: str
+    path: str
+    modified: str | None
+    n: int | None = None
+    has_positions: bool | None = None
+    sha256: str | None = None
+    error: str | None = None
+
+
+class ProjectInfo(_Out):
+    name: str
+    description: str
+    path: str
+    session_id: str
+    datasets: list[DatasetListing]
+
+
+class RecentProject(_Out):
+    path: str
+    name: str
+    opened: str | None
+    modified: str
+
+
+class FolderEntry(_Out):
+    name: str
+    path: str
+    is_project: bool
+
+
+class FolderListing(_Out):
+    path: str
+    parent: str | None
+    is_project: bool
+    entries: list[FolderEntry]
+
+
+class UpdateStatus(_Out):
+    current: str
+    released: str | None
+    automatic: bool | None
+    latest: str | None
+    available: bool
+    old: bool
+    checked_at: str | None
+    error: str | None
+    notes_url: str | None
+    command: str
+
+
+class ServerState(_Out):
+    version: str
+    project: ProjectInfo | None
+    model: ModelStatus
+    job: JobInfo | None
+    first_run: bool
+    theme: Literal["light", "dark", "system"]
+    projects_parent: str
+    update: UpdateStatus
+    sae: dict[str, Any] = {"state": "none"}
+
+
+class Settings(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    system_check_seen: bool | None = None
+    theme: str | None = None
+    update_check: bool | None = None
+
+
+class OkOut(_Out):
+    ok: bool
+
+
+# -- datasets and prompts ----------------------------------------------------------------------
+
+
+class PromptIssue(_Out):
+    index: int
+    kind: str
+    message: str
+
+
+class DatasetDetail(_Out):
+    name: str
+    path: str
+    sha256: str
+    n: int
+    records: list[dict[str, Any]]
+    issues: list[PromptIssue] | None = None
+    lengths: list[int] | None = None
+
+
+class DatasetCreated(_Out):
+    name: str
+    path: str
+    n: int
+
+
+class IOITemplateOut(_Out):
+    id: str
+    text: str
+    default: bool
+
+
+class TokenRow(_Out):
+    tokens: list[str]
+    ids: list[int]
+
+
+class AnswerTokens(_Out):
+    text: str
+    tokens: list[str]
+    id: int | None
+
+
+class TokenStrip(_Out):
+    clean: TokenRow
+    corrupt: TokenRow
+    aligned: bool
+    differs: list[int]
+    labels: dict[str, int]
+    answer: AnswerTokens
+    distractor: AnswerTokens
+    issues: list[PromptIssue]
+
+
+# -- analyses ----------------------------------------------------------------------------------
+
+
+class TopToken(_Out):
+    token: str
+    id: int
+    prob: float
+
+
+class BaselinePrompt(_Out):
+    index: int
+    clean: str
+    corrupt: str
+    answer: str
+    distractor: str
+    clean_logit_diff: float | None
+    corrupt_logit_diff: float | None
+    clean_answer_prob: float | None
+    corrupt_answer_prob: float | None
+    clean_top: list[TopToken]
+    corrupt_top: list[TopToken]
+
+
+class BaselineSummaryOut(_Out):
+    clean_logit_diff: float | None
+    corrupt_logit_diff: float | None
+    gap: float | None
+    clean_prefers_answer: int
+    corrupt_prefers_answer: int
+    clean_answer_prob: float | None
+    corrupt_answer_prob: float | None
+
+
+class BaselineReport(_Out):
+    n: int
+    prompts: list[BaselinePrompt]
+    issues: list[PromptIssue]
+    summary: BaselineSummaryOut | None
+
+
+class AttentionData(_Out):
+    layer: int
+    head: int
+    which: Literal["clean", "corrupt"]
+    index: int
+    tokens: list[str]
+    labels: dict[str, int]
+    pattern: list[list[float]]
+    average: list[list[float]]
+    average_tokens: list[str | None]
+    average_labels: dict[str, int]
+    n_average: int
+    n_total: int
+    length: int
+
+
+class PredictionLayer(_Out):
+    layer: int
+    top: list[TopToken]
+    answer_prob: float
+    distractor_prob: float
+    logit_diff: float
+
+
+class PredictionReport(_Out):
+    index: int
+    position: int
+    tokens: list[str]
+    settings: PredictionSettings
+    layers: list[PredictionLayer]
+    batch_members: list[int]
+    answer: str
+    distractor: str
+    issues: list[PromptIssue]
+
+
+class Presets(_Out):
+    presets: list[dict[str, Any]]
+    suggestions: list[str]
+
+
+class ArchitectureOut(_Out):
+    n_layers: int
+    n_heads: int
+    d_model: int
+    d_mlp: int
+    d_vocab: int
+    n_ctx: int
+    architecture: str
+
+
+class MemoryEstimateOut(_Out):
+    device: str
+    dtype: str
+    n_params: int
+    weights: int
+    activations: int
+    margin: int
+    total: int
+    available: int
+    verdict: Literal["fits", "tight", "wont_fit"]
+    explanation: str
+
+
+class EstimateOut(_Out):
+    id: str
+    revision: str
+    architecture: ArchitectureOut
+    download_bytes: int
+    total_bytes: int
+    gated: bool
+    # Whether TransformerLens can load this architecture, and if not, why (before downloading).
+    supported: bool
+    support_note: str | None
+    estimate: MemoryEstimateOut
+
+
+# -- runs --------------------------------------------------------------------------------------
+
+
+class RunDetail(_Out):
+    """A run's files as stored. They were validated when written; older files are passed through."""
+
+    id: str
+    spec: dict[str, Any] | None
+    summary: dict[str, Any] | None
+    manifest: dict[str, Any] | None
+    predictions: PredictionReport | None = None
+    listing: RunListing | None
+    folder: str
+
+
+class PromptEvidence(_Out):
+    index: int
+    clean: str | None
+    corrupt: str | None
+    answer: str | None
+    distractor: str | None
+    effect: float | None
+    delta: float | None
+    patched_logit_diff: float | None
+    receiver_logit_diff: float | None
+    reference_logit_diff: float | None
+    patched_answer_prob: float | None
+    receiver_answer_prob: float | None
+    flipped: bool
+
+
+class SiteDetail(_Out):
+    run_id: str
+    site: SiteSummary
+    receiver: Literal["clean", "corrupt"]
+    reference: Literal["clean", "corrupt"]
+    metric: dict[str, Any]
+    statistics: dict[str, Any]
+    prompts: list[PromptEvidence]
+    strongest: list[int]
+    weakest: list[int]
+    dataset_changed: bool
+    dataset_available: bool
+
+
+class ComparisonChange(_Out):
+    label: str
+    kind: str
+    layer: int
+    head: int | None
+    feature: int | None = None
+    position_key: str
+    variant_key: str | None = None
+    index_a: int
+    index_b: int
+    row: int
+    col: int
+    effect_a: dict[str, float | None]
+    effect_b: dict[str, float | None]
+    rank_a: int
+    rank_b: int
+    flags: list[Literal["sign", "left_top", "entered_top"]]
+
+
+class DiffCell(_Out):
+    index_a: int
+    row: int
+    col: int
+    value: float
+
+
+class SpecDifference(_Out):
+    path: str
+    a: Any
+    b: Any
+
+
+class Comparison(_Out):
+    run_a: str
+    run_b: str
+    n_common: int
+    n_a: int
+    n_b: int
+    spearman: float | None
+    top_k: int
+    top_overlap: int
+    top_a: list[str]
+    top_b: list[str]
+    changes: list[ComparisonChange]
+    flagged: list[int]
+    n_sign_changes: int
+    diff: list[DiffCell]
+    same_layout: bool
+    spec_differences: list[SpecDifference]

logogram/server/security.py

@@ -0,0 +1,212 @@
+"""Local server security.
+
+* The server binds to 127.0.0.1 only.
+* Every HTTP and WebSocket request must carry the session token: as the HttpOnly cookie set when
+  the launch URL (``/?token=...``) is opened, or as an ``Authorization: Bearer`` header.
+* The browser is opened with a single-use launch code instead of the token, because a browser's
+  command line can be read by other users of the machine. The URL printed in the terminal carries
+  the token itself, for opening more tabs or another browser.
+* The Host header must name this server (defeats DNS rebinding). Requests that carry an Origin
+  must come from this server's origin; WebSocket handshakes must carry one, and so must
+  cookie-authenticated requests that change anything (or a Fetch Metadata header). Fetch Metadata
+  headers must say same-origin. No CORS headers are sent.
+"""
+
+from __future__ import annotations
+
+import hmac
+import secrets
+import threading
+import time
+from dataclasses import dataclass, field
+from urllib.parse import urlsplit
+
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+SAFE_METHODS = {"GET", "HEAD", "OPTIONS"}
+# The launch URL. /api/session also works behind the Vite dev proxy.
+LOGIN_PATHS = {"/", "/api/session"}
+
+
+def new_token() -> str:
+    return secrets.token_urlsafe(32)
+
+
+@dataclass
+class SecurityConfig:
+    token: str
+    port: int
+    dev_origins: list[str] = field(default_factory=list)
+    _launch: tuple[str, float] | None = field(default=None, repr=False)
+    _launch_lock: threading.Lock = field(default_factory=threading.Lock, repr=False)
+
+    @property
+    def cookie_name(self) -> str:
+        return f"logogram_{self.port}"
+
+    def allowed_hosts(self) -> set[str]:
+        hosts = {f"127.0.0.1:{self.port}", f"localhost:{self.port}", f"[::1]:{self.port}"}
+        for origin in self.dev_origins:
+            hosts.add(urlsplit(origin).netloc)
+        return hosts
+
+    def allowed_origins(self) -> set[str]:
+        origins = {f"http://{h}" for h in self.allowed_hosts()}
+        origins.update(self.dev_origins)
+        return origins
+
+    def token_ok(self, candidate: str | None) -> bool:
+        return bool(candidate) and hmac.compare_digest(candidate.encode(), self.token.encode())
+
+    def new_launch_code(self, ttl: float = 600.0) -> str:
+        """A code that logs in once, within ``ttl`` seconds, for the URL the browser is opened
+        with."""
+        code = secrets.token_urlsafe(24)
+        with self._launch_lock:
+            self._launch = (code, time.monotonic() + ttl)
+        return code
+
+    def use_launch_code(self, candidate: str) -> bool:
+        with self._launch_lock:
+            if self._launch is None or not candidate:
+                return False
+            code, expires = self._launch
+            if time.monotonic() > expires:
+                self._launch = None
+                return False
+            if not hmac.compare_digest(candidate.encode(), code.encode()):
+                return False
+            self._launch = None  # once only
+            return True
+
+
+def _header(scope: Scope, name: bytes) -> str | None:
+    for key, value in scope.get("headers", []):
+        if key == name:
+            return value.decode("latin-1")
+    return None
+
+
+def _cookies(scope: Scope, name: str) -> list[str]:
+    """Every cookie called ``name``. Cookies aren't separated by port, so a page served from
+    another local port can add one with the same name; it mustn't hide the real one."""
+    values = []
+    for key, value in scope.get("headers", []):
+        if key != b"cookie":
+            continue
+        for part in value.decode("latin-1").split(";"):
+            k, _, v = part.strip().partition("=")
+            if k == name:
+                values.append(v)
+    return values
+
+
+def _query_token(scope: Scope) -> str | None:
+    query = scope.get("query_string", b"").decode("latin-1")
+    for part in query.split("&"):
+        key, _, value = part.partition("=")
+        if key == "token":
+            return value
+    return None
+
+
+class SecurityMiddleware:
+    """Pure ASGI middleware so it covers HTTP and WebSocket alike."""
+
+    def __init__(self, app: ASGIApp, config: SecurityConfig) -> None:
+        self.app = app
+        self.config = config
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] not in ("http", "websocket"):
+            await self.app(scope, receive, send)
+            return
+        cfg = self.config
+        host = _header(scope, b"host")
+        if host not in cfg.allowed_hosts():
+            await self._deny(scope, receive, send, 421, "Unknown host.")
+            return
+        origin = _header(scope, b"origin")
+        method = scope.get("method", "GET")
+        needs_origin = scope["type"] == "websocket" or method not in SAFE_METHODS
+        if origin is not None and origin not in cfg.allowed_origins():
+            await self._deny(scope, receive, send, 403, "Cross-origin requests are not allowed.")
+            return
+        if needs_origin and origin is None and scope["type"] == "websocket":
+            await self._deny(scope, receive, send, 403, "Missing origin.")
+            return
+        # Fetch Metadata: refuse requests started by another site, including other local ports
+        # (which count as the same *site*, so SameSite cookies alone don't stop them).
+        fetch_site = _header(scope, b"sec-fetch-site")
+        if fetch_site is not None and fetch_site not in ("same-origin", "none"):
+            await self._deny(scope, receive, send, 403, "Cross-site requests are not allowed.")
+            return
+
+        # Opening the launch URL exchanges the token (or the single-use launch code) for a
+        # cookie, then drops it from the URL.
+        query_token = _query_token(scope)
+        if scope["type"] == "http" and query_token is not None and scope["path"] in LOGIN_PATHS:
+            if cfg.token_ok(query_token) or cfg.use_launch_code(query_token):
+                await self._login(send)
+            else:
+                await self._deny(scope, receive, send, 401, _LOGIN_HELP)
+            return
+
+        bearer = _header(scope, b"authorization")
+        bearer_token = bearer[7:] if bearer and bearer.lower().startswith("bearer ") else None
+        if cfg.token_ok(bearer_token):
+            await self.app(scope, receive, send)
+            return
+        if any(cfg.token_ok(c) for c in _cookies(scope, cfg.cookie_name)):
+            # A browser sends the cookie by itself, so a request that changes something must
+            # show where it came from (every current browser sends one of these).
+            if needs_origin and origin is None and fetch_site is None:
+                await self._deny(scope, receive, send, 403, "Missing origin.")
+                return
+            await self.app(scope, receive, send)
+            return
+        await self._deny(scope, receive, send, 401, _LOGIN_HELP)
+
+    async def _login(self, send: Send) -> None:
+        cookie = f"{self.config.cookie_name}={self.config.token}; HttpOnly; SameSite=Strict; Path=/"
+        await send(
+            {
+                "type": "http.response.start",
+                "status": 303,
+                "headers": [
+                    (b"location", b"/"),
+                    (b"set-cookie", cookie.encode("latin-1")),
+                    (b"cache-control", b"no-store"),
+                    (b"referrer-policy", b"no-referrer"),
+                ],
+            }
+        )
+        await send({"type": "http.response.body", "body": b""})
+
+    async def _deny(
+        self, scope: Scope, receive: Receive, send: Send, status: int, text: str
+    ) -> None:
+        if scope["type"] == "websocket":
+            await send({"type": "websocket.close", "code": 4401 if status == 401 else 4403})
+            return
+        body = text.encode("utf-8")
+        await send(
+            {
+                "type": "http.response.start",
+                "status": status,
+                "headers": [
+                    (b"content-type", b"text/plain; charset=utf-8"),
+                    (b"content-length", str(len(body)).encode()),
+                    (b"cache-control", b"no-store"),
+                ],
+            }
+        )
+        await send({"type": "http.response.body", "body": body})
+
+
+_LOGIN_HELP = (
+    "Logogram needs its session token.\n\n"
+    "Open the link printed in the terminal where you started `logogram` "
+    "(it ends with ?token=...). Each start creates a new token, and the link the browser was "
+    "first opened with works only once.\n"
+)