caudate-cli 0.1.0__py3-none-any.whl

core/anthropic_auth.py

@@ -0,0 +1,91 @@
+"""Subscription-OAuth helper for Cognos's web UI calls to Anthropic.
+
+Background
+----------
+Claude Code (the CLI) authenticates to Anthropic with an OAuth Bearer
+token stored at ``~/.claude/.credentials.json``. The Cognos passthrough
+proxy at ``/v1/messages`` already forwards that token verbatim from
+incoming Claude Code requests — it never has to read the file itself.
+
+The web UI is a different story: it calls Cognos's own ``/chat`` endpoint,
+which goes through ``LLMProvider`` → LiteLLM → Anthropic. LiteLLM expects
+``ANTHROPIC_API_KEY`` in the environment. The user has a subscription, not
+an API key — so without help, the web UI silently fails on Opus.
+
+Scope
+-----
+This helper is intentionally **scoped to the web UI only**. It uses a
+``ContextVar`` set by the ``/chat`` and ``/chat/stream`` endpoints, so
+that the same ``LLMProvider`` instance behaves normally for every other
+caller (background agentic loops, voice, MCP). Other code paths that hit
+Anthropic without a real API key will still fail with the original auth
+error — that's the desired blast radius.
+
+Token freshness
+---------------
+We read the credentials file on every call rather than caching, since
+Claude Code may rotate the access token. The file is ~470 bytes; the
+parse cost is negligible next to the network round-trip to Anthropic.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from contextlib import contextmanager
+from contextvars import ContextVar
+from pathlib import Path
+from typing import Iterator
+
+logger = logging.getLogger(__name__)
+
+
+_CREDENTIALS_PATH = Path.home() / ".claude" / ".credentials.json"
+
+# Set by web-UI endpoints during their request scope. LLMProvider checks
+# this; if False (the default), behaviour is identical to before.
+_use_subscription: ContextVar[bool] = ContextVar(
+    "cognos_use_subscription_auth", default=False,
+)
+
+
+def is_active() -> bool:
+    """True if the current async-context is allowed to use the subscription."""
+    return _use_subscription.get()
+
+
+@contextmanager
+def subscription_auth_scope() -> Iterator[None]:
+    """Enable subscription-OAuth auth for the duration of the ``with`` block.
+
+    Usage (inside a request handler)::
+
+        with subscription_auth_scope():
+            reply = await agent.chat(...)
+    """
+    token = _use_subscription.set(True)
+    try:
+        yield
+    finally:
+        _use_subscription.reset(token)
+
+
+def read_subscription_token() -> str | None:
+    """Read the current Claude Code OAuth access token, or None if missing.
+
+    Returns None on any failure (file missing, bad JSON, missing key) so
+    callers can fall through gracefully rather than crash.
+    """
+    try:
+        with _CREDENTIALS_PATH.open() as f:
+            data = json.load(f)
+    except (OSError, json.JSONDecodeError) as e:
+        logger.debug(f"could not read {_CREDENTIALS_PATH}: {e}")
+        return None
+    oauth = data.get("claudeAiOauth")
+    if not isinstance(oauth, dict):
+        return None
+    token = oauth.get("accessToken")
+    if not isinstance(token, str) or not token:
+        return None
+    return token

core/background.py

@@ -0,0 +1,113 @@
+"""Background task management — run agent turns asynchronously.
+
+A background task is a CognosAgent.chat call that's been kicked off
+without blocking the REPL. The pool tracks each task by id, exposes
+status (pending/running/done/failed), and stores the final reply so the
+user can `watch` it later.
+
+This is intentionally process-local: when the CLI exits, in-flight
+tasks are abandoned. Persistence would require a separate worker
+process or supervisor — outside scope of a UX feature.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+import uuid
+from dataclasses import dataclass, field
+from typing import Any, Awaitable, Callable
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class BackgroundTask:
+    id: str = field(default_factory=lambda: uuid.uuid4().hex[:8])
+    label: str = ""
+    status: str = "pending"          # pending | running | done | failed | cancelled
+    started_at: float = 0.0
+    finished_at: float = 0.0
+    result: str | None = None
+    error: str | None = None
+    _task: asyncio.Task | None = field(default=None, repr=False)
+
+    @property
+    def duration(self) -> float:
+        if not self.started_at:
+            return 0.0
+        end = self.finished_at or time.time()
+        return end - self.started_at
+
+
+class BackgroundTaskPool:
+    """In-process pool of CognosAgent invocations."""
+
+    def __init__(self) -> None:
+        self._tasks: dict[str, BackgroundTask] = {}
+
+    def submit(
+        self,
+        coro_factory: Callable[[], Awaitable[Any]],
+        label: str = "",
+    ) -> BackgroundTask:
+        """Schedule `coro_factory()` on the running loop.
+
+        Must be called from within a running asyncio loop. The factory
+        is invoked once, immediately, and the resulting awaitable is
+        wrapped in a Task. The BackgroundTask handle is returned
+        synchronously so the caller can show an id.
+        """
+        bg = BackgroundTask(label=label)
+        bg.status = "running"
+        bg.started_at = time.time()
+        loop = asyncio.get_running_loop()
+
+        async def _runner() -> None:
+            try:
+                result = await coro_factory()
+                bg.result = str(result) if result is not None else ""
+                bg.status = "done"
+            except asyncio.CancelledError:
+                bg.status = "cancelled"
+                raise
+            except Exception as e:
+                bg.error = str(e)
+                bg.status = "failed"
+                logger.exception(f"Background task {bg.id} failed: {e}")
+            finally:
+                bg.finished_at = time.time()
+
+        bg._task = loop.create_task(_runner(), name=f"cognos-bg-{bg.id}")
+        self._tasks[bg.id] = bg
+        return bg
+
+    def get(self, task_id: str) -> BackgroundTask | None:
+        return self._tasks.get(task_id)
+
+    def list(self) -> list[BackgroundTask]:
+        return sorted(self._tasks.values(), key=lambda t: t.started_at, reverse=True)
+
+    def cancel(self, task_id: str) -> bool:
+        bg = self._tasks.get(task_id)
+        if bg is None or bg._task is None or bg._task.done():
+            return False
+        bg._task.cancel()
+        return True
+
+    def prune_finished(self) -> int:
+        before = len(self._tasks)
+        self._tasks = {
+            tid: t for tid, t in self._tasks.items()
+            if t.status in ("pending", "running")
+        }
+        return before - len(self._tasks)
+
+
+# Process-wide pool — mirrors usage tracker.
+_GLOBAL = BackgroundTaskPool()
+
+
+def get_global_pool() -> BackgroundTaskPool:
+    return _GLOBAL

