feat: persist episodes and journal to the lantern-episodes dataset

app.py

@@ -1,175 +1,173 @@
-"""Lantern web server: runs Beam episodes and streams them to the browser over SSE.
-
-Live runs spend the Space owner's inference credits, so they are limited to one at a
-time and a daily cap. Every finished episode is saved and can be replayed for free.
-"""
-
-from __future__ import annotations
-
-import asyncio
-import json
-import os
-import threading
-import time
-import uuid
-from datetime import date
-from pathlib import Path
-
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import FileResponse, StreamingResponse
-from fastapi.staticfiles import StaticFiles
-from pydantic import BaseModel, Field
-
-from lantern.agent import EpisodeConfig, run_episode
-from lantern.env import TASKS, DungeonEnv
-from lantern.journal import Journal
-from lantern.llm import ChatClient
-
-DATA_DIR = Path(os.environ.get("LANTERN_DATA", "data"))
-EPISODES_DIR = DATA_DIR / "episodes"
-DAILY_LIVE_RUNS = int(os.environ.get("LANTERN_DAILY_RUNS", "20"))
-WEB_DIR = Path(__file__).parent / "web"
-
-app = FastAPI(title="Lantern")
-journal = Journal(DATA_DIR / "journal.json")
-
-
-class Run:
-    """One live episode running in a background thread, buffering its events."""
-
-    def __init__(self, task_id: str, seed: int, use_lessons: bool):
-        self.id = uuid.uuid4().hex[:12]
-        self.task_id, self.seed, self.use_lessons = task_id, seed, use_lessons
-        self.events: list[dict] = []
-        self.done = False
-        self.stop = threading.Event()
-
-    def start(self) -> None:
-        threading.Thread(target=self._play, daemon=True).start()
-
-    def _play(self) -> None:
-        env = None
-        try:
-            env = DungeonEnv(TASKS[self.task_id], seed=self.seed)
-            lessons = journal.lessons(self.task_id)
-            config = EpisodeConfig(use_lessons=self.use_lessons)
-            for event in run_episode(env, ChatClient(), lessons, config, self.stop):
-                if event["type"] == "end" and event["new_lessons"]:
-                    event["journal"] = journal.add(self.task_id, event["new_lessons"])
-                self.events.append(event)
-        except Exception as exc:  # surface the failure to the viewer instead of hanging
-            self.events.append({"type": "error", "message": f"{type(exc).__name__}: {exc}"})
-        finally:
-            if env is not None:
-                env.close()
-            self.done = True
-            self._save()
-
-    def _save(self) -> None:
-        EPISODES_DIR.mkdir(parents=True, exist_ok=True)
-        end = next((e for e in reversed(self.events) if e["type"] == "end"), {})
-        record = {"id": self.id, "task": self.task_id, "seed": self.seed,
-                  "use_lessons": self.use_lessons, "finished_at": time.time(),
-                  "outcome": end.get("outcome", "error"), "steps": end.get("steps", 0),
-                  "events": self.events}
-        (EPISODES_DIR / f"{self.id}.json").write_text(json.dumps(record), "utf-8")
-
-
-runs: dict[str, Run] = {}
-_quota = {"day": date.today(), "used": 0}
-_lock = threading.Lock()
-
-
-class RunRequest(BaseModel):
-    task: str
-    seed: int = Field(default=0, ge=0, le=1_000_000)
-    use_lessons: bool = True
-
-
-@app.get("/api/tasks")
-def list_tasks() -> list[dict]:
-    return [{"id": t.id, "title": t.title, "goal": t.goal, "max_steps": t.max_steps,
-             "lessons": len(journal.lessons(t.id))} for t in TASKS.values()]
-
-
-@app.get("/api/status")
-def status() -> dict:
-    busy = any(not r.done for r in runs.values())
-    return {"busy": busy, "runs_left_today": max(0, DAILY_LIVE_RUNS - _quota["used"])}
-
-
-@app.post("/api/runs")
-def start_run(req: RunRequest) -> dict:
-    if req.task not in TASKS:
-        raise HTTPException(404, f"Unknown level '{req.task}'.")
-    with _lock:
-        if _quota["day"] != date.today():
-            _quota.update(day=date.today(), used=0)
-        if any(not r.done for r in runs.values()):
-            raise HTTPException(409, "Beam is already exploring. Watch that run or open a replay.")
-        if _quota["used"] >= DAILY_LIVE_RUNS:
-            raise HTTPException(429, "Today's live runs are used up. Replays are still available.")
-        _quota["used"] += 1
-        run = Run(req.task, req.seed, req.use_lessons)
-        runs[run.id] = run
-    run.start()
-    return {"id": run.id}
-
-
-@app.post("/api/runs/{run_id}/stop")
-def stop_run(run_id: str) -> dict:
-    run = runs.get(run_id)
-    if run is None:
-        raise HTTPException(404, "No such run.")
-    run.stop.set()
-    return {"stopping": True}
-
-
-@app.get("/api/runs/{run_id}/events")
-async def run_events(run_id: str) -> StreamingResponse:
-    run = runs.get(run_id)
-    if run is None:
-        raise HTTPException(404, "No such run.")
-
-    async def stream():
-        sent = 0
-        while True:
-            while sent < len(run.events):
-                yield f"data: {json.dumps(run.events[sent])}\n\n"
-                sent += 1
-            if run.done and sent >= len(run.events):
-                yield "event: close\ndata: {}\n\n"
-                return
-            await asyncio.sleep(0.25)
-
-    return StreamingResponse(stream(), media_type="text/event-stream",
-                             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
-
-
-@app.get("/api/episodes")
-def list_episodes() -> list[dict]:
-    if not EPISODES_DIR.exists():
-        return []
-    out = []
-    for path in sorted(EPISODES_DIR.glob("*.json"), key=lambda p: p.stat().st_mtime, reverse=True)[:50]:
-        rec = json.loads(path.read_text("utf-8"))
-        out.append({k: rec[k] for k in ("id", "task", "seed", "use_lessons", "outcome", "steps", "finished_at")})
-    return out
-
-
-@app.get("/api/episodes/{episode_id}")
-def get_episode(episode_id: str) -> dict:
-    if not episode_id.isalnum():
-        raise HTTPException(400, "Bad episode id.")
-    path = EPISODES_DIR / f"{episode_id}.json"
-    if not path.exists():
-        raise HTTPException(404, "No such episode.")
-    return json.loads(path.read_text("utf-8"))
-
-
-@app.get("/")
-def index() -> FileResponse:
-    return FileResponse(WEB_DIR / "index.html")
-
-
-app.mount("/static", StaticFiles(directory=WEB_DIR), name="static")
+"""Lantern web server: runs Beam episodes and streams them to the browser over SSE.
+
+Live runs spend the Space owner's inference credits, so they are limited to one at a
+time and a daily cap. Every finished episode is saved and can be replayed for free.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import threading
+import time
+import uuid
+from datetime import date
+from pathlib import Path
+
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import FileResponse, StreamingResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel, Field
+
+from lantern.agent import EpisodeConfig, run_episode
+from lantern.env import TASKS, DungeonEnv
+from lantern.journal import Journal
+from lantern.llm import ChatClient
+from lantern.store import EpisodeStore
+
+DATA_DIR = Path(os.environ.get("LANTERN_DATA", "data"))
+DAILY_LIVE_RUNS = int(os.environ.get("LANTERN_DAILY_RUNS", "20"))
+WEB_DIR = Path(__file__).parent / "web"
+
+app = FastAPI(title="Lantern")
+# Episodes and the journal are mirrored to a Hub dataset so they survive restarts.
+store = EpisodeStore(DATA_DIR, os.environ.get("LANTERN_DATASET"), os.environ.get("LANTERN_DATASET_TOKEN"))
+store.pull()
+journal = Journal(store.journal_path)
+
+
+class Run:
+    """One live episode running in a background thread, buffering its events."""
+
+    def __init__(self, task_id: str, seed: int, use_lessons: bool):
+        self.id = uuid.uuid4().hex[:12]
+        self.task_id, self.seed, self.use_lessons = task_id, seed, use_lessons
+        self.events: list[dict] = []
+        self.done = False
+        self.stop = threading.Event()
+
+    def start(self) -> None:
+        threading.Thread(target=self._play, daemon=True).start()
+
+    def _play(self) -> None:
+        env = None
+        try:
+            env = DungeonEnv(TASKS[self.task_id], seed=self.seed)
+            lessons = journal.lessons(self.task_id)
+            config = EpisodeConfig(use_lessons=self.use_lessons)
+            for event in run_episode(env, ChatClient(), lessons, config, self.stop):
+                if event["type"] == "end" and event["new_lessons"]:
+                    event["journal"] = journal.add(self.task_id, event["new_lessons"])
+                    store.push_journal()
+                self.events.append(event)
+        except Exception as exc:  # surface the failure to the viewer instead of hanging
+            self.events.append({"type": "error", "message": f"{type(exc).__name__}: {exc}"})
+        finally:
+            if env is not None:
+                env.close()
+            self.done = True
+            self._save()
+
+    def _save(self) -> None:
+        end = next((e for e in reversed(self.events) if e["type"] == "end"), {})
+        store.save({"id": self.id, "task": self.task_id, "seed": self.seed, "source": "space",
+                    "use_lessons": self.use_lessons, "finished_at": time.time(),
+                    "outcome": end.get("outcome", "error"), "steps": end.get("steps", 0),
+                    "events": self.events})
+
+
+runs: dict[str, Run] = {}
+_quota = {"day": date.today(), "used": 0}
+_lock = threading.Lock()
+
+
+class RunRequest(BaseModel):
+    task: str
+    seed: int = Field(default=0, ge=0, le=1_000_000)
+    use_lessons: bool = True
+
+
+@app.get("/api/tasks")
+def list_tasks() -> list[dict]:
+    return [{"id": t.id, "title": t.title, "goal": t.goal, "max_steps": t.max_steps,
+             "lessons": len(journal.lessons(t.id))} for t in TASKS.values()]
+
+
+@app.get("/api/status")
+def status() -> dict:
+    busy = any(not r.done for r in runs.values())
+    return {"busy": busy, "runs_left_today": max(0, DAILY_LIVE_RUNS - _quota["used"])}
+
+
+@app.post("/api/runs")
+def start_run(req: RunRequest) -> dict:
+    if req.task not in TASKS:
+        raise HTTPException(404, f"Unknown level '{req.task}'.")
+    with _lock:
+        if _quota["day"] != date.today():
+            _quota.update(day=date.today(), used=0)
+        if any(not r.done for r in runs.values()):
+            raise HTTPException(409, "Beam is already exploring. Watch that run or open a replay.")
+        if _quota["used"] >= DAILY_LIVE_RUNS:
+            raise HTTPException(429, "Today's live runs are used up. Replays are still available.")
+        _quota["used"] += 1
+        run = Run(req.task, req.seed, req.use_lessons)
+        runs[run.id] = run
+    run.start()
+    return {"id": run.id}
+
+
+@app.post("/api/runs/{run_id}/stop")
+def stop_run(run_id: str) -> dict:
+    run = runs.get(run_id)
+    if run is None:
+        raise HTTPException(404, "No such run.")
+    run.stop.set()
+    return {"stopping": True}
+
+
+@app.get("/api/runs/{run_id}/events")
+async def run_events(run_id: str) -> StreamingResponse:
+    run = runs.get(run_id)
+    if run is None:
+        raise HTTPException(404, "No such run.")
+
+    async def stream():
+        sent = 0
+        while True:
+            while sent < len(run.events):
+                yield f"data: {json.dumps(run.events[sent])}\n\n"
+                sent += 1
+            if run.done and sent >= len(run.events):
+                yield "event: close\ndata: {}\n\n"
+                return
+            await asyncio.sleep(0.25)
+
+    return StreamingResponse(stream(), media_type="text/event-stream",
+                             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
+
+
+@app.get("/api/episodes")
+def list_episodes() -> list[dict]:
+    keys = ("id", "task", "seed", "use_lessons", "outcome", "steps", "finished_at")
+    return [{k: rec.get(k) for k in keys} | {"source": rec.get("source", "space")}
+            for rec in store.episodes()[:50]]
+
+
+@app.get("/api/episodes/{episode_id}")
+def get_episode(episode_id: str) -> dict:
+    if not episode_id.isalnum():
+        raise HTTPException(400, "Bad episode id.")
+    record = store.load(episode_id)
+    if record is None:
+        raise HTTPException(404, "No such episode.")
+    return record
+
+
+@app.get("/")
+def index() -> FileResponse:
+    return FileResponse(WEB_DIR / "index.html")
+
+
+app.mount("/static", StaticFiles(directory=WEB_DIR), name="static")

lantern/agent.py

@@ -136,6 +136,8 @@ def run_episode(env: Env, client: ChatClient, lessons: list[str] | None = None,
 
     yield {"type": "start", "task": task.id, "title": task.title, "goal": task.goal,
            "actions": env.action_names, "lessons": list(lessons or []) if config.use_lessons else [],
+           # Exact teacher context, so trajectories can be replayed as training data.
+           "model": getattr(client, "model", None), "system": system, "tools": tools,
            "frame": frame.to_dict()}
 
     while outcome is None:

lantern/bench.py

@@ -0,0 +1,65 @@
+"""Benchmark series: play one level across seeds, optionally carrying the journal forward."""
+
+from __future__ import annotations
+
+import time
+import uuid
+
+from .agent import EpisodeConfig, run_episode
+from .env import TASKS, DungeonEnv
+from .journal import merge_lessons
+from .llm import ChatClient
+from .store import EpisodeStore
+
+SEED_BASE = 1000  # benchmark seeds stay clear of the low seeds visitors try in the Space
+
+
+def bench_seeds(n: int) -> list[int]:
+    return [SEED_BASE + i for i in range(n)]
+
+
+def play_series(task_id: str, seeds: list[int], use_lessons: bool, tag: str,
+                client: ChatClient, store: EpisodeStore, source: str = "bench") -> list[dict]:
+    """Play `seeds` in order on one level and save every episode; returns per-episode summaries.
+
+    Each series starts with an empty journal; with `use_lessons`, lessons written after
+    each episode are merged in before the next seed.
+    """
+    config = EpisodeConfig(use_lessons=use_lessons, reflect=use_lessons)
+    lessons: list[str] = []
+    summaries = []
+    for seed in seeds:
+        env = DungeonEnv(TASKS[task_id], seed=seed)
+        try:
+            events = list(run_episode(env, client, lessons, config))
+        except Exception as exc:  # one bad episode must not sink the series
+            events = [{"type": "error", "message": f"{type(exc).__name__}: {exc}"}]
+        finally:
+            env.close()
+        end = next((e for e in events if e["type"] == "end"), {})
+        lessons_in = len(lessons)
+        if use_lessons and end.get("new_lessons"):
+            lessons = merge_lessons(lessons, end["new_lessons"], 12)
+        record = {"id": uuid.uuid4().hex[:12], "task": task_id, "seed": seed, "source": source,
+                  "bench": tag, "model": client.model, "use_lessons": use_lessons,
+                  "finished_at": time.time(), "outcome": end.get("outcome", "error"),
+                  "steps": end.get("steps", 0), "lessons_in": lessons_in, "events": events}
+        store.save(record, wait=True)
+        summaries.append({k: record[k] for k in ("id", "task", "seed", "use_lessons", "outcome", "steps")}
+                         | {"usage": end.get("usage", {})})
+    return summaries
+
+
+def summary_table(results: list[dict], task_ids: list[str], conditions: list[str]) -> str:
+    """Plain-text results table: wins, average steps, and tokens per level and condition."""
+    lines = [f"{'level':14}{'journal':>8}{'won':>8}{'avg steps':>11}{'tokens':>12}"]
+    for t in task_ids:
+        for cond in conditions:
+            rows = [r for r in results if r["task"] == t and r["use_lessons"] == (cond == "on")]
+            if not rows:
+                continue
+            won = sum(r["outcome"] == "success" for r in rows)
+            steps = sum(r["steps"] for r in rows) / len(rows)
+            tok = sum(r["usage"].get("prompt_tokens", 0) + r["usage"].get("completion_tokens", 0) for r in rows)
+            lines.append(f"{t:14}{cond:>8}{won:>4}/{len(rows):<3}{steps:>11.1f}{tok:>12,}")
+    return "\n".join(lines)

lantern/export.py

@@ -0,0 +1,69 @@
+"""Turn recorded episodes into supervised fine-tuning rows for a student model.
+
+Each Beam call becomes one chat example: the system prompt and observation Beam saw, and
+the tool calls it made. Only calls that executed something are kept, and by default only
+episodes Beam won, so the student imitates successful play.
+
+    python -m lantern.export data/ train.jsonl
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+from .store import EpisodeStore
+
+
+def tool_calls_for(executed: list[dict]) -> list[dict]:
+    """Rebuild OpenAI-style tool calls from the agent's executed-tool records."""
+    calls = []
+    for i, e in enumerate(executed):
+        if e["tool"] == "act":
+            args = {"actions": e["planned"]}
+        elif e["tool"] == "write_notes":
+            args = {"notes": e["notes"]}
+        else:
+            continue
+        calls.append({"id": f"call_{i}", "type": "function",
+                      "function": {"name": e["tool"], "arguments": json.dumps(args)}})
+    return calls
+
+
+def episode_rows(record: dict) -> list[dict]:
+    """SFT rows for one episode, or [] if it lacks the teacher context (older recordings)."""
+    start = next((e for e in record["events"] if e["type"] == "start"), None)
+    if not start or "system" not in start:
+        return []
+    rows = []
+    for e in record["events"]:
+        if e["type"] != "turn" or not e["executed"]:
+            continue
+        rows.append({
+            "messages": [
+                {"role": "system", "content": start["system"]},
+                {"role": "user", "content": e["prompt"]},
+                {"role": "assistant", "content": "", "tool_calls": tool_calls_for(e["executed"])},
+            ],
+            "tools": start["tools"],
+            "task": record["task"], "seed": record.get("seed"), "episode": record["id"],
+        })
+    return rows
+
+
+def export(root: Path, out: Path, outcomes: tuple[str, ...] = ("success",)) -> int:
+    """Write SFT rows from every stored episode with a matching outcome; returns the row count."""
+    n = 0
+    with out.open("w", encoding="utf-8") as f:
+        for record in EpisodeStore(root).episodes():
+            if record.get("outcome") not in outcomes:
+                continue
+            for row in episode_rows(record):
+                f.write(json.dumps(row) + "\n")
+                n += 1
+    return n
+
+
+if __name__ == "__main__":
+    print(export(Path(sys.argv[1]), Path(sys.argv[2])), "rows")

lantern/llm.py

@@ -7,6 +7,7 @@ production) or the local `hf auth login` cache. It is never logged.
 from __future__ import annotations
 
 import os
+import random
 import time
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -47,27 +48,39 @@ class ChatClient:
     """
 
     def __init__(self, model: str = BEAM_MODEL, url: str = ROUTER_URL,
-                 token: str | None = None, timeout_s: float = 300.0, retries: int = 3):
+                 token: str | None = None, timeout_s: float = 300.0, retries: int = 6,
+                 max_tokens: int = 16384, extra_body: dict | None = None):
+        """`max_tokens` is the default reply budget; `extra_body` is merged into every
+        request (e.g. `chat_template_kwargs` for a self-hosted student model)."""
         self.model = model
         self.url = url
+        self.max_tokens = max_tokens
+        self.extra_body = extra_body or {}
         self._headers = {"Authorization": f"Bearer {token or _load_token()}"}
         self._http = httpx.Client(timeout=timeout_s)
         self._retries = retries
 
     def chat(self, messages: list[dict], tools: list[dict] | None = None,
-             temperature: float = 0.7, max_tokens: int = 16384) -> Reply:
+             temperature: float = 0.7, max_tokens: int | None = None) -> Reply:
         """Send one chat completion request, retrying transient failures."""
         body = {"model": self.model, "messages": messages, "temperature": temperature,
-                "top_p": 0.99, "max_tokens": max_tokens}
+                "top_p": 0.99, "max_tokens": min(max_tokens or self.max_tokens, self.max_tokens),
+                **self.extra_body}
         if tools:
             body["tools"] = tools
             body["tool_choice"] = "auto"
 
         for attempt in range(self._retries):
             start = time.monotonic()
-            resp = self._http.post(self.url, json=body, headers=self._headers)
+            try:
+                resp = self._http.post(self.url, json=body, headers=self._headers)
+            except httpx.TransportError:
+                if attempt == self._retries - 1:
+                    raise
+                time.sleep(2 ** attempt + random.random())
+                continue
             if resp.status_code in (429, 500, 502, 503, 504) and attempt < self._retries - 1:
-                time.sleep(2 ** attempt)
+                time.sleep(2 ** attempt + random.random())  # jitter so parallel runs spread out
                 continue
             resp.raise_for_status()
             data = resp.json()

lantern/store.py

@@ -0,0 +1,84 @@
+"""Episode storage: local JSON files mirrored to a Hugging Face dataset repo.
+
+Layout (same locally and on the Hub):
+    episodes/{task}/{episode_id}.json
+    journal.json
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import threading
+import time
+from pathlib import Path
+
+log = logging.getLogger("lantern.store")
+
+
+class EpisodeStore:
+    """Writes episodes to `root`, and, when `repo_id` is set, mirrors them to the Hub."""
+
+    def __init__(self, root: Path, repo_id: str | None = None, token: str | None = None):
+        self.root = Path(root)
+        self.repo_id = repo_id
+        self.token = token
+        (self.root / "episodes").mkdir(parents=True, exist_ok=True)
+
+    @property
+    def journal_path(self) -> Path:
+        return self.root / "journal.json"
+
+    def pull(self) -> None:
+        """Download existing episodes and the journal from the Hub, if configured."""
+        if not self.repo_id:
+            return
+        from huggingface_hub import snapshot_download
+        try:
+            snapshot_download(self.repo_id, repo_type="dataset", local_dir=self.root, token=self.token,
+                              allow_patterns=["episodes/**", "journal.json"])
+        except Exception as exc:  # the app still works from local files
+            log.warning("Could not pull episodes from %s: %s", self.repo_id, exc)
+
+    def save(self, record: dict, wait: bool = False) -> Path:
+        """Write one episode record locally, then upload it (in the background unless `wait`)."""
+        rel = Path("episodes") / record["task"] / f"{record['id']}.json"
+        path = self.root / rel
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(json.dumps(record), "utf-8")
+        self._upload([(path, rel.as_posix())], f"Add episode {record['id']} ({record['task']})", wait)
+        return path
+
+    def push_journal(self) -> None:
+        if self.journal_path.exists():
+            self._upload([(self.journal_path, "journal.json")], "Update journal")
+
+    def episodes(self) -> list[dict]:
+        """All stored episode records, newest first."""
+        records = [json.loads(p.read_text("utf-8")) for p in (self.root / "episodes").rglob("*.json")]
+        return sorted(records, key=lambda r: r.get("finished_at", 0), reverse=True)
+
+    def load(self, episode_id: str) -> dict | None:
+        match = next((self.root / "episodes").rglob(f"{episode_id}.json"), None)
+        return json.loads(match.read_text("utf-8")) if match else None
+
+    def _upload(self, files: list[tuple[Path, str]], message: str, wait: bool = False) -> None:
+        if not (self.repo_id and self.token):
+            return
+
+        def run() -> None:
+            from huggingface_hub import CommitOperationAdd, HfApi
+            for attempt in range(4):  # parallel writers can race on the repo head
+                try:
+                    HfApi(token=self.token).create_commit(
+                        self.repo_id, repo_type="dataset", commit_message=message,
+                        operations=[CommitOperationAdd(path_in_repo=rel, path_or_fileobj=str(p)) for p, rel in files])
+                    return
+                except Exception as exc:
+                    log.warning("Upload to %s failed (attempt %d): %s", self.repo_id, attempt + 1, exc)
+                    time.sleep(2 + attempt * 3)
+
+        if wait:
+            run()
+        else:
+            threading.Thread(target=run, daemon=True).start()

requirements.txt

@@ -4,3 +4,4 @@ httpx==0.28.1
 pydantic>=2.7,<3
 nle==1.3.0
 minihack==1.0.2
+huggingface_hub==0.36.2