From 103a899f897d64472491000299b21e8e767477e5 Mon Sep 17 00:00:00 2001 From: Juniper Bevensee Date: Mon, 22 Jun 2026 16:07:55 +1200 Subject: [PATCH 1/2] feat(pipeline): stateful, resumable study store + runner (+ bibliography study) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Long analyses currently run inline in one agent turn — they die on OOM and lose all state on container rebuild. This adds durable, checkpointed analysis state so a study resumes from the last completed step instead of vanishing. - engine/store.py: StudyStore (stdlib sqlite3, WAL, atomic). studies/steps/ artifacts/findings tables; create/get/list studies, add_steps, step status + result, artifacts, findings, study_summary. Injected db path (MATILDE_STUDY_DB / HERMES_HOME / repo engagements dir). State survives a fresh StudyStore instance on the same file (the container-rebuild guarantee, under test). - engine/pipeline.py: resumable runner. Step/StepContext/StepResult; run()/resume() skip done steps, persist result+artifacts+findings per step, stop on failure leaving the study blocked + resumable. Pure/injected. - engine/bibliography_study.py: first concrete study (parse_refs -> verify_each -> summarize) reusing the citations/parsing engine, verifier injected. verify_each is mid-step resumable (skips refs that already have findings). - tools.py / __init__.py: matilde_study_create/run/status/list, matching the existing tool-dict + thin-handler conventions. Tests (TDD): store CRUD + status transitions + persistence across a new instance; runner skip-done / stop-on-failure / resume-without-rerun; bibliography happy path and fail-on-ref-#3-then-resume (refs 1-2 NOT re-verified); tool arg coercion + output shape. Full suite: 123 passed, 6 skipped. Engine stays stdlib-only (no mne/numpy) — the heavy MEG data-analysis study is a documented follow-up. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/stateful-study-pipeline.md | 97 ++++++ matilde_plugin/__init__.py | 12 + matilde_plugin/engine/bibliography_study.py | 130 ++++++++ matilde_plugin/engine/pipeline.py | 119 +++++++ matilde_plugin/engine/store.py | 326 ++++++++++++++++++++ matilde_plugin/tools.py | 212 +++++++++++++ tests/test_bibliography_study.py | 104 +++++++ tests/test_pipeline.py | 140 +++++++++ tests/test_plugin_tools.py | 4 + tests/test_store.py | 179 +++++++++++ tests/test_study_tools.py | 103 +++++++ 11 files changed, 1426 insertions(+) create mode 100644 docs/stateful-study-pipeline.md create mode 100644 matilde_plugin/engine/bibliography_study.py create mode 100644 matilde_plugin/engine/pipeline.py create mode 100644 matilde_plugin/engine/store.py create mode 100644 tests/test_bibliography_study.py create mode 100644 tests/test_pipeline.py create mode 100644 tests/test_store.py create mode 100644 tests/test_study_tools.py diff --git a/docs/stateful-study-pipeline.md b/docs/stateful-study-pipeline.md new file mode 100644 index 0000000..e7ced4e --- /dev/null +++ b/docs/stateful-study-pipeline.md @@ -0,0 +1,97 @@ +# Stateful Study Pipeline (design) + +## Motivation + +Long analyses — e.g. validating a paper's findings against an OpenNeuro dataset — +currently run **inline in a single agent turn**: one-shot scripts that + +- **die on OOM** (observed 2026-06-22: a live MEG M100 validation full-preloaded a + 2.5 GB dataset into the agent container and was OOM-killed), and +- **lose all state when the container restarts** (the same session's analysis was + wiped when the agent image was rebuilt mid-run — no resume, no memory of it). + +Both failure modes have the same remedy: **durable, checkpointed analysis state** +so an analysis resumes from the last completed step after an OOM *or* a container +rebuild, instead of vanishing. + +## Hard requirements + +1. **Survives container rebuild** — study state lives on the mounted data volume + (`HERMES_HOME`/Matilde engagements dir), not in process memory. +2. **Survives OOM / mid-step failure** — a failed step leaves the study resumable; + completed steps are never re-run. +3. **Stdlib-only core** — mirrors the rest of `matilde_plugin/engine` (pure, + injected-I/O, SQLite from stdlib). Heavy scientific deps (mne, etc.) live only + in optional step implementations, never in the framework. +4. **Agent-drivable** — the agent kicks off / advances / inspects a study through + tools, so it never has to hold a long analysis in a single turn. +5. **TDD** — injected I/O, deterministic, matches the existing ~99-test suite. + +## Components + +### `matilde_plugin/engine/store.py` — `StudyStore` (SQLite) + +Persistent store; injected db path (default under the engagements dir). WAL, +atomic writes. Tables: + +- `studies(id, slug UNIQUE, title, status, plan_json, created_at, updated_at, meta_json)` + - status: `created | running | done | failed | blocked` +- `steps(id, study_id, idx, name, status, result_json, error, started_at, finished_at)` + - status: `pending | running | done | failed | skipped` + - UNIQUE(study_id, name) +- `artifacts(id, study_id, step_name, path, kind, sha256, bytes, meta_json, created_at)` + - references on-disk files (downloads, intermediate `.npy`/`.fif`, plots) — **never** blobs in the DB +- `findings(id, study_id, step_name, claim, verdict, score, evidence_json, created_at)` + - the scientific output: claim + verdict (`supported|refuted|inconclusive`) + evidence + +API (all idempotent / upsert where sensible): `create_study`, `get_study`, +`list_studies`, `add_steps(plan)`, `set_step_status`, `record_step_result`, +`add_artifact`, `add_finding`, `get_findings`, `study_summary`. + +### `matilde_plugin/engine/pipeline.py` — resumable runner + +- `Step = (name: str, fn: Callable[[StepContext], StepResult])` +- `StepContext` carries the store, study_id, results of prior steps, and injected + I/O handles (so steps are testable without network/fs side effects). +- `StepResult` = `{data: dict, artifacts: [...], findings: [...]}`. +- `run(store, study_id, steps, *, resume=True)`: + - for each step in order: if its stored status is `done`, **skip** (resume); + else mark `running`, call `fn(ctx)`, persist result + artifacts + findings, + mark `done`; on exception mark `failed`, record error, **stop** (study + `blocked`/`failed`, resumable later). Returns a per-step summary. +- `resume(store, study_id, steps)` = `run(..., resume=True)` — continues from the + first non-`done` step. Re-running a fully-done study is a no-op. +- Idempotency contract: step fns must be safe to re-run from scratch (the runner + only guarantees done-steps are skipped, not partial-step rollback). + +### Agent tools (`matilde_plugin/tools.py`) + +Match the existing tool-dict + handler conventions in this file: + +- `matilde_study_create(slug, title, plan)` → `{study_id}` (plan = ordered step names + params) +- `matilde_study_run(study_id)` → advances pending steps; returns status + per-step summary (so an OOM/restart mid-run is resumed by simply calling this again) +- `matilde_study_status(study_id)` → full status + findings + artifacts +- `matilde_study_list()` → recent studies + +### First concrete pipeline (this PR — lightweight, proves the framework, no heavy deps) + +**Bibliography-validation study**: steps `parse_refs → verify_each → summarize`, +reusing the existing citations engine. `verify_each` records one finding per +reference. Demonstrated resumability: run with the verifier failing on ref #3, +assert the study is `blocked` with refs 1–2 `done`; fix the injected verifier; +`resume`; assert it completes from ref #3 without re-verifying 1–2. + +### Follow-up (separate PR — heavy data analysis) + +**Data-validation study** (the MEG M100 case): steps `fetch_dataset → +preprocess → epoch → evoked → validate_finding`, each **memory-bounded** +(chunked/streamed, intermediates checkpointed as artifacts) so an OOM resumes +from the last completed step. Requires deciding the agent container memory budget +and where heavy deps (mne) are installed. Out of scope for the framework PR. + +## Testing + +stdlib + injected I/O. Cover: store CRUD + status transitions + persistence +across a fresh `StudyStore` instance (simulates restart); runner skip-done / +stop-on-failure / resume; the bibliography study end-to-end incl. the +fail-then-resume scenario; the agent tools' arg coercion + outputs. diff --git a/matilde_plugin/__init__.py b/matilde_plugin/__init__.py index ecd3f98..540b0e3 100644 --- a/matilde_plugin/__init__.py +++ b/matilde_plugin/__init__.py @@ -55,6 +55,10 @@ OPENNEURO_SEARCH_SCHEMA = _tools_mod.OPENNEURO_SEARCH_SCHEMA OPENNEURO_FILES_SCHEMA = _tools_mod.OPENNEURO_FILES_SCHEMA FETCH_FULLTEXT_SCHEMA = _tools_mod.FETCH_FULLTEXT_SCHEMA +STUDY_CREATE_SCHEMA = _tools_mod.STUDY_CREATE_SCHEMA +STUDY_RUN_SCHEMA = _tools_mod.STUDY_RUN_SCHEMA +STUDY_STATUS_SCHEMA = _tools_mod.STUDY_STATUS_SCHEMA +STUDY_LIST_SCHEMA = _tools_mod.STUDY_LIST_SCHEMA _check_available = _tools_mod._check_available _handle_verify_citation = _tools_mod._handle_verify_citation _handle_verify_bibliography = _tools_mod._handle_verify_bibliography @@ -63,6 +67,10 @@ _handle_openneuro_search = _tools_mod._handle_openneuro_search _handle_openneuro_list_files = _tools_mod._handle_openneuro_list_files _handle_fetch_fulltext = _tools_mod._handle_fetch_fulltext +_handle_study_create = _tools_mod._handle_study_create +_handle_study_run = _tools_mod._handle_study_run +_handle_study_status = _tools_mod._handle_study_status +_handle_study_list = _tools_mod._handle_study_list # --------------------------------------------------------------------------- # Tool registry — (tool_name, schema, handler, emoji) @@ -75,6 +83,10 @@ ("matilde_openneuro_search", OPENNEURO_SEARCH_SCHEMA, _handle_openneuro_search, "🔎"), ("matilde_openneuro_list_files", OPENNEURO_FILES_SCHEMA, _handle_openneuro_list_files, "🗂"), ("matilde_fetch_fulltext", FETCH_FULLTEXT_SCHEMA, _handle_fetch_fulltext, "📄"), + ("matilde_study_create", STUDY_CREATE_SCHEMA, _handle_study_create, "🧪"), + ("matilde_study_run", STUDY_RUN_SCHEMA, _handle_study_run, "▶"), + ("matilde_study_status", STUDY_STATUS_SCHEMA, _handle_study_status, "📊"), + ("matilde_study_list", STUDY_LIST_SCHEMA, _handle_study_list, "🗒"), ) diff --git a/matilde_plugin/engine/bibliography_study.py b/matilde_plugin/engine/bibliography_study.py new file mode 100644 index 0000000..f157e7b --- /dev/null +++ b/matilde_plugin/engine/bibliography_study.py @@ -0,0 +1,130 @@ +"""Bibliography-validation study — the first concrete pipeline. + +Three steps, reusing the existing parsing + citations engine: + + 1. ``parse_refs`` — parse a BibTeX blob into ``Reference`` objects. + 2. ``verify_each`` — verify each reference, recording one *finding* per ref. + 3. ``summarize`` — roll the findings up into a per-verdict tally. + +This is deliberately lightweight (no heavy scientific deps) and exists to *prove +the framework*: it is the worked example for resumability. ``verify_each`` is +written to be safe to re-run from scratch — on resume it skips references that +already have a persisted finding — so when the verifier fails partway through, a +later ``resume`` continues from the unfinished reference without re-verifying the +ones that already passed. (The runner only skips done *steps*; mid-step +idempotency is the step's own responsibility — see ``pipeline`` docstring.) + +The verifier is **injected**: ``verifier(ref) -> {claim, verdict, score, +evidence}``. Production passes :func:`default_verifier` (which wraps the real +citations engine); tests pass a fake that can force a failure on a chosen ref. +""" +from __future__ import annotations + +from typing import Callable, List, Optional + +from .citations import Reference, verify_reference +from .parsing import parse_bibtex +from .pipeline import Step, StepContext, StepResult + +VerifierFn = Callable[[Reference], dict] + +# Map the citation engine's verdicts onto the finding vocabulary +# (supported | refuted | inconclusive) the study records. +_VERDICT_MAP = { + "verified": "supported", + "warnings": "inconclusive", + "unverifiable": "inconclusive", + "not_found": "refuted", + "retracted": "refuted", +} + + +def _ref_key(ref: Reference) -> str: + """A stable identifier for a reference (for resume de-duplication + claims).""" + return ref.doi or ref.title or ref.raw[:80] or "(unidentified)" + + +def default_verifier(ref: Reference) -> dict: + """Verify *ref* with the real citations engine (network). Injected in prod.""" + result = verify_reference(ref) + verdict = _VERDICT_MAP.get(result.verdict, "inconclusive") + return { + "claim": f"reference is valid: {_ref_key(ref)}", + "verdict": verdict, + "score": result.score, + "evidence": { + "citation_verdict": result.verdict, + "axes": { + "existence": result.existence.status, + "metadata_match": result.metadata_match.status, + "retraction": result.retraction.status, + "url_liveness": result.url_liveness.status, + }, + }, + } + + +def _parse_refs_step(bibtex: str) -> Step: + def fn(ctx: StepContext) -> StepResult: + refs = parse_bibtex(bibtex) + return StepResult(data={ + "count": len(refs), + "refs": [r.raw or _ref_key(r) for r in refs], + }) + return Step(name="parse_refs", fn=fn) + + +def _verify_each_step(bibtex: str, verifier: VerifierFn) -> Step: + def fn(ctx: StepContext) -> StepResult: + refs = parse_bibtex(bibtex) + # Resume-safe: skip refs that already have a persisted finding. + already = {f["evidence"].get("ref_key") + for f in ctx.store.get_findings(ctx.study_id) + if f["step_name"] == "verify_each"} + for ref in refs: + key = _ref_key(ref) + if key in already: + continue + v = verifier(ref) # may raise -> step fails, study resumable + evidence = dict(v.get("evidence") or {}) + evidence["ref_key"] = key + # Persist immediately so a later failure doesn't lose this ref. + ctx.store.add_finding( + ctx.study_id, "verify_each", + claim=v.get("claim", f"reference: {key}"), + verdict=v.get("verdict", "inconclusive"), + score=v.get("score"), + evidence=evidence) + return StepResult(data={"verified": len(refs)}) + return Step(name="verify_each", fn=fn) + + +def _summarize_step() -> Step: + def fn(ctx: StepContext) -> StepResult: + findings = [f for f in ctx.store.get_findings(ctx.study_id) + if f["step_name"] == "verify_each"] + counts: dict = {} + for f in findings: + counts[f["verdict"]] = counts.get(f["verdict"], 0) + 1 + return StepResult(data={ + "total": len(findings), + "supported": counts.get("supported", 0), + "refuted": counts.get("refuted", 0), + "inconclusive": counts.get("inconclusive", 0), + "by_verdict": counts, + }) + return Step(name="summarize", fn=fn) + + +def build_steps(bibtex: str, verifier: Optional[VerifierFn] = None) -> List[Step]: + """Build the ``parse_refs -> verify_each -> summarize`` steps for *bibtex*. + + Pass *verifier* to inject the per-reference check (defaults to the real + citations engine). The returned steps are fed to ``pipeline.run`` / ``resume``. + """ + verifier = verifier or default_verifier + return [ + _parse_refs_step(bibtex), + _verify_each_step(bibtex, verifier), + _summarize_step(), + ] diff --git a/matilde_plugin/engine/pipeline.py b/matilde_plugin/engine/pipeline.py new file mode 100644 index 0000000..9021d83 --- /dev/null +++ b/matilde_plugin/engine/pipeline.py @@ -0,0 +1,119 @@ +"""Resumable study runner — checkpoints each step into a ``StudyStore``. + +A study is an ordered list of :class:`Step`\\s. The runner walks them in order: +a step whose stored status is already ``done`` is **skipped** (resume); otherwise +it is marked ``running``, its ``fn`` is called, its result + artifacts + findings +are persisted, and it is marked ``done``. On exception the step is marked +``failed`` (with the error), the study is marked ``blocked``, and the run **stops** +— leaving everything before it ``done`` and the study resumable later. This is the +durability contract: an OOM or container rebuild mid-run is recovered by simply +calling ``run``/``resume`` again on the same store. + +Pure and injected: a :class:`Step`'s ``fn`` receives a :class:`StepContext` +carrying the store, the study id, and the results of prior steps, so steps are +unit-testable with in-memory fakes (no network/fs needed). + +Idempotency contract: the runner only guarantees *done steps are skipped*, not +partial-step rollback. A step that does internal sub-work (e.g. verifying N refs) +must itself be safe to re-run from scratch — typically by checking the store for +what it already persisted and skipping that (see ``bibliography_study``). +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional + +from .store import StudyStore + + +@dataclass +class StepResult: + """What a step produces. All optional — a step may only record findings.""" + + data: dict = field(default_factory=dict) + artifacts: List[dict] = field(default_factory=list) + findings: List[dict] = field(default_factory=list) + + +@dataclass +class StepContext: + """Handle passed to a step's ``fn``: the store, the study, and prior results.""" + + store: StudyStore + study_id: int + step_name: str + results: Dict[str, Any] = field(default_factory=dict) + + +@dataclass +class Step: + name: str + fn: Callable[[StepContext], StepResult] + + +def run(store: StudyStore, study_id: int, steps: List[Step], *, + resume: bool = True) -> dict: + """Execute *steps* in order against *study_id*, persisting after each. + + Returns a summary dict: ``{study_id, status, steps: [...], failed_step}``. + Done steps are skipped when ``resume`` is True (the default). + """ + store.add_steps(study_id, [s.name for s in steps]) + + # Carry forward results of steps already done (so a resumed step can read them). + prior: Dict[str, Any] = {} + for s in store.get_steps(study_id): + if s["status"] == "done" and s["result"] is not None: + prior[s["name"]] = s["result"] + + store.set_study_status(study_id, "running") + per_step: List[dict] = [] + failed_step: Optional[str] = None + + for step in steps: + stored = store.get_step(study_id, step.name) + if resume and stored is not None and stored["status"] == "done": + per_step.append({"name": step.name, "status": "done", "skipped": True}) + continue + + store.set_step_status(study_id, step.name, "running") + ctx = StepContext(store=store, study_id=study_id, + step_name=step.name, results=dict(prior)) + try: + result = step.fn(ctx) + except Exception as exc: # stop, leave resumable + store.set_step_status(study_id, step.name, "failed", + error=f"{type(exc).__name__}: {exc}") + store.set_study_status(study_id, "blocked") + per_step.append({"name": step.name, "status": "failed", + "error": f"{type(exc).__name__}: {exc}"}) + failed_step = step.name + return {"study_id": study_id, "status": "blocked", + "steps": per_step, "failed_step": failed_step} + + result = result or StepResult() + store.record_step_result(study_id, step.name, result.data) + for art in result.artifacts: + store.add_artifact( + study_id, step.name, + path=art.get("path", ""), kind=art.get("kind", ""), + sha256=art.get("sha256", ""), bytes=art.get("bytes"), + meta=art.get("meta")) + for fnd in result.findings: + store.add_finding( + study_id, step.name, + claim=fnd.get("claim", ""), verdict=fnd.get("verdict", ""), + score=fnd.get("score"), evidence=fnd.get("evidence")) + store.set_step_status(study_id, step.name, "done") + prior[step.name] = result.data + per_step.append({"name": step.name, "status": "done"}) + + store.set_study_status(study_id, "done") + return {"study_id": study_id, "status": "done", + "steps": per_step, "failed_step": None} + + +def resume(store: StudyStore, study_id: int, steps: List[Step]) -> dict: + """Continue a study from its first non-``done`` step. A fully-done study is a + no-op. Equivalent to ``run(..., resume=True)``.""" + return run(store, study_id, steps, resume=True) diff --git a/matilde_plugin/engine/store.py b/matilde_plugin/engine/store.py new file mode 100644 index 0000000..d85864a --- /dev/null +++ b/matilde_plugin/engine/store.py @@ -0,0 +1,326 @@ +"""StudyStore — durable, checkpointed analysis state for Matilde studies. + +A *study* is a long-running analysis (e.g. validating a paper against a dataset) +decomposed into ordered *steps*. Each step can emit *artifacts* (references to +on-disk files — never blobs in the DB) and *findings* (the scientific output: +claim + verdict + evidence). All of it is persisted to a SQLite file so the work + + - **survives an OOM / mid-step failure** — a failed step leaves the study + resumable; completed steps are never lost, and + - **survives a container rebuild** — the db lives on the mounted data volume, + so a brand-new ``StudyStore`` pointing at the same file sees all prior state. + +Design mirrors the rest of ``matilde_plugin/engine``: stdlib-only (``sqlite3``), +pure data in / data out, injected db path. WAL mode + autocommit make each write +atomic and crash-safe. + +The db path is injected. ``default_db_path()`` puts it under the engagements data +dir (``MATILDE_STUDY_DB`` override, else ``$HERMES_HOME``/engagements, else the +repo's ``engagements/`` dir) — the same on-volume location the rest of Matilde's +private operational data uses. +""" +from __future__ import annotations + +import json +import os +import sqlite3 +import time +from typing import List, Optional + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS studies ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + slug TEXT UNIQUE NOT NULL, + title TEXT NOT NULL DEFAULT '', + status TEXT NOT NULL DEFAULT 'created', + plan_json TEXT NOT NULL DEFAULT '[]', + meta_json TEXT NOT NULL DEFAULT '{}', + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL +); + +CREATE TABLE IF NOT EXISTS steps ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + study_id INTEGER NOT NULL REFERENCES studies(id), + idx INTEGER NOT NULL, + name TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', + result_json TEXT, + error TEXT, + started_at TEXT, + finished_at TEXT, + UNIQUE(study_id, name) +); + +CREATE TABLE IF NOT EXISTS artifacts ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + study_id INTEGER NOT NULL REFERENCES studies(id), + step_name TEXT NOT NULL DEFAULT '', + path TEXT NOT NULL, + kind TEXT NOT NULL DEFAULT '', + sha256 TEXT NOT NULL DEFAULT '', + bytes INTEGER, + meta_json TEXT NOT NULL DEFAULT '{}', + created_at TEXT NOT NULL +); + +CREATE TABLE IF NOT EXISTS findings ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + study_id INTEGER NOT NULL REFERENCES studies(id), + step_name TEXT NOT NULL DEFAULT '', + claim TEXT NOT NULL DEFAULT '', + verdict TEXT NOT NULL DEFAULT '', + score REAL, + evidence_json TEXT NOT NULL DEFAULT '{}', + created_at TEXT NOT NULL +); +""" + + +def _now() -> str: + return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + + +def default_db_path() -> str: + """Resolve the on-volume db path the rest of Matilde's private data uses. + + Precedence: ``MATILDE_STUDY_DB`` (explicit override) → ``$HERMES_HOME``/ + engagements → the repo's ``engagements/`` dir. The directory is created if + needed; this keeps study state on the mounted data volume so it survives a + container rebuild. + """ + override = os.environ.get("MATILDE_STUDY_DB", "").strip() + if override: + return override + home = os.environ.get("HERMES_HOME", "").strip() + if home: + base = os.path.join(home, "engagements") + else: + # Repo-relative engagements/ (gitignored) — the default private data dir. + engine_dir = os.path.dirname(os.path.realpath(__file__)) + repo_root = os.path.normpath(os.path.join(engine_dir, "..", "..")) + base = os.path.join(repo_root, "engagements") + return os.path.join(base, "studies.db") + + +class StudyStore: + """SQLite-backed store for studies, steps, artifacts, and findings. + + Open many instances against the same db file — each opens its own connection; + WAL mode keeps reads/writes consistent. Every mutating method commits + immediately (atomic), so state is durable the instant a method returns. + """ + + def __init__(self, db_path: Optional[str] = None): + self.db_path = db_path or default_db_path() + parent = os.path.dirname(os.path.abspath(self.db_path)) + if parent: + os.makedirs(parent, exist_ok=True) + self._conn = sqlite3.connect(self.db_path, isolation_level=None) + self._conn.row_factory = sqlite3.Row + self._conn.execute("PRAGMA journal_mode=WAL;") + self._conn.execute("PRAGMA foreign_keys=ON;") + self._conn.executescript(_SCHEMA) + + # ---- studies --------------------------------------------------------- + + def create_study(self, slug: str, title: str = "", + plan: Optional[List[str]] = None, + meta: Optional[dict] = None) -> int: + """Create (or return the existing) study for *slug*. Idempotent on slug.""" + plan = list(plan or []) + existing = self._conn.execute( + "SELECT id FROM studies WHERE slug = ?", (slug,)).fetchone() + if existing is not None: + return int(existing["id"]) + now = _now() + cur = self._conn.execute( + "INSERT INTO studies (slug, title, status, plan_json, meta_json, " + "created_at, updated_at) VALUES (?, ?, 'created', ?, ?, ?, ?)", + (slug, title, json.dumps(plan), json.dumps(meta or {}), now, now), + ) + return int(cur.lastrowid) + + def get_study(self, study_id: int) -> Optional[dict]: + row = self._conn.execute( + "SELECT * FROM studies WHERE id = ?", (study_id,)).fetchone() + return self._study_row(row) if row else None + + def get_study_by_slug(self, slug: str) -> Optional[dict]: + row = self._conn.execute( + "SELECT * FROM studies WHERE slug = ?", (slug,)).fetchone() + return self._study_row(row) if row else None + + def list_studies(self, limit: int = 50) -> List[dict]: + rows = self._conn.execute( + "SELECT * FROM studies ORDER BY id DESC LIMIT ?", (int(limit),) + ).fetchall() + return [self._study_row(r) for r in rows] + + def set_study_status(self, study_id: int, status: str) -> None: + self._conn.execute( + "UPDATE studies SET status = ?, updated_at = ? WHERE id = ?", + (status, _now(), study_id)) + + @staticmethod + def _study_row(row: sqlite3.Row) -> dict: + return { + "id": int(row["id"]), + "slug": row["slug"], + "title": row["title"], + "status": row["status"], + "plan": json.loads(row["plan_json"] or "[]"), + "meta": json.loads(row["meta_json"] or "{}"), + "created_at": row["created_at"], + "updated_at": row["updated_at"], + } + + # ---- steps ----------------------------------------------------------- + + def add_steps(self, study_id: int, names: List[str]) -> None: + """Create pending steps in order. Idempotent: existing steps keep their + status; new names are appended after the current max index.""" + existing = {r["name"]: r["idx"] for r in self._conn.execute( + "SELECT name, idx FROM steps WHERE study_id = ?", (study_id,))} + next_idx = (max(existing.values()) + 1) if existing else 0 + for name in names: + if name in existing: + continue + self._conn.execute( + "INSERT INTO steps (study_id, idx, name, status) " + "VALUES (?, ?, ?, 'pending')", + (study_id, next_idx, name)) + next_idx += 1 + + def get_steps(self, study_id: int) -> List[dict]: + rows = self._conn.execute( + "SELECT * FROM steps WHERE study_id = ? ORDER BY idx ASC", (study_id,) + ).fetchall() + return [self._step_row(r) for r in rows] + + def get_step(self, study_id: int, name: str) -> Optional[dict]: + row = self._conn.execute( + "SELECT * FROM steps WHERE study_id = ? AND name = ?", + (study_id, name)).fetchone() + return self._step_row(row) if row else None + + def set_step_status(self, study_id: int, name: str, status: str, + error: Optional[str] = None) -> None: + now = _now() + sets = ["status = ?"] + params: list = [status] + if status == "running": + sets.append("started_at = ?") + params.append(now) + if status in ("done", "failed", "skipped"): + sets.append("finished_at = ?") + params.append(now) + if error is not None: + sets.append("error = ?") + params.append(error) + params += [study_id, name] + self._conn.execute( + f"UPDATE steps SET {', '.join(sets)} WHERE study_id = ? AND name = ?", + params) + + def record_step_result(self, study_id: int, name: str, result: dict) -> None: + self._conn.execute( + "UPDATE steps SET result_json = ? WHERE study_id = ? AND name = ?", + (json.dumps(result, default=str), study_id, name)) + + @staticmethod + def _step_row(row: sqlite3.Row) -> dict: + return { + "id": int(row["id"]), + "study_id": int(row["study_id"]), + "idx": int(row["idx"]), + "name": row["name"], + "status": row["status"], + "result": json.loads(row["result_json"]) if row["result_json"] else None, + "error": row["error"], + "started_at": row["started_at"], + "finished_at": row["finished_at"], + } + + # ---- artifacts ------------------------------------------------------- + + def add_artifact(self, study_id: int, step_name: str, path: str, + kind: str = "", sha256: str = "", + bytes: Optional[int] = None, # noqa: A002 (mirror schema) + meta: Optional[dict] = None) -> int: + cur = self._conn.execute( + "INSERT INTO artifacts (study_id, step_name, path, kind, sha256, " + "bytes, meta_json, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", + (study_id, step_name, path, kind, sha256, bytes, + json.dumps(meta or {}), _now())) + return int(cur.lastrowid) + + def get_artifacts(self, study_id: int) -> List[dict]: + rows = self._conn.execute( + "SELECT * FROM artifacts WHERE study_id = ? ORDER BY id ASC", + (study_id,)).fetchall() + return [{ + "id": int(r["id"]), + "step_name": r["step_name"], + "path": r["path"], + "kind": r["kind"], + "sha256": r["sha256"], + "bytes": r["bytes"], + "meta": json.loads(r["meta_json"] or "{}"), + "created_at": r["created_at"], + } for r in rows] + + # ---- findings -------------------------------------------------------- + + def add_finding(self, study_id: int, step_name: str, claim: str, + verdict: str, score: Optional[float] = None, + evidence: Optional[dict] = None) -> int: + cur = self._conn.execute( + "INSERT INTO findings (study_id, step_name, claim, verdict, score, " + "evidence_json, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)", + (study_id, step_name, claim, verdict, score, + json.dumps(evidence or {}), _now())) + return int(cur.lastrowid) + + def get_findings(self, study_id: int) -> List[dict]: + rows = self._conn.execute( + "SELECT * FROM findings WHERE study_id = ? ORDER BY id ASC", + (study_id,)).fetchall() + return [{ + "id": int(r["id"]), + "step_name": r["step_name"], + "claim": r["claim"], + "verdict": r["verdict"], + "score": r["score"], + "evidence": json.loads(r["evidence_json"] or "{}"), + "created_at": r["created_at"], + } for r in rows] + + # ---- summary --------------------------------------------------------- + + def study_summary(self, study_id: int) -> Optional[dict]: + study = self.get_study(study_id) + if study is None: + return None + steps = self.get_steps(study_id) + findings = self.get_findings(study_id) + counts: dict = {} + for s in steps: + counts[s["status"]] = counts.get(s["status"], 0) + 1 + return { + "study": study, + "steps": steps, + "step_status_counts": counts, + "artifacts": self.get_artifacts(study_id), + "findings": findings, + "finding_count": len(findings), + } + + def close(self) -> None: + try: + self._conn.close() + except Exception: + pass + + def __del__(self): # best-effort cleanup on GC (simulated restart in tests) + self.close() diff --git a/matilde_plugin/tools.py b/matilde_plugin/tools.py index 83da41d..30c7a44 100644 --- a/matilde_plugin/tools.py +++ b/matilde_plugin/tools.py @@ -433,3 +433,215 @@ def _handle_fetch_fulltext(args: dict, **kwargs: Any) -> str: return _tool_result(payload) except Exception as exc: return _tool_error(f"fetch_fulltext failed: {type(exc).__name__}: {exc}") + + +# --------------------------------------------------------------------------- +# Study pipeline — durable, resumable, agent-drivable analyses +# +# These tools let the agent kick off / advance / inspect a long analysis without +# holding it in a single turn. State lives in a SQLite ``StudyStore`` on the +# mounted data volume, so an OOM or container rebuild mid-run is recovered by +# simply calling matilde_study_run again. Logic lives in engine/store.py + +# engine/pipeline.py; these handlers stay thin. +# --------------------------------------------------------------------------- + +def _open_store() -> Any: + """Open the StudyStore at the injected/default on-volume db path.""" + from .engine.store import StudyStore + return StudyStore() # default_db_path() honors MATILDE_STUDY_DB / HERMES_HOME + + +def _coerce_plan(plan: Any) -> list: + """Accept a list of step names or a comma/newline-separated string.""" + if isinstance(plan, list): + return [str(p).strip() for p in plan if str(p).strip()] + if isinstance(plan, str): + sep = "," if "," in plan else "\n" + return [p.strip() for p in plan.split(sep) if p.strip()] + return [] + + +STUDY_CREATE_SCHEMA = { + "name": "matilde_study_create", + "description": ( + "Create a durable, resumable study — a long analysis decomposed into " + "ordered steps whose state is checkpointed to disk so it survives an OOM " + "or container rebuild. Provide a unique 'slug', a human 'title', and a " + "'plan' (ordered step names). Returns the study_id. For the built-in " + "bibliography-validation study, set kind='bibliography' and pass 'bibtex'; " + "then call matilde_study_run to verify each reference. Creating with an " + "existing slug returns that study (idempotent)." + ), + "parameters": { + "type": "object", + "properties": { + "slug": {"type": "string", "description": "Unique short id for the study (e.g. 'paper42-bib')."}, + "title": {"type": "string", "description": "Human-readable title."}, + "plan": {"type": "array", "items": {"type": "string"}, + "description": "Ordered step names (e.g. ['parse_refs','verify_each','summarize'])."}, + "kind": {"type": "string", "description": "Optional study kind. 'bibliography' wires the built-in reference-validation steps."}, + "bibtex": {"type": "string", "description": "For kind='bibliography': the BibTeX to validate (stored on the study)."}, + }, + "required": ["slug"], + }, +} + + +def _handle_study_create(args: dict, **kwargs: Any) -> str: + slug = str(args.get("slug", "")).strip() + if not slug: + return _tool_error("'slug' is required (a unique short id for the study).") + try: + title = str(args.get("title", "")).strip() or slug + plan = _coerce_plan(args.get("plan")) + kind = str(args.get("kind", "")).strip() + bibtex = args.get("bibtex") + if kind == "bibliography" and not plan: + plan = ["parse_refs", "verify_each", "summarize"] + meta: dict = {} + if kind: + meta["kind"] = kind + if isinstance(bibtex, str) and bibtex.strip(): + meta["bibtex"] = bibtex + store = _open_store() + sid = store.create_study(slug=slug, title=title, plan=plan, meta=meta) + store.add_steps(sid, plan) + return _tool_result( + study_id=sid, slug=slug, plan=plan, + message=f"Created study {sid} ('{slug}') with {len(plan)} step(s).") + except Exception as exc: + return _tool_error(f"study_create failed: {type(exc).__name__}: {exc}") + + +STUDY_RUN_SCHEMA = { + "name": "matilde_study_run", + "description": ( + "Advance a study: run its pending steps in order, checkpointing after " + "each. Already-completed steps are skipped, so if a previous run was " + "OOM-killed or the container was rebuilt mid-run, just call this again to " + "resume from the last completed step. Returns the study status and a " + "per-step summary; on a step failure the study is left 'blocked' and " + "resumable. For bibliography studies the BibTeX stored at create time is " + "used automatically." + ), + "parameters": { + "type": "object", + "properties": { + "study_id": {"type": "integer", "description": "The study to advance (from matilde_study_create)."}, + }, + "required": ["study_id"], + }, +} + + +def _handle_study_run(args: dict, **kwargs: Any) -> str: + sid_raw = args.get("study_id") + try: + sid = int(sid_raw) + except (TypeError, ValueError): + return _tool_error("'study_id' is required and must be an integer.") + try: + from .engine.pipeline import run + store = _open_store() + study = store.get_study(sid) + if study is None: + return _tool_error(f"No study with id {sid}.", study_id=sid) + meta = study.get("meta") or {} + kind = meta.get("kind") + if kind == "bibliography": + from .engine.bibliography_study import build_steps + steps = build_steps(bibtex=meta.get("bibtex", "")) + else: + return _tool_error( + f"Study {sid} has no runnable step implementations " + f"(kind={kind!r}). Built-in runner currently supports " + f"kind='bibliography'.", study_id=sid) + summary = run(store, sid, steps) + return _tool_result( + study_id=sid, status=summary["status"], steps=summary["steps"], + failed_step=summary.get("failed_step"), + message=f"Study {sid} is now '{summary['status']}'.") + except Exception as exc: + return _tool_error(f"study_run failed: {type(exc).__name__}: {exc}") + + +STUDY_STATUS_SCHEMA = { + "name": "matilde_study_status", + "description": ( + "Get the full status of a study: its steps (with per-step status), " + "recorded findings (claim + verdict + evidence), artifacts, and a " + "step-status tally. Use this to inspect progress or read out the " + "scientific findings after a run." + ), + "parameters": { + "type": "object", + "properties": { + "study_id": {"type": "integer", "description": "The study to inspect."}, + }, + "required": ["study_id"], + }, +} + + +def _handle_study_status(args: dict, **kwargs: Any) -> str: + sid_raw = args.get("study_id") + try: + sid = int(sid_raw) + except (TypeError, ValueError): + return _tool_error("'study_id' is required and must be an integer.") + try: + store = _open_store() + summary = store.study_summary(sid) + if summary is None: + return _tool_error(f"No study with id {sid}.", study_id=sid) + study = summary["study"] + return _tool_result( + study_id=sid, + slug=study["slug"], + status=study["status"], + steps=summary["steps"], + findings=summary["findings"], + finding_count=summary["finding_count"], + artifacts=summary["artifacts"], + step_status_counts=summary["step_status_counts"], + message=(f"Study {sid} ('{study['slug']}') is '{study['status']}' — " + f"{summary['finding_count']} finding(s).")) + except Exception as exc: + return _tool_error(f"study_status failed: {type(exc).__name__}: {exc}") + + +STUDY_LIST_SCHEMA = { + "name": "matilde_study_list", + "description": ( + "List recent studies (most recent first) with their slug, title, and " + "status. Use to discover existing studies — e.g. to find one to resume " + "after a restart." + ), + "parameters": { + "type": "object", + "properties": { + "limit": {"type": "integer", "description": "How many studies to return (default 20)."}, + }, + "required": [], + }, +} + + +def _handle_study_list(args: dict, **kwargs: Any) -> str: + try: + limit = args.get("limit") or 20 + try: + limit = max(1, min(int(limit), 200)) + except (TypeError, ValueError): + limit = 20 + store = _open_store() + studies = [ + {"id": s["id"], "slug": s["slug"], "title": s["title"], + "status": s["status"], "updated_at": s["updated_at"]} + for s in store.list_studies(limit=limit) + ] + return _tool_result( + count=len(studies), studies=studies, + message=f"Listed {len(studies)} study(ies).") + except Exception as exc: + return _tool_error(f"study_list failed: {type(exc).__name__}: {exc}") diff --git a/tests/test_bibliography_study.py b/tests/test_bibliography_study.py new file mode 100644 index 0000000..115c2b7 --- /dev/null +++ b/tests/test_bibliography_study.py @@ -0,0 +1,104 @@ +"""Tests for the first concrete pipeline: a bibliography-validation study. + +Steps: parse_refs -> verify_each -> summarize, reusing the citations/parsing +engine. The verifier is INJECTED so a test can force a failure on a specific +reference, prove the study blocks with the earlier refs done, then swap in a +working verifier and resume — without re-verifying the refs that already passed. +This is the end-to-end proof of the resumability guarantee. +""" +from __future__ import annotations + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.normpath(os.path.join(os.path.dirname(__file__), ".."))) + +from matilde_plugin.engine.bibliography_study import build_steps # noqa: E402 +from matilde_plugin.engine.pipeline import run, resume # noqa: E402 +from matilde_plugin.engine.store import StudyStore # noqa: E402 + +# A tiny BibTeX with three entries (all PUBLIC reference material). +BIB = """ +@article{a2017, title={Attention Is All You Need}, author={Vaswani, A}, year={2017}, doi={10.48550/arXiv.1706.03762}} +@article{w1953, title={Molecular structure of nucleic acids}, author={Watson, J and Crick, F}, year={1953}, doi={10.1038/171737a0}} +@article{wakefield1998, title={Ileal-lymphoid-nodular hyperplasia}, author={Wakefield, A}, year={1998}, doi={10.1016/S0140-6736(97)11096-0}} +""" + + +def _ok_verifier(ref): + """A fake verifier: every ref is 'supported'. Records which refs it saw.""" + return {"verdict": "supported", "score": 0.9, + "claim": f"reference exists: {ref.title or ref.doi}", + "evidence": {"doi": ref.doi}} + + +def _failing_on_third(seen): + """Returns a verifier that raises on the 3rd distinct ref it is asked about.""" + def verify(ref): + seen.append(ref.doi or ref.title) + if len(seen) >= 3: + raise RuntimeError("verifier provider error on ref #3") + return {"verdict": "supported", "score": 0.9, + "claim": f"reference exists: {ref.title or ref.doi}", + "evidence": {"doi": ref.doi}} + return verify + + +@pytest.fixture() +def store(tmp_path): + return StudyStore(str(tmp_path / "studies.db")) + + +def test_happy_path_records_a_finding_per_reference(store): + sid = store.create_study(slug="bib", title="Bib", + plan=["parse_refs", "verify_each", "summarize"]) + steps = build_steps(bibtex=BIB, verifier=_ok_verifier) + summary = run(store, sid, steps) + + assert store.get_study(sid)["status"] == "done" + findings = store.get_findings(sid) + assert len(findings) == 3 + assert all(f["verdict"] == "supported" for f in findings) + # summarize step produced a roll-up + summarize_result = store.get_step(sid, "summarize")["result"] + assert summarize_result["total"] == 3 + assert summarize_result["supported"] == 3 + + +def test_fail_then_resume_does_not_reverify_done_refs(store): + sid = store.create_study(slug="bib", title="Bib", + plan=["parse_refs", "verify_each", "summarize"]) + + # --- First pass: verifier fails on ref #3 --- + seen = [] + failing_steps = build_steps(bibtex=BIB, verifier=_failing_on_third(seen)) + run(store, sid, failing_steps) + + # parse_refs done; verify_each failed; study resumable. + assert store.get_step(sid, "parse_refs")["status"] == "done" + assert store.get_step(sid, "verify_each")["status"] == "failed" + assert store.get_step(sid, "summarize")["status"] == "pending" + assert store.get_study(sid)["status"] in ("blocked", "failed") + # Refs 1-2 were verified and persisted as findings before the failure. + findings_after_fail = store.get_findings(sid) + assert len(findings_after_fail) == 2 + + # --- Second pass: swap in a working verifier and resume --- + reverified = [] + + def working(ref): + reverified.append(ref.doi or ref.title) + return _ok_verifier(ref) + + fixed_steps = build_steps(bibtex=BIB, verifier=working) + resume(store, sid, fixed_steps) + + # Study completes. + assert store.get_study(sid)["status"] == "done" + # parse_refs was NOT re-run (it was done) — verifier only saw the unfinished refs. + # Refs 1-2 already had findings; only ref #3 (and any after) get verified now. + assert len(reverified) == 1 + findings = store.get_findings(sid) + assert len(findings) == 3 diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py new file mode 100644 index 0000000..ec47fb2 --- /dev/null +++ b/tests/test_pipeline.py @@ -0,0 +1,140 @@ +"""Tests for the resumable pipeline runner. + +Steps are in-memory fakes with side-effect counters, so we can prove the runner +(a) skips already-done steps on resume and (b) stops on a failing step leaving the +study resumable, then continues from the first non-done step on resume — WITHOUT +re-running the steps that already completed. +""" +from __future__ import annotations + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.normpath(os.path.join(os.path.dirname(__file__), ".."))) + +from matilde_plugin.engine.pipeline import Step, StepResult, run, resume # noqa: E402 +from matilde_plugin.engine.store import StudyStore # noqa: E402 + + +@pytest.fixture() +def store(tmp_path): + return StudyStore(str(tmp_path / "studies.db")) + + +def _counter_step(name, calls, *, data=None, fail=False): + def fn(ctx): + calls.append(name) + if fail: + raise RuntimeError(f"{name} blew up") + return StepResult(data=data or {"step": name}) + return Step(name=name, fn=fn) + + +def test_run_executes_all_steps_and_marks_done(store): + sid = store.create_study(slug="s", title="S", plan=[]) + calls = [] + steps = [_counter_step("a", calls), _counter_step("b", calls)] + summary = run(store, sid, steps) + assert calls == ["a", "b"] + assert store.get_step(sid, "a")["status"] == "done" + assert store.get_step(sid, "b")["status"] == "done" + assert store.get_study(sid)["status"] == "done" + assert [s["name"] for s in summary["steps"]] == ["a", "b"] + + +def test_run_persists_step_data_artifacts_and_findings(store): + sid = store.create_study(slug="s", title="S", plan=[]) + + def fn(ctx): + return StepResult( + data={"k": "v"}, + artifacts=[{"path": "/x", "kind": "blob", "bytes": 1}], + findings=[{"claim": "c", "verdict": "supported", "score": 1.0}], + ) + + run(store, sid, [Step(name="only", fn=fn)]) + assert store.get_step(sid, "only")["result"] == {"k": "v"} + assert len(store.get_artifacts(sid)) == 1 + assert store.get_findings(sid)[0]["claim"] == "c" + + +def test_run_skips_already_done_steps(store): + sid = store.create_study(slug="s", title="S", plan=[]) + store.add_steps(sid, ["a", "b"]) + store.set_step_status(sid, "a", "done") # pretend 'a' already ran + calls = [] + steps = [_counter_step("a", calls), _counter_step("b", calls)] + run(store, sid, steps) + assert calls == ["b"] # 'a' was skipped + + +def test_run_stops_on_failure_and_leaves_study_resumable(store): + sid = store.create_study(slug="s", title="S", plan=[]) + calls = [] + steps = [ + _counter_step("a", calls), + _counter_step("b", calls, fail=True), + _counter_step("c", calls), + ] + summary = run(store, sid, steps) + assert calls == ["a", "b"] # stopped at b; c never ran + assert store.get_step(sid, "a")["status"] == "done" + assert store.get_step(sid, "b")["status"] == "failed" + assert store.get_step(sid, "c")["status"] == "pending" + assert store.get_study(sid)["status"] in ("blocked", "failed") + assert summary["failed_step"] == "b" + + +def test_resume_continues_from_first_non_done_without_rerunning(store): + sid = store.create_study(slug="s", title="S", plan=[]) + + # First pass: 'b' fails. + calls1 = [] + failing = [ + _counter_step("a", calls1), + _counter_step("b", calls1, fail=True), + _counter_step("c", calls1), + ] + run(store, sid, failing) + assert calls1 == ["a", "b"] + + # Second pass: swap in a WORKING 'b'. Resume must NOT re-run 'a'. + calls2 = [] + fixed = [ + _counter_step("a", calls2), + _counter_step("b", calls2), # now succeeds + _counter_step("c", calls2), + ] + resume(store, sid, fixed) + assert calls2 == ["b", "c"] # 'a' skipped (already done), continues from b + assert store.get_study(sid)["status"] == "done" + assert store.get_step(sid, "c")["status"] == "done" + + +def test_resume_of_completed_study_is_noop(store): + sid = store.create_study(slug="s", title="S", plan=[]) + calls = [] + steps = [_counter_step("a", calls), _counter_step("b", calls)] + run(store, sid, steps) + assert calls == ["a", "b"] + calls.clear() + resume(store, sid, steps) + assert calls == [] # nothing re-ran + assert store.get_study(sid)["status"] == "done" + + +def test_step_context_exposes_prior_results(store): + sid = store.create_study(slug="s", title="S", plan=[]) + seen = {} + + def a(ctx): + return StepResult(data={"value": 42}) + + def b(ctx): + seen["prior"] = ctx.results.get("a") + return StepResult(data={}) + + run(store, sid, [Step(name="a", fn=a), Step(name="b", fn=b)]) + assert seen["prior"] == {"value": 42} diff --git a/tests/test_plugin_tools.py b/tests/test_plugin_tools.py index 1482974..82d96a2 100644 --- a/tests/test_plugin_tools.py +++ b/tests/test_plugin_tools.py @@ -35,6 +35,10 @@ def test_plugin_registers_expected_tools(): "matilde_openneuro_search", "matilde_openneuro_list_files", "matilde_fetch_fulltext", + "matilde_study_create", + "matilde_study_run", + "matilde_study_status", + "matilde_study_list", ] diff --git a/tests/test_store.py b/tests/test_store.py new file mode 100644 index 0000000..8532614 --- /dev/null +++ b/tests/test_store.py @@ -0,0 +1,179 @@ +"""Tests for the StudyStore — durable, checkpointed analysis state (SQLite). + +All I/O is a temp db path (injected), so these run offline and deterministically. +The load-bearing test is ``test_state_survives_fresh_store_instance``: it proves +the store's "survives container rebuild" guarantee by opening a NEW StudyStore on +the same db file and reading back everything the first instance wrote. +""" +from __future__ import annotations + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.normpath(os.path.join(os.path.dirname(__file__), ".."))) + +from matilde_plugin.engine.store import StudyStore # noqa: E402 + + +@pytest.fixture() +def db_path(tmp_path): + return str(tmp_path / "studies.db") + + +def test_create_and_get_study(db_path): + store = StudyStore(db_path) + sid = store.create_study(slug="bib-check", title="Bibliography check", + plan=["parse_refs", "verify_each", "summarize"]) + assert isinstance(sid, int) + study = store.get_study(sid) + assert study["slug"] == "bib-check" + assert study["title"] == "Bibliography check" + assert study["status"] == "created" + assert study["plan"] == ["parse_refs", "verify_each", "summarize"] + assert study["created_at"] + assert study["updated_at"] + + +def test_create_study_is_idempotent_on_slug(db_path): + store = StudyStore(db_path) + sid1 = store.create_study(slug="dup", title="First", plan=["a"]) + sid2 = store.create_study(slug="dup", title="Second", plan=["a", "b"]) + # Same slug returns the same study (upsert), not a duplicate row. + assert sid1 == sid2 + assert len(store.list_studies()) == 1 + + +def test_add_steps_creates_pending_steps_in_order(db_path): + store = StudyStore(db_path) + sid = store.create_study(slug="s", title="S", plan=[]) + store.add_steps(sid, ["one", "two", "three"]) + steps = store.get_steps(sid) + assert [s["name"] for s in steps] == ["one", "two", "three"] + assert [s["idx"] for s in steps] == [0, 1, 2] + assert all(s["status"] == "pending" for s in steps) + + +def test_add_steps_is_idempotent(db_path): + store = StudyStore(db_path) + sid = store.create_study(slug="s", title="S", plan=[]) + store.add_steps(sid, ["one", "two"]) + store.add_steps(sid, ["one", "two", "three"]) # re-add + extend + steps = store.get_steps(sid) + assert [s["name"] for s in steps] == ["one", "two", "three"] + + +def test_set_step_status_transitions(db_path): + store = StudyStore(db_path) + sid = store.create_study(slug="s", title="S", plan=[]) + store.add_steps(sid, ["one"]) + store.set_step_status(sid, "one", "running") + assert store.get_step(sid, "one")["status"] == "running" + store.set_step_status(sid, "one", "done") + assert store.get_step(sid, "one")["status"] == "done" + store.set_step_status(sid, "one", "failed", error="boom") + step = store.get_step(sid, "one") + assert step["status"] == "failed" + assert step["error"] == "boom" + + +def test_record_step_result_persists_json(db_path): + store = StudyStore(db_path) + sid = store.create_study(slug="s", title="S", plan=[]) + store.add_steps(sid, ["one"]) + store.record_step_result(sid, "one", {"count": 3, "items": ["a", "b"]}) + step = store.get_step(sid, "one") + assert step["result"] == {"count": 3, "items": ["a", "b"]} + + +def test_add_and_get_artifact(db_path): + store = StudyStore(db_path) + sid = store.create_study(slug="s", title="S", plan=[]) + store.add_steps(sid, ["one"]) + store.add_artifact(sid, "one", path="/data/out.npy", kind="array", + sha256="abc", bytes=128, meta={"shape": [4]}) + arts = store.get_artifacts(sid) + assert len(arts) == 1 + a = arts[0] + assert a["path"] == "/data/out.npy" + assert a["kind"] == "array" + assert a["sha256"] == "abc" + assert a["bytes"] == 128 + assert a["meta"] == {"shape": [4]} + + +def test_add_and_get_finding(db_path): + store = StudyStore(db_path) + sid = store.create_study(slug="s", title="S", plan=[]) + store.add_steps(sid, ["verify"]) + store.add_finding(sid, "verify", claim="ref #1 exists", verdict="supported", + score=0.95, evidence={"doi": "10.1/x"}) + findings = store.get_findings(sid) + assert len(findings) == 1 + f = findings[0] + assert f["claim"] == "ref #1 exists" + assert f["verdict"] == "supported" + assert f["score"] == 0.95 + assert f["evidence"] == {"doi": "10.1/x"} + + +def test_study_summary_aggregates_steps_and_findings(db_path): + store = StudyStore(db_path) + sid = store.create_study(slug="s", title="S", plan=["a", "b"]) + store.add_steps(sid, ["a", "b"]) + store.set_step_status(sid, "a", "done") + store.add_finding(sid, "a", claim="c", verdict="supported", score=1.0) + summary = store.study_summary(sid) + assert summary["study"]["slug"] == "s" + assert summary["step_status_counts"]["done"] == 1 + assert summary["step_status_counts"]["pending"] == 1 + assert summary["finding_count"] == 1 + assert len(summary["findings"]) == 1 + assert any(s["name"] == "a" and s["status"] == "done" for s in summary["steps"]) + + +def test_list_studies_returns_recent_first(db_path): + store = StudyStore(db_path) + a = store.create_study(slug="a", title="A", plan=[]) + b = store.create_study(slug="b", title="B", plan=[]) + studies = store.list_studies() + ids = [s["id"] for s in studies] + # most-recent-first ordering + assert ids.index(b) < ids.index(a) + + +def test_get_study_missing_returns_none(db_path): + store = StudyStore(db_path) + assert store.get_study(9999) is None + + +# --- The load-bearing guarantee: state survives a fresh StudyStore instance --- + +def test_state_survives_fresh_store_instance(db_path): + """Simulates a container rebuild: a brand-new StudyStore on the same db file + must see everything the first instance committed.""" + store1 = StudyStore(db_path) + sid = store1.create_study(slug="durable", title="Durable", + plan=["parse", "verify"]) + store1.add_steps(sid, ["parse", "verify"]) + store1.set_step_status(sid, "parse", "done") + store1.record_step_result(sid, "parse", {"n": 2}) + store1.add_finding(sid, "verify", claim="x", verdict="supported", score=0.9) + del store1 # drop the connection — simulate process death / rebuild + + store2 = StudyStore(db_path) # fresh instance, same file + study = store2.get_study(sid) + assert study is not None + assert study["slug"] == "durable" + assert store2.get_step(sid, "parse")["status"] == "done" + assert store2.get_step(sid, "parse")["result"] == {"n": 2} + findings = store2.get_findings(sid) + assert len(findings) == 1 + assert findings[0]["claim"] == "x" + + +def test_wal_mode_is_enabled(db_path): + store = StudyStore(db_path) + mode = store._conn.execute("PRAGMA journal_mode;").fetchone()[0] + assert mode.lower() == "wal" diff --git a/tests/test_study_tools.py b/tests/test_study_tools.py new file mode 100644 index 0000000..02cc136 --- /dev/null +++ b/tests/test_study_tools.py @@ -0,0 +1,103 @@ +"""Offline tests for the study-pipeline agent tools. + +Mirrors test_plugin_tools.py: loads the plugin, checks the new tools register, +coerce args, and return well-formed JSON envelopes. The db path is injected via +MATILDE_STUDY_DB so tests never touch the real engagements dir. +""" +from __future__ import annotations + +import importlib.util +import json +import os +import sys + +ROOT = os.path.normpath(os.path.join(os.path.dirname(__file__), "..")) +sys.path.insert(0, ROOT) + + +def _load_plugin(): + path = os.path.join(ROOT, "matilde_plugin", "__init__.py") + spec = importlib.util.spec_from_file_location("matilde_plugin", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def test_plugin_registers_study_tools(): + plugin = _load_plugin() + names = [t[0] for t in plugin._TOOLS] + for expected in ("matilde_study_create", "matilde_study_run", + "matilde_study_status", "matilde_study_list"): + assert expected in names + + +def test_study_create_requires_slug(monkeypatch, tmp_path): + plugin = _load_plugin() + monkeypatch.setenv("MATILDE_STUDY_DB", str(tmp_path / "s.db")) + out = json.loads(plugin._handle_study_create({})) + assert out["success"] is False + + +def test_study_create_returns_study_id(monkeypatch, tmp_path): + plugin = _load_plugin() + monkeypatch.setenv("MATILDE_STUDY_DB", str(tmp_path / "s.db")) + out = json.loads(plugin._handle_study_create( + {"slug": "bib", "title": "Bib", "plan": ["parse_refs", "verify_each"]})) + assert out["success"] is True + assert isinstance(out["study_id"], int) + + +def test_study_create_coerces_comma_separated_plan(monkeypatch, tmp_path): + plugin = _load_plugin() + monkeypatch.setenv("MATILDE_STUDY_DB", str(tmp_path / "s.db")) + out = json.loads(plugin._handle_study_create( + {"slug": "p", "title": "P", "plan": "a, b, c"})) + assert out["success"] is True + status = json.loads(plugin._handle_study_status({"study_id": out["study_id"]})) + assert [s["name"] for s in status["steps"]] == ["a", "b", "c"] + + +def test_study_status_requires_id(monkeypatch, tmp_path): + plugin = _load_plugin() + monkeypatch.setenv("MATILDE_STUDY_DB", str(tmp_path / "s.db")) + out = json.loads(plugin._handle_study_status({})) + assert out["success"] is False + + +def test_study_status_unknown_id_errors(monkeypatch, tmp_path): + plugin = _load_plugin() + monkeypatch.setenv("MATILDE_STUDY_DB", str(tmp_path / "s.db")) + out = json.loads(plugin._handle_study_status({"study_id": 9999})) + assert out["success"] is False + + +def test_study_list_returns_created_studies(monkeypatch, tmp_path): + plugin = _load_plugin() + monkeypatch.setenv("MATILDE_STUDY_DB", str(tmp_path / "s.db")) + plugin._handle_study_create({"slug": "one", "title": "One", "plan": ["a"]}) + plugin._handle_study_create({"slug": "two", "title": "Two", "plan": ["a"]}) + out = json.loads(plugin._handle_study_list({})) + assert out["success"] is True + slugs = [s["slug"] for s in out["studies"]] + assert "one" in slugs and "two" in slugs + + +def test_study_run_advances_and_is_resumable(monkeypatch, tmp_path): + """End-to-end through the tools using the bibliography study. Coerces a string + id, runs to completion, and a second run is a no-op (resume of a done study).""" + plugin = _load_plugin() + monkeypatch.setenv("MATILDE_STUDY_DB", str(tmp_path / "s.db")) + bib = ("@article{w1953, title={Molecular structure of nucleic acids}, " + "author={Watson, J}, year={1953}, doi={10.1038/171737a0}}") + created = json.loads(plugin._handle_study_create({ + "slug": "bibrun", "title": "Bib run", + "plan": ["parse_refs", "verify_each", "summarize"], + "kind": "bibliography", "bibtex": bib, + })) + sid = created["study_id"] + # study_run accepts a string id (arg coercion). + out = json.loads(plugin._handle_study_run({"study_id": str(sid)})) + assert out["success"] is True + assert out["status"] == "done" + status = json.loads(plugin._handle_study_status({"study_id": sid})) + assert status["finding_count"] == 1 From ba75eeeb9ba619df31d15a1e82f8e5e0cd747380 Mon Sep 17 00:00:00 2001 From: Juniper Bevensee Date: Mon, 22 Jun 2026 16:12:45 +1200 Subject: [PATCH 2/2] docs(pipeline): generalize motivation for public repo (drop dated incident specifics); note plan-evolution semantics Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/stateful-study-pipeline.md | 35 ++++++++++++++++++++------------- 1 file changed, 21 insertions(+), 14 deletions(-) diff --git a/docs/stateful-study-pipeline.md b/docs/stateful-study-pipeline.md index e7ced4e..c42f988 100644 --- a/docs/stateful-study-pipeline.md +++ b/docs/stateful-study-pipeline.md @@ -2,17 +2,19 @@ ## Motivation -Long analyses — e.g. validating a paper's findings against an OpenNeuro dataset — -currently run **inline in a single agent turn**: one-shot scripts that +Long analyses — for example validating a paper's reported findings against an +open dataset — tend to run **inline in a single agent turn** as one-shot +scripts. Two failure modes follow from that: -- **die on OOM** (observed 2026-06-22: a live MEG M100 validation full-preloaded a - 2.5 GB dataset into the agent container and was OOM-killed), and -- **lose all state when the container restarts** (the same session's analysis was - wiped when the agent image was rebuilt mid-run — no resume, no memory of it). +- **Memory exhaustion** — loading a multi-gigabyte dataset into the agent's + process can exceed the container's memory budget and get the process killed + mid-analysis. +- **Loss on restart** — if the container is recreated (deploy, rebuild, crash), + an in-flight analysis and its progress are gone, with no way to resume. -Both failure modes have the same remedy: **durable, checkpointed analysis state** -so an analysis resumes from the last completed step after an OOM *or* a container -rebuild, instead of vanishing. +Both have the same remedy: **durable, checkpointed analysis state** so an +analysis resumes from the last completed step after a memory kill *or* a +container restart, instead of vanishing. ## Hard requirements @@ -63,6 +65,10 @@ API (all idempotent / upsert where sensible): `create_study`, `get_study`, first non-`done` step. Re-running a fully-done study is a no-op. - Idempotency contract: step fns must be safe to re-run from scratch (the runner only guarantees done-steps are skipped, not partial-step rollback). +- Plan evolution: the runner only iterates the steps passed to it. If a step is + dropped from a later plan, its stored row is left as-is (orphaned, never + re-run) rather than reconciled. Fine for fixed-plan studies; a study whose + plan changes between runs should reconcile stored-vs-passed steps explicitly. ### Agent tools (`matilde_plugin/tools.py`) @@ -83,11 +89,12 @@ assert the study is `blocked` with refs 1–2 `done`; fix the injected verifier; ### Follow-up (separate PR — heavy data analysis) -**Data-validation study** (the MEG M100 case): steps `fetch_dataset → -preprocess → epoch → evoked → validate_finding`, each **memory-bounded** -(chunked/streamed, intermediates checkpointed as artifacts) so an OOM resumes -from the last completed step. Requires deciding the agent container memory budget -and where heavy deps (mne) are installed. Out of scope for the framework PR. +**Data-validation study** (e.g. checking a reported evoked-response peak in an +open MEG dataset): steps `fetch_dataset → preprocess → epoch → evoked → +validate_finding`, each **memory-bounded** (chunked/streamed, intermediates +checkpointed as artifacts) so a memory kill resumes from the last completed step. +Requires deciding the agent container memory budget and where heavy scientific +deps (e.g. mne) are installed. Out of scope for the framework PR. ## Testing