core/banner.py

@@ -0,0 +1,134 @@
+"""Branded boot banner for the Cognos CLI.
+
+Lands on stdout when `cognos` (or `cognos interactive`) starts so the
+launch feels like a real CLI tool — ASCII logo, brain status, a one-
+line tip on slash commands. Replaces the prior ad-hoc `console.print`
+blob in `main.py:interactive`.
+
+The banner is plain Rich — no Textual, no full-screen takeover — so
+piping (`cognos -p ...`) and non-TTY contexts still work cleanly.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from rich.align import Align
+from rich.console import Console, Group
+from rich.panel import Panel
+from rich.text import Text
+
+
+_LOGO = r"""
+  ____
+ / ___|___   __ _ _ __   ___  ___
+| |   / _ \ / _` | '_ \ / _ \/ __|
+| |__| (_) | (_| | | | | (_) \__ \
+ \____\___/ \__, |_| |_|\___/|___/
+            |___/
+""".rstrip("\n").lstrip("\n")
+
+
+def _version() -> str:
+    try:
+        from importlib.metadata import version
+        return version("cognos")
+    except Exception:
+        return "0.1.0"
+
+
+def _brain_summary(agent: Any) -> Text:
+    """Two-line summary of which models are wired in."""
+    from llm.router import DualLLMProvider
+    t = Text()
+    if isinstance(agent.llm, DualLLMProvider):
+        t.append("S1 ", style="bold cyan")
+        t.append(getattr(agent.llm_fast, "model", "?"), style="white")
+        t.append("   S2 ", style="bold magenta")
+        t.append(getattr(agent.llm_slow, "model", "?"), style="white")
+    else:
+        t.append("model ", style="bold")
+        t.append(getattr(agent.llm, "model", "?"), style="white")
+    return t
+
+
+def _caudate_summary(agent: Any) -> Text | None:
+    """Caudate router trust + active model — only if she's loaded."""
+    t = Text()
+    try:
+        caudate = getattr(agent, "caudate", None)
+        if caudate is None:
+            return None
+        trust = getattr(caudate, "trust_level", None) or getattr(caudate, "trust", None) or "whisper"
+        params = getattr(caudate, "n_params", None)
+        t.append("caudate ", style="bold green")
+        t.append(str(trust).lower(), style="white")
+        if params:
+            t.append(f"  {params/1e6:.2f}M params", style="dim")
+        return t
+    except Exception:
+        return None
+
+
+def _session_summary(agent: Any, resumed: bool) -> Text:
+    t = Text()
+    try:
+        sid = agent.session.id
+        n = len(agent.session.messages)
+    except Exception:
+        sid, n = "?", 0
+    if resumed:
+        t.append("resumed ", style="yellow")
+        t.append(f"{sid[:8]} · {n} messages", style="dim")
+    else:
+        t.append("new session ", style="green")
+        t.append(sid[:8], style="dim")
+    return t
+
+
+def print_startup_banner(
+    agent: Any,
+    console: Console,
+    *,
+    resumed: bool = False,
+    extras: list[str] | None = None,
+) -> None:
+    """Render the boot screen.
+
+    `extras` is a list of mode flags shown as a footer (e.g. ["streaming",
+    "thinking"]). Empty list / None hides that line.
+    """
+    logo = Text(_LOGO, style="bold cyan")
+    version_line = Text()
+    version_line.append(f"v{_version()}", style="dim")
+    version_line.append("  ·  ", style="dim")
+    version_line.append("local-first cognitive agent", style="italic dim")
+
+    rows: list[Any] = [
+        Align.center(logo),
+        Align.center(version_line),
+        Text(""),
+        _brain_summary(agent),
+    ]
+    caudate = _caudate_summary(agent)
+    if caudate is not None:
+        rows.append(caudate)
+    rows.append(_session_summary(agent, resumed))
+    if extras:
+        flags = Text("  ".join(extras), style="dim italic")
+        rows.append(flags)
+    rows.append(Text(""))
+
+    tip = Text()
+    tip.append("type ", style="dim")
+    tip.append("/", style="bold")
+    tip.append(" for the command palette  ·  ", style="dim")
+    tip.append("/quit", style="bold")
+    tip.append(" to exit", style="dim")
+    rows.append(tip)
+
+    console.print(Panel(
+        Group(*rows),
+        border_style="cyan",
+        padding=(1, 2),
+    ))

core/bootstrap.py

@@ -0,0 +1,292 @@
+"""First-run bootstrap + health checks.
+
+A stranger who runs `pip install cognos` shouldn't have to know about
+Ollama, HuggingFace, Anthropic keys, or `data/nn/caudate.pt`. This
+module covers the cold-start path:
+
+- `cognos init`   — interactive wizard: picks system1/system2, downloads
+                    Caudate weights, writes `~/.cognos/settings.json`,
+                    creates a starter sandbox directory.
+- `cognos doctor` — read-only diagnostic: which dependencies are wired,
+                    which are missing, exact one-line fixes per gap.
+
+Caudate weights live on HuggingFace as a public model repo (set via the
+`COGNOS_CAUDATE_HF_REPO` env var, defaulting to `raveuk/cognos-caudate`).
+The repo is expected to contain `caudate.pt` and `caudate.meta.json`.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import socket
+import subprocess
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+from rich.console import Console
+from rich.panel import Panel
+from rich.prompt import Confirm, Prompt
+from rich.table import Table
+
+
+_DEFAULT_HF_REPO = "raveuk/cognos-caudate"
+_CAUDATE_FILES = ("caudate.pt", "caudate.meta.json")
+
+
+# ---------------------------------------------------------------------
+# Locations
+# ---------------------------------------------------------------------
+def cognos_home() -> Path:
+    """`~/.cognos`. Created on demand."""
+    p = Path(os.path.expanduser("~")) / ".cognos"
+    p.mkdir(parents=True, exist_ok=True)
+    return p
+
+
+def settings_path() -> Path:
+    return cognos_home() / "settings.json"
+
+
+def caudate_dir() -> Path:
+    """Where Caudate weights live for an installed cognos.
+
+    Repo checkouts have `<repo>/data/nn/caudate.pt`; pip-installed users
+    get `~/.cognos/nn/caudate.pt`.
+    """
+    p = cognos_home() / "nn"
+    p.mkdir(parents=True, exist_ok=True)
+    return p
+
+
+def is_first_run() -> bool:
+    return not settings_path().exists()
+
+
+# ---------------------------------------------------------------------
+# Checks
+# ---------------------------------------------------------------------
+@dataclass
+class CheckResult:
+    name: str
+    ok: bool
+    detail: str
+    fix: str | None = None
+
+
+def check_ollama() -> CheckResult:
+    """Is Ollama reachable on the default port?"""
+    host = os.environ.get("OLLAMA_HOST", "127.0.0.1:11434")
+    if "://" in host:
+        host = host.split("://", 1)[1]
+    h, _, port = host.partition(":")
+    port = int(port or "11434")
+    try:
+        with socket.create_connection((h, port), timeout=1.5):
+            pass
+        return CheckResult("ollama", True, f"reachable at {h}:{port}")
+    except Exception:
+        return CheckResult(
+            "ollama", False, f"not reachable at {h}:{port}",
+            fix="install from https://ollama.com and run `ollama serve`",
+        )
+
+
+def check_caudate() -> CheckResult:
+    """Are Caudate weights present (either in repo or in ~/.cognos)?"""
+    repo = Path.cwd() / "data" / "nn" / "caudate.pt"
+    user = caudate_dir() / "caudate.pt"
+    for p in (repo, user):
+        if p.exists():
+            mb = p.stat().st_size / 1e6
+            return CheckResult("caudate", True, f"{p} ({mb:.1f} MB)")
+    return CheckResult(
+        "caudate", False, "no caudate.pt found",
+        fix="run `cognos init` (downloads weights from HuggingFace)",
+    )
+
+
+def check_anthropic_key() -> CheckResult:
+    if os.environ.get("ANTHROPIC_API_KEY"):
+        return CheckResult("anthropic_key", True, "ANTHROPIC_API_KEY set in env")
+    return CheckResult(
+        "anthropic_key", False, "ANTHROPIC_API_KEY not set",
+        fix="export ANTHROPIC_API_KEY=sk-ant-... (only needed if "
+            "system1 or system2 is an `anthropic/...` model)",
+    )
+
+
+def check_settings() -> CheckResult:
+    p = settings_path()
+    if not p.exists():
+        return CheckResult(
+            "settings", False, f"{p} missing",
+            fix="run `cognos init`",
+        )
+    try:
+        data = json.loads(p.read_text())
+        s1 = data.get("system1") or "?"
+        s2 = data.get("system2") or "?"
+        return CheckResult("settings", True, f"system1={s1}, system2={s2}")
+    except Exception as e:
+        return CheckResult(
+            "settings", False, f"{p} unreadable ({e})",
+            fix="delete the file and re-run `cognos init`",
+        )
+
+
+def run_checks() -> list[CheckResult]:
+    return [
+        check_settings(),
+        check_ollama(),
+        check_caudate(),
+        check_anthropic_key(),
+    ]
+
+
+# ---------------------------------------------------------------------
+# Doctor — read-only diagnostic
+# ---------------------------------------------------------------------
+def cmd_doctor(console: Console) -> int:
+    """Print a status table + fix hints. Returns 0 if all green."""
+    results = run_checks()
+    table = Table(show_header=True, header_style="bold", expand=False)
+    table.add_column("check", style="cyan")
+    table.add_column("status")
+    table.add_column("detail", style="dim")
+    for r in results:
+        mark = "[green]✓[/green]" if r.ok else "[red]✗[/red]"
+        table.add_row(r.name, mark, r.detail)
+    console.print(Panel(table, title="cognos doctor", border_style="cyan"))
+
+    missing = [r for r in results if not r.ok]
+    if missing:
+        console.print()
+        console.print("[bold]fixes:[/bold]")
+        for r in missing:
+            if r.fix:
+                console.print(f"  [yellow]{r.name}[/yellow]: {r.fix}")
+        return 1
+    console.print("[green]all checks passed[/green]")
+    return 0
+
+
+# ---------------------------------------------------------------------
+# Init — interactive wizard
+# ---------------------------------------------------------------------
+_MODEL_PRESETS = [
+    ("local-only",
+     "ollama/glm-5.1:cloud", "ollama/qwen3.6:35b",
+     "fully local; both brains run on Ollama"),
+    ("hybrid",
+     "ollama/glm-5.1:cloud", "anthropic/claude-haiku-4-5",
+     "local fast, hosted slow (recommended)"),
+    ("hosted-only",
+     "anthropic/claude-haiku-4-5", "anthropic/claude-opus-4-7",
+     "no Ollama needed; requires ANTHROPIC_API_KEY"),
+]
+
+
+def _download_caudate(console: Console, repo: str) -> bool:
+    """Pull caudate.pt + caudate.meta.json from a HF repo into ~/.cognos/nn/.
+
+    Returns True on success. Best-effort: if `huggingface_hub` isn't
+    installed we fall back to printing the manual download URLs.
+    """
+    target = caudate_dir()
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError:
+        console.print(
+            f"[yellow]huggingface_hub not installed.[/yellow] Download "
+            f"these two files into {target}:"
+        )
+        for fname in _CAUDATE_FILES:
+            console.print(
+                f"  https://huggingface.co/{repo}/resolve/main/{fname}"
+            )
+        return False
+    try:
+        for fname in _CAUDATE_FILES:
+            console.print(f"  fetching [cyan]{fname}[/cyan] from {repo}…")
+            src = hf_hub_download(repo_id=repo, filename=fname)
+            shutil.copy2(src, target / fname)
+        console.print(f"[green]caudate weights ready at {target}[/green]")
+        return True
+    except Exception as e:
+        console.print(f"[red]download failed:[/red] {e}")
+        console.print(
+            f"manual fallback — get the files from "
+            f"https://huggingface.co/{repo} and drop them in {target}"
+        )
+        return False
+
+
+def cmd_init(console: Console, *, force: bool = False) -> int:
+    """Interactive first-run setup."""
+    console.print(Panel(
+        "[bold]cognos init[/bold] — first-run setup\n"
+        "Picks your fast/slow models, downloads Caudate, "
+        "writes ~/.cognos/settings.json",
+        border_style="cyan",
+    ))
+
+    sp = settings_path()
+    if sp.exists() and not force:
+        if not Confirm.ask(
+            f"settings already exist at {sp} — overwrite?", default=False,
+        ):
+            console.print("[dim]aborted — existing settings preserved[/dim]")
+            return 1
+
+    # --- model preset
+    console.print("\n[bold]model preset[/bold]")
+    table = Table(show_header=True, header_style="bold")
+    table.add_column("#", style="cyan")
+    table.add_column("name")
+    table.add_column("system1 (fast)")
+    table.add_column("system2 (slow)")
+    table.add_column("notes", style="dim")
+    for i, (name, s1, s2, note) in enumerate(_MODEL_PRESETS, 1):
+        table.add_row(str(i), name, s1, s2, note)
+    console.print(table)
+    choice = Prompt.ask(
+        "pick a preset (1-3) or [c]ustom", choices=["1", "2", "3", "c"], default="2",
+    )
+    if choice == "c":
+        s1 = Prompt.ask("system1 (fast model)", default="ollama/glm-5.1:cloud")
+        s2 = Prompt.ask("system2 (slow model)", default="anthropic/claude-haiku-4-5")
+    else:
+        _, s1, s2, _ = _MODEL_PRESETS[int(choice) - 1]
+
+    # --- caudate download
+    download = Confirm.ask(
+        "\ndownload caudate weights (~62 MB) from HuggingFace?", default=True,
+    )
+    if download:
+        repo = os.environ.get("COGNOS_CAUDATE_HF_REPO", _DEFAULT_HF_REPO)
+        _download_caudate(console, repo)
+
+    # --- write settings
+    settings: dict[str, object] = {"system1": s1, "system2": s2}
+    if sp.exists():
+        try:
+            existing = json.loads(sp.read_text())
+            existing.update(settings)
+            settings = existing
+        except Exception:
+            pass
+    sp.write_text(json.dumps(settings, indent=2) + "\n")
+    console.print(f"\n[green]wrote {sp}[/green]")
+
+    # --- key reminder
+    if any("anthropic" in m for m in (s1, s2)) and not os.environ.get("ANTHROPIC_API_KEY"):
+        console.print(
+            "\n[yellow]heads up:[/yellow] you picked an anthropic model. "
+            "set [bold]ANTHROPIC_API_KEY[/bold] in your shell before running."
+        )
+
+    console.print("\n[bold]done.[/bold] try `cognos doctor` to verify, then `cognos`.")
+    return 0