litesquad 0.0.1__py3-none-any.whl

litesquad/__init__.py

@@ -0,0 +1,5 @@
+from .cli import app
+
+
+def main() -> None:
+    app()

litesquad/check_keys.py

@@ -0,0 +1,57 @@
+"""Dead-simple check that each provider API key is valid and active.
+
+Run:  uv run litesquad-keys      (or: uv run python -m litesquad.check_keys)
+
+Pings one cheap model per provider with a 1-token request and prints a
+checkmark per key. Independent of the squad config, so it tests the key
+itself, not whichever model your squad happens to use.
+"""
+
+import os
+
+import litellm
+from rich.console import Console
+
+from .llm import load_env  # loads .env with override=True; importing also quiets litellm
+
+console = Console()
+
+# One cheap, broadly-available probe model per provider key.
+PROBES = {
+    "ANTHROPIC_API_KEY": "anthropic/claude-haiku-4-5",
+    "OPENAI_API_KEY": "openai/gpt-4o-mini",
+    "GEMINI_API_KEY": "gemini/gemini-2.5-flash",
+    "OPENROUTER_API_KEY": "openrouter/openai/gpt-4o-mini",
+}
+
+
+def check(model: str) -> tuple[bool, str]:
+    try:
+        litellm.completion(
+            model=model,
+            messages=[{"role": "user", "content": "ping"}],
+            max_tokens=1,
+            drop_params=True,
+        )
+        return True, "ok"
+    except Exception as exc:  # noqa: BLE001 - any failure means the key isn't usable
+        return False, str(exc).splitlines()[0].strip()
+
+
+def main() -> None:
+    load_env()
+    for key, model in PROBES.items():
+        # An unset key is a distinct state from an invalid one: say "not set"
+        # instead of probing and printing a provider auth exception.
+        if not os.environ.get(key):
+            console.print(f"[red]✗[/] {key} — not set")
+            continue
+        ok, detail = check(model)
+        if ok:
+            console.print(f"[green]✓[/] {key} — active")
+        else:
+            console.print(f"[red]✗[/] {key} — invalid: {detail}")
+
+
+if __name__ == "__main__":
+    main()

litesquad/cli.py

@@ -0,0 +1,276 @@
+"""Command-line interface: one fixed-flow run, then interactive follow-ups."""
+
+from datetime import datetime
+from pathlib import Path
+
+import typer
+from rich.console import Console
+from rich.markdown import Markdown
+from rich.panel import Panel
+from rich.prompt import Prompt
+from rich.status import Status
+
+from . import paths
+from .config import RunConfig, SquadConfig, ensure_starter, load_config
+from .llm import LLMError, MissingKeysError, call_model, load_env, mock_call_model, preflight
+from .models import Conversation, Stage, TranscriptEvent
+from .squad import run_quick, run_turn
+
+app = typer.Typer(add_completion=False, help="Ask a small ensemble of LLMs anything.")
+console = Console()
+
+STAGE_LABEL: dict[Stage, str] = {
+    "propose": "responding",
+    "critique": "critiquing",
+    "revise": "revising",
+    "extract": "de-stylizing",
+    "cluster": "clustering",
+    "judge": "judging",
+    "reply": "answering",
+}
+QUIT_WORDS = {":quit", ":q", "quit", "exit"}
+SMOKE_PROMPT = "What is 1 + 1? Answer in one short sentence."
+
+
+class ConsoleReporter:
+    """Renders each stage as a Rich panel and appends events to a JSONL file.
+
+    Worker chains run in parallel, so their stages complete interleaved. Printing them
+    as they land is unreadable, so chain stages (propose/critique/revise) are buffered
+    per worker and printed as one coherent block the moment that worker's chain finishes
+    (or dies) -- blocks appear in chain-completion order. Everything else (extract,
+    cluster, judge, reply) prints as it completes, and the spinner always shows everyone
+    currently in flight. The transcript JSONL is still appended per event in completion
+    order: it is the durable stream, and the web UI regroups it by worker anyway.
+
+    Per the Reporter contract, calls arrive serialized, so no lock is needed here;
+    printing while the Status runs is safe (rich renders it above the spinner).
+    """
+
+    def __init__(self, transcript_path: Path | None) -> None:
+        self.transcript_path = transcript_path
+        self._status: Status | None = None
+        self._in_flight: dict[tuple[str, Stage], str] = {}  # (role, stage) -> spinner text
+        self._pending_chains: dict[str, list[Panel]] = {}  # worker key -> panels so far
+
+    @staticmethod
+    def _worker_key(role: str) -> str:
+        """propose/revise carry ``worker_N``; critique carries ``critic->worker_N``."""
+        return role.split("->")[-1]
+
+    @staticmethod
+    def _panel(event: TranscriptEvent) -> Panel:
+        title = f"{event.role} | {event.model}"
+        if event.error:
+            return Panel(event.error, title=f"{title} (error)", border_style="red")
+        border = "green" if event.stage in ("judge", "reply") else "cyan"
+        return Panel(Markdown(event.output), title=title, border_style=border)
+
+    def _spinner_text(self) -> str:
+        return "  |  ".join(self._in_flight.values())
+
+    def _flush_chain(self, worker: str) -> None:
+        for panel in self._pending_chains.pop(worker, []):
+            console.print(panel)
+
+    def stage_start(self, stage: Stage, role: str, model: str) -> None:
+        self._in_flight[(role, stage)] = f"[bold]{role}[/] ({model}) - {STAGE_LABEL[stage]}..."
+        if self._status is None:
+            self._status = console.status(self._spinner_text(), spinner="dots")
+            self._status.start()
+        else:
+            self._status.update(self._spinner_text())
+
+    def stage_done(self, event: TranscriptEvent) -> None:
+        self._in_flight.pop((event.role, event.stage), None)
+        if self.transcript_path is not None:
+            with self.transcript_path.open("a", encoding="utf-8") as fh:
+                fh.write(event.to_jsonl() + "\n")
+        if event.stage in ("propose", "critique", "revise"):
+            worker = self._worker_key(event.role)
+            self._pending_chains.setdefault(worker, []).append(self._panel(event))
+            if event.stage == "revise" or event.error:  # chain finished, or died here
+                self._flush_chain(worker)
+        else:
+            console.print(self._panel(event))
+        if self._status is not None:
+            if self._in_flight:
+                self._status.update(self._spinner_text())
+            else:
+                self._status.stop()
+                self._status = None
+
+    def close(self) -> None:
+        """Flush any unfinished chains and stop the spinner. Call in a ``finally``: an
+        aborted turn must still show what its surviving chains produced, and must not
+        leave a live display running (hidden cursor, endless spinner)."""
+        for worker in list(self._pending_chains):
+            self._flush_chain(worker)
+        if self._status is not None:
+            self._status.stop()
+            self._status = None
+        self._in_flight.clear()
+
+
+def _transcript_path(save: bool) -> Path | None:
+    if not save:
+        return None
+    stamp = datetime.now().strftime("%Y-%m-%d_%H%M%S")
+    return paths.transcripts_dir() / f"{stamp}.jsonl"
+
+
+def _model_roles(config: SquadConfig) -> list[tuple[str, str]]:
+    """Distinct configured models in pipeline order (workers first, then critic ->
+    extractor -> clusterer -> judge), each with the role(s) it holds -- so --check
+    doubles as a who-does-what roster card."""
+    assignments = [
+        *(("worker", worker.model) for worker in config.workers),
+        ("critic", config.critic.model),
+        ("extractor", config.extractor.model),
+        ("clusterer", config.clusterer.model),
+        ("judge", config.judge.model),
+    ]
+    roles_by_model: dict[str, list[str]] = {}
+    for role, model in assignments:
+        roles = roles_by_model.setdefault(model, [])
+        if role not in roles:  # two workers on one model still read as one "worker"
+            roles.append(role)
+    return [(model, ", ".join(roles)) for model, roles in roles_by_model.items()]
+
+
+def _check_models(config: SquadConfig) -> bool:
+    """Ping each distinct configured model with a tiny request. Returns all-ok.
+
+    ``max_tokens`` is generous (not 5) because reasoning models (GPT-5,
+    Gemini 2.5 Pro) spend output budget on hidden reasoning before any visible
+    text — too small a cap returns empty content. The cap only prevents
+    truncation; actual usage stays tiny since the model stops after "ok".
+    """
+    run_cfg = RunConfig(max_tokens=1024, save_transcript=False)
+    messages = [{"role": "user", "content": "Reply with the single word: ok"}]
+    all_ok = True
+    for model, roles in _model_roles(config):
+        try:
+            reply = call_model(model, messages, run_cfg, role="check").strip()
+            console.print(f"[green]✓[/] {model} [dim]({roles})[/] — {reply[:40]}")
+        except LLMError as exc:
+            all_ok = False
+            console.print(f"[red]✗[/] {model} [dim]({roles})[/] — {exc}")
+    return all_ok
+
+
+@app.command()
+def run(
+    task: str = typer.Argument(None, help="Your question or task for the ensemble."),
+    quick: bool = typer.Option(
+        False, "--quick", help="Talk to just the judge (Opus), skipping the ensemble."
+    ),
+    mock: bool = typer.Option(
+        False, "--mock", help="Use canned offline responses; no API keys needed."
+    ),
+    check: bool = typer.Option(
+        False, "--check", help="Ping each configured model with a tiny request and exit."
+    ),
+    smoke: bool = typer.Option(
+        False, "--smoke", help="Run one real turn on a fixed tiny prompt and exit (cheap end-to-end)."
+    ),
+    web: bool = typer.Option(
+        False, "--web", help="Serve the web UI instead of the terminal (needs the 'web' extra)."
+    ),
+    port: int = typer.Option(8050, help="Port for the web UI (only used with --web)."),
+) -> None:
+    """Ask the ensemble (or, with --quick, just the judge), then take follow-ups."""
+    if web:
+        if task is not None:
+            console.print("[red]With --web, submit tasks in the browser, not on the command line.[/]")
+            raise typer.Exit(2)
+        from .web import serve  # lazy: dash is an optional extra
+
+        serve(port=port, mock=mock)
+        return
+
+    load_env()
+
+    cfg_path = paths.config_path()
+    if ensure_starter(cfg_path):
+        console.print(
+            f"[dim]Wrote a starter config (all defaults, commented) you can edit at {cfg_path}[/]"
+        )
+    config = load_config(cfg_path)
+
+    if check:
+        try:
+            preflight(config)
+        except MissingKeysError as exc:
+            console.print(f"[red]{exc}[/]")
+            raise typer.Exit(1) from exc
+        console.print("Pinging configured models…")
+        raise typer.Exit(0 if _check_models(config) else 1)
+
+    if smoke:
+        caller = mock_call_model if mock else call_model
+        if not mock:
+            try:
+                preflight(config)
+            except MissingKeysError as exc:
+                console.print(f"[red]{exc}[/]")
+                raise typer.Exit(1) from exc
+        console.print(f'Smoke test: one turn on [dim]"{SMOKE_PROMPT}"[/]\n')
+        transcript_path = _transcript_path(config.run.save_transcript)
+        reporter = ConsoleReporter(transcript_path)
+        try:
+            run_turn(Conversation(), SMOKE_PROMPT, config, reporter, caller=caller)
+        except Exception as exc:  # noqa: BLE001 - smoke surfaces ANY failure (call or save)
+            console.print(f"[red]Smoke test FAILED: {exc}[/]")
+            raise typer.Exit(1) from exc
+        finally:
+            reporter.close()
+        if transcript_path is not None:
+            console.print(f"[dim]Transcript: {transcript_path}[/]")
+        console.print("[green]Smoke test passed - all stages produced output.[/]")
+        raise typer.Exit(0)
+
+    if task is None:
+        console.print('[red]Provide a task, e.g. litesquad "Plan my week", or use --check.[/]')
+        raise typer.Exit(2)
+
+    caller = call_model
+    if mock:
+        caller = mock_call_model
+        console.print("[yellow]Running in --mock mode: canned responses, no API calls.[/]")
+    else:
+        try:
+            preflight(config)
+        except MissingKeysError as exc:
+            console.print(f"[red]{exc}[/]")
+            raise typer.Exit(1) from exc
+
+    conversation = Conversation()
+    transcript_path = _transcript_path(config.run.save_transcript)
+    reporter = ConsoleReporter(transcript_path)
+    run_one = run_quick if quick else run_turn
+    if quick:
+        console.print("[dim]Quick mode: just the judge, no ensemble.[/]")
+
+    current_task = task
+    while True:
+        try:
+            run_one(conversation, current_task, config, reporter, caller=caller)
+        except LLMError as exc:
+            console.print(f"[red]Turn aborted: {exc}[/]")
+        finally:
+            reporter.close()
+
+        if transcript_path is not None:
+            console.print(f"[dim]Transcript: {transcript_path}[/]")
+
+        try:
+            reply = Prompt.ask(
+                "\n[bold magenta]Follow-up[/] ([bold yellow]:quit to exit[/])"
+            ).strip()
+        except (EOFError, KeyboardInterrupt):
+            console.print()
+            break
+        if not reply or reply.lower() in QUIT_WORDS:
+            break
+        current_task = reply

litesquad/config.py

@@ -0,0 +1,192 @@
+"""Typed configuration: good defaults in code, user overrides from ~/.litesquad.
+
+The defaults below are authoritative and versioned. On load, an optional user
+config at :func:`litesquad.paths.config_path` is shallow-merged on top, so the
+user's file only needs the handful of things they want to change.
+"""
+
+import tomllib
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+
+from . import paths
+
+DEFAULT_CONFIG_TOML = """\
+# litesquad default squad. Models are LiteLLM model strings.
+
+[run]
+# Caps output per stage. Reasoning models (GPT-5, Gemini 2.5 Pro) spend part of
+# this on hidden reasoning before the visible answer, so keep it roomy.
+max_tokens = 8000
+save_transcript = true
+# Randomize the order the judge sees the responses so no worker is permanently
+# first (LLM judges have a primacy bias). Turn off for deterministic-order tests.
+shuffle = true
+# How many worker chains run at once (threads; model calls are I/O-bound).
+# 1 runs everything serially, which is useful when debugging. Higher values
+# speed up wide rosters but hit providers harder: the critic model receives
+# roughly this many concurrent requests.
+max_parallel = 4
+# temperature is omitted: frontier models (Opus 4.7+, GPT-5) reject it with a 400.
+# Set it only if every model in your squad supports it (Sonnet, Gemini, Opus 4.6-).
+# temperature = 0.4
+
+# The judge writes the final answer from the clustered content map (extract -> cluster -> judge).
+[agents.judge]
+model = "anthropic/claude-opus-4-8"
+
+[agents.critic]
+model = "openai/gpt-5"
+
+# The extractor de-stylizes each revised response into content units (JSON). A mechanical step,
+# so a cheap model is fine -- swap to openai/gpt-4.1-mini to cut cost.
+[agents.extractor]
+model = "openai/gpt-5"
+
+# The clusterer groups equivalent units across responses and flags conflicts, building the
+# content map the judge writes from. It makes no quality judgment (that is the judge's job).
+# Its load scales with roster width (every worker's units land in one pool), and a reasoning
+# model in this seat spends hidden budget before emitting the JSON, so it gets extra headroom.
+[agents.clusterer]
+model = "anthropic/claude-opus-4-8"
+max_tokens = 24000
+
+# Each worker responds independently (blind to the others), the critic gives each one feedback,
+# and the worker revises. Then the revised answers are extracted into units, clustered into a
+# content map, and the judge writes the final answer from it. Any agent may add an
+# `instructions` string that is appended to its system prompt, and a `max_tokens` that
+# overrides [run] max_tokens for its own calls -- e.g. a reasoning model clustering a wide
+# roster's units can need max_tokens = 24000 while the workers stay at the run default.
+[[agents.workers]]
+model = "anthropic/claude-sonnet-4-6"
+
+[[agents.workers]]
+model = "gemini/gemini-2.5-pro"
+
+# openrouter/* workers: one OPENROUTER_API_KEY reaches every provider on
+# openrouter.ai (deepseek, mistral, llama, grok, qwen, ...), which is how the
+# roster grows wide without needing a key per provider. Prefer non-reasoning
+# models as workers here: through OpenRouter a reasoning model can spend the
+# whole max_tokens budget on hidden reasoning and return empty content. If you
+# do add one, tell it to answer tersely -- instructions cap the visible answer
+# reliably, the hidden reasoning only partly.
+[[agents.workers]]
+model = "openrouter/deepseek/deepseek-chat"
+
+[[agents.workers]]
+model = "openrouter/mistralai/mistral-large"
+
+[[agents.workers]]
+model = "openrouter/meta-llama/llama-3.3-70b-instruct"
+"""
+
+_STARTER_HEADER = """\
+# litesquad overrides. This file mirrors the built-in defaults, fully commented.
+# Uncomment and edit a line to override that default; anything left commented
+# keeps following the library default, so you only carry the deltas you care about.
+
+"""
+
+
+class AgentConfig(BaseModel):
+    model: str
+    instructions: str | None = None  # optional, appended to this agent's system prompt
+    max_tokens: int | None = None  # optional, overrides run.max_tokens for this agent's calls
+
+
+class RunConfig(BaseModel):
+    temperature: float | None = None
+    max_tokens: int = 8000
+    save_transcript: bool = True
+    shuffle: bool = True  # randomize the order the judge sees responses (kills primacy bias)
+    max_parallel: int = 4  # concurrent worker chains (1 = serial); keep equal to the TOML default
+
+
+class SquadConfig(BaseModel):
+    run: RunConfig = Field(default_factory=RunConfig)
+    judge: AgentConfig
+    critic: AgentConfig
+    extractor: AgentConfig
+    clusterer: AgentConfig
+    workers: list[AgentConfig] = Field(min_length=1)
+
+    def models(self) -> list[str]:
+        """Every distinct model string referenced by the squad."""
+        seen: list[str] = []
+        roles = [self.judge.model, self.critic.model, self.extractor.model, self.clusterer.model]
+        for model in [*roles, *(w.model for w in self.workers)]:
+            if model not in seen:
+                seen.append(model)
+        return seen
+
+
+def _build(raw: dict) -> SquadConfig:
+    """Build a SquadConfig from a raw TOML dict (flattening [agents.*])."""
+    agents = raw.get("agents", {})
+    return SquadConfig(
+        run=raw.get("run", {}),
+        judge=agents.get("judge", {}),
+        critic=agents.get("critic", {}),
+        extractor=agents.get("extractor", {}),
+        clusterer=agents.get("clusterer", {}),
+        workers=agents.get("workers", []),
+    )
+
+
+def default_config() -> SquadConfig:
+    """The authoritative, versioned default squad."""
+    return _build(tomllib.loads(DEFAULT_CONFIG_TOML))
+
+
+def _merge(defaults: dict, overrides: dict) -> dict:
+    """Shallow-merge user overrides onto the default raw config.
+
+    ``run`` keys merge individually; a provided single agent (judge/critic/extractor/clusterer)
+    replaces that agent; a provided (non-empty) workers list replaces the default workers.
+    Nothing deeper.
+    """
+    d_agents = defaults.get("agents", {})
+    o_agents = overrides.get("agents", {})
+    return {
+        "run": {**defaults.get("run", {}), **overrides.get("run", {})},
+        "agents": {
+            "judge": o_agents.get("judge") or d_agents.get("judge", {}),
+            "critic": o_agents.get("critic") or d_agents.get("critic", {}),
+            "extractor": o_agents.get("extractor") or d_agents.get("extractor", {}),
+            "clusterer": o_agents.get("clusterer") or d_agents.get("clusterer", {}),
+            "workers": o_agents.get("workers") or d_agents.get("workers", []),
+        },
+    }
+
+
+def load_config(path: Path | None = None) -> SquadConfig:
+    """Load the default squad, with any user overrides merged on top."""
+    defaults = tomllib.loads(DEFAULT_CONFIG_TOML)
+    cfg_path = path or paths.config_path()
+    if cfg_path.exists():
+        overrides = tomllib.loads(cfg_path.read_text(encoding="utf-8"))
+        return _build(_merge(defaults, overrides))
+    return _build(defaults)
+
+
+def _commented_starter() -> str:
+    """The defaults rendered as an all-commented override file."""
+    lines = []
+    for line in DEFAULT_CONFIG_TOML.splitlines():
+        stripped = line.lstrip()
+        if stripped == "" or stripped.startswith("#"):
+            lines.append(line)
+        else:
+            lines.append(f"# {line}")
+    return _STARTER_HEADER + "\n".join(lines) + "\n"
+
+
+def ensure_starter(path: Path | None = None) -> bool:
+    """Write the commented starter override if none exists. True if it was written."""
+    cfg_path = path or paths.config_path()
+    if cfg_path.exists():
+        return False
+    cfg_path.parent.mkdir(parents=True, exist_ok=True)
+    cfg_path.write_text(_commented_starter(), encoding="utf-8")
+    return True