Files
stack/src/cli/prisma.py
kert 49a2f6f723 feat(prisma): fleet mode for eligibility and extract via a shared run_pool driver (refs #655)
prisma.fleet.run_pool owns the multi-host thread-pool loop that screen
had inline: at most --per-host in-flight requests per Ollama host, LLM
calls only on the pool, every SQLite touch and per-item commit on the
calling thread, and a host that fails twice in a row is retired for the
rest of the run (its queued items go to the surviving hosts). screen,
eligibility and extract each supply prepare/apply/on_error glue; items
with no retrievable full text resolve without a fleet call. CLI:
prisma eligible/extract gain --hosts/--fleet-model/--num-ctx/--per-host
matching prisma screen. Docs regenerated.
2026-09-11 18:23:37 -04:00

617 lines
20 KiB
Python

"""stack prisma — LLM-driven PRISMA screening workflow.
Stages:
1 identified (implicit — everything in the project tag)
2 screen title/abstract LLM pass
3 eligibility full-text LLM pass
3+ extract structured data extraction on included items
4 included derived from tag state
Commands are idempotent: tags record the transition a row has already
made, and the screen / eligibility loops skip items already past that
stage.
"""
from __future__ import annotations
import subprocess
from pathlib import Path
import typer
app = typer.Typer(no_args_is_help=True)
_DEFAULT_ZOT = Path("data/zotero/data/zotero.sqlite")
_DEFAULT_STORAGE = Path("data/zotero/data/storage")
def _hold(fn, *, hold: bool):
if not hold:
return fn()
subprocess.run(
["docker", "stop", "zotero"], capture_output=True, check=False, timeout=30
)
try:
return fn()
finally:
subprocess.run(
["docker", "start", "zotero"], capture_output=True, check=False, timeout=30
)
@app.command(name="init")
def init_project(
name: str = typer.Argument(..., help="Project slug, e.g. skin-subs."),
db: Path = typer.Option(_DEFAULT_ZOT, "--db", help="Zotero SQLite path."),
hold: bool = typer.Option(True, "--hold/--no-hold"),
) -> None:
"""Create (or re-load) the project's three Zotero anchor items:
criteria, extraction-template, and the reasons codebook."""
from prisma.project import init as _init
from zot.db import Db
def _go():
with Db(str(db)) as zdb:
project = _init(zdb, name)
typer.echo(f"project: {project.name}")
typer.echo(f" criteria: {len(project.criteria):>6} chars")
typer.echo(
f" extraction_template: {len(project.extraction_template):>6} chars"
)
typer.echo(f" reasons: {len(project.reasons):>6} chars")
_hold(_go, hold=hold)
@app.command(name="export")
def export_item(
zot_id: int = typer.Argument(..., help="Zotero itemID to export."),
fulltext: bool = typer.Option(False, "--fulltext"),
db: Path = typer.Option(_DEFAULT_ZOT, "--db"),
storage: Path = typer.Option(_DEFAULT_STORAGE, "--storage"),
hold: bool = typer.Option(True, "--hold/--no-hold"),
) -> None:
"""Dump one Zotero item as the LLM-ready markdown snapshot."""
from prisma.export import load_item, to_markdown
from zot.db import Db
def _go() -> str:
with Db(str(db)) as zdb:
snap = load_item(zdb, zot_id, storage)
return to_markdown(snap, include_fulltext=fulltext)
typer.echo(_hold(_go, hold=hold))
@app.command(name="ping-llm")
def ping_llm() -> None:
"""Send a cheap ping through the configured provider.
Useful for verifying ``PRISMA_LLM_PROVIDER`` + auth before kicking
off a long screening run.
"""
from prisma.llm import LLMCall, LLMMessage, make_provider
provider = make_provider()
call = LLMCall(
messages=[
LLMMessage(role="user", content="Respond with exactly: pong."),
],
max_tokens=8,
)
result = provider.complete(call)
typer.echo(f"text: {result.text.strip()}")
typer.echo(f"usage: {result.usage}")
# ── Stage loops ─────────────────────────────────────────────────
@app.command(name="screen")
def screen(
name: str = typer.Argument(..., help="Project slug."),
limit: int = typer.Option(
0, "--limit", "-n", help="0 = screen every unscreened item."
),
db: Path = typer.Option(_DEFAULT_ZOT, "--db"),
storage: Path = typer.Option(_DEFAULT_STORAGE, "--storage"),
hold: bool = typer.Option(True, "--hold/--no-hold"),
hosts: str = typer.Option(
"",
"--hosts",
help="Comma-separated Ollama base URLs — fleet mode (#655): "
"parallel screening via native /api/chat with per-request "
"num_ctx; bypasses PRISMA_LLM_* env config.",
),
fleet_model: str = typer.Option(
"qwen2.5:14b", "--fleet-model", help="Model name on every fleet host."
),
num_ctx: int = typer.Option(
8192, "--num-ctx", help="Context window per fleet request."
),
) -> None:
"""Stage 2 — LLM-driven title/abstract screening pass."""
from prisma import screen as _screen
from prisma.llm import make_provider
from prisma.project import load
from zot.db import Db
def _go():
with Db(str(db)) as zdb:
project = load(zdb, name)
if hosts:
return _screen.run_fleet(
zdb,
project,
hosts=[h.strip() for h in hosts.split(",") if h.strip()],
model=fleet_model,
num_ctx=num_ctx,
storage_dir=storage,
limit=limit or None,
)
return _screen.run(
zdb,
make_provider(),
project,
storage_dir=storage,
limit=limit or None,
)
stats = _hold(_go, hold=hold)
for k, v in stats.items():
typer.echo(f" {k}: {v}")
@app.command(name="eligible")
def eligible(
name: str = typer.Argument(..., help="Project slug."),
limit: int = typer.Option(0, "--limit", "-n"),
max_chars: int = typer.Option(
0,
"--max-chars",
help="Cap per-item markdown (full text) at N chars so bounded-"
"context local models keep the criteria in window. 0 = no cap.",
),
db: Path = typer.Option(_DEFAULT_ZOT, "--db"),
storage: Path = typer.Option(_DEFAULT_STORAGE, "--storage"),
hold: bool = typer.Option(True, "--hold/--no-hold"),
hosts: str = typer.Option(
"",
"--hosts",
help="Comma-separated Ollama base URLs — fleet mode (#655): "
"parallel eligibility via native /api/chat with per-request "
"num_ctx; bypasses PRISMA_LLM_* env config.",
),
fleet_model: str = typer.Option(
"qwen2.5:14b", "--fleet-model", help="Model name on every fleet host."
),
num_ctx: int = typer.Option(
8192, "--num-ctx", help="Context window per fleet request."
),
per_host: int = typer.Option(
1, "--per-host", help="Max in-flight requests per fleet host."
),
) -> None:
"""Stage 3 — LLM-driven full-text eligibility pass."""
from prisma import eligibility
from prisma.llm import make_provider
from prisma.project import load
from zot.db import Db
def _go():
with Db(str(db)) as zdb:
project = load(zdb, name)
if hosts:
return eligibility.run_fleet(
zdb,
project,
hosts=[h.strip() for h in hosts.split(",") if h.strip()],
model=fleet_model,
num_ctx=num_ctx,
per_host=per_host,
storage_dir=storage,
limit=limit or None,
max_fulltext_chars=max_chars or None,
)
return eligibility.run(
zdb,
make_provider(),
project,
storage_dir=storage,
limit=limit or None,
max_fulltext_chars=max_chars or None,
)
stats = _hold(_go, hold=hold)
for k, v in stats.items():
typer.echo(f" {k}: {v}")
@app.command(name="extract")
def extract(
name: str = typer.Argument(..., help="Project slug."),
limit: int = typer.Option(0, "--limit", "-n"),
db: Path = typer.Option(_DEFAULT_ZOT, "--db"),
storage: Path = typer.Option(_DEFAULT_STORAGE, "--storage"),
hold: bool = typer.Option(True, "--hold/--no-hold"),
hosts: str = typer.Option(
"",
"--hosts",
help="Comma-separated Ollama base URLs — fleet mode (#655): "
"parallel extraction via native /api/chat with per-request "
"num_ctx; bypasses PRISMA_LLM_* env config.",
),
fleet_model: str = typer.Option(
"qwen2.5:14b", "--fleet-model", help="Model name on every fleet host."
),
num_ctx: int = typer.Option(
8192, "--num-ctx", help="Context window per fleet request."
),
per_host: int = typer.Option(
1, "--per-host", help="Max in-flight requests per fleet host."
),
) -> None:
"""Stage 3+ — structured data extraction on included studies."""
from prisma import extract as _extract
from prisma.llm import make_provider
from prisma.project import load
from zot.db import Db
def _go():
with Db(str(db)) as zdb:
project = load(zdb, name)
if hosts:
return _extract.run_fleet(
zdb,
project,
hosts=[h.strip() for h in hosts.split(",") if h.strip()],
model=fleet_model,
num_ctx=num_ctx,
per_host=per_host,
storage_dir=storage,
limit=limit or None,
)
return _extract.run(
zdb,
make_provider(),
project,
storage_dir=storage,
limit=limit or None,
)
stats = _hold(_go, hold=hold)
for k, v in stats.items():
typer.echo(f" {k}: {v}")
@app.command(name="flow")
def flow(
name: str = typer.Argument(..., help="Project slug."),
mermaid: bool = typer.Option(True, "--mermaid/--text"),
db: Path = typer.Option(_DEFAULT_ZOT, "--db"),
hold: bool = typer.Option(True, "--hold/--no-hold"),
) -> None:
"""Emit the PRISMA 2020 flow diagram from current tag counts."""
from prisma import flow as _flow
from zot.db import Db
def _go() -> str:
with Db(str(db)) as zdb:
counts = _flow.count(zdb, name)
return (
_flow.mermaid(counts, project=name)
if mermaid
else _flow.text_summary(counts)
)
typer.echo(_hold(_go, hold=hold))
@app.command(name="fetch")
def fetch(
name: str = typer.Argument(..., help="Project slug."),
limit: int = typer.Option(0, "--limit", "-n"),
workers: int = typer.Option(
8,
"--workers",
"-w",
help="Concurrent fetch workers (fallback tier stays serialized).",
),
db: Path = typer.Option(_DEFAULT_ZOT, "--db"),
storage: Path = typer.Option(_DEFAULT_STORAGE, "--storage"),
scratch: Path = typer.Option(Path(".state/prisma-fetch"), "--scratch"),
proxy: str = typer.Option(
"", "--proxy", help="Override PRISMA_FETCH_PROXY (else env)."
),
use_vpn: bool = typer.Option(
True,
"--vpn/--no-vpn",
help="Auto-open an SSH tunnel to the provisioned droplet when one "
"exists. --no-vpn skips the tunnel and uses only the "
"unproxied sources.",
),
verify: bool = typer.Option(
True,
"--verify/--no-verify",
help="Curl through the proxy once before fetching and confirm "
"egress is non-US.",
),
hold: bool = typer.Option(True, "--hold/--no-hold"),
) -> None:
"""Fetch PDFs for every non-excluded item lacking an attachment.
Source cascade: Unpaywall → PMC → fallback (via the VPN droplet).
Items excluded at stage 2 are untouched — the queue filter only
passes `screen:include` and `screen:uncertain` through.
Tunnel lifecycle is automatic: if a droplet is tracked in
``.state/prisma-vpn/`` this command opens an SSH port-forward for
the duration of the fetch and closes it on exit. No sudo required.
"""
import contextlib
import os
from prisma import fetch as _fetch
from prisma import vpn as _vpn
from zot.db import Db
email = os.environ.get("UNPAYWALL_EMAIL", "dev@fhirworx.io")
explicit_proxy = proxy or os.environ.get("PRISMA_FETCH_PROXY") or None
droplet_up = _vpn.status().get("status") == "up"
auto_tunnel = use_vpn and droplet_up and not explicit_proxy
@contextlib.contextmanager
def _proxy_cm():
if explicit_proxy:
yield explicit_proxy
elif auto_tunnel:
with _vpn.active() as url:
if verify:
info = _vpn.verify_egress(url)
typer.echo(
f" egress: {info.get('ip', '?')} ({info.get('country', '?')})"
)
yield url
else:
if not droplet_up:
typer.echo(
" note: no VPN droplet tracked — fallback tier disabled. "
"Run `stack prisma vpn up` to provision one."
)
yield None
def _go():
with _proxy_cm() as fetch_proxy, Db(str(db)) as zdb:
return _fetch.run(
zdb,
project=name,
storage_dir=storage,
scratch_dir=scratch,
email=email,
fetch_proxy=fetch_proxy,
limit=limit or None,
workers=workers,
)
stats = _hold(_go, hold=hold)
for k, v in stats.items():
typer.echo(f" {k}: {v}")
# ── VPN droplet lifecycle ───────────────────────────────────────
vpn_app = typer.Typer(no_args_is_help=True)
app.add_typer(
vpn_app,
name="vpn",
help="DigitalOcean SOCKS5 exit for fallback fetch (SSH tunnel, no sudo needed).",
)
@vpn_app.command(name="up")
def vpn_up(
region: str = typer.Option(
"", "--region", help="DO region slug. Defaults to PRISMA_VPN_REGION."
),
attach_zotero: bool = typer.Option(
True,
"--attach-zotero/--no-attach-zotero",
help="Also start a sidecar SSH tunnel on the stack gateway network "
"and rewrite Zotero's prefs.js to proxy through it, so the "
"fallback plugin downloads via the droplet.",
),
) -> None:
"""Provision the droplet, wait for dante, wire Zotero up.
End state: Zotero's fallback-PDF plugin downloads through
``socks5://fetch-proxy:1080``; stack's own `fetch` command pulls
through the same route via ``vpn.active()``. Zero client-side sudo.
"""
from prisma import vpn as _vpn
info = _vpn.up(region=region or None, attach_zotero=attach_zotero)
typer.echo(f" droplet: {info['name']} ({info['droplet_id']})")
typer.echo(f" region: {info['region']}")
typer.echo(f" public IP: {info['public_ip']}")
typer.echo(f" ssh key: {info['ssh_key']}")
typer.echo(f" proxy URL: {info['proxy_url']}")
if info.get("sidecar"):
typer.echo(
f" zotero proxy: {info['zotero_proxy']} (sidecar: {info['sidecar']})"
)
typer.echo("")
typer.echo("Next: `stack prisma fetch <project>` — tunnel opens automatically.")
@vpn_app.command(name="down")
def vpn_down() -> None:
"""Destroy the droplet and wipe local state."""
from prisma import vpn as _vpn
info = _vpn.down()
for k, v in info.items():
typer.echo(f" {k}: {v}")
@vpn_app.command(name="status")
def vpn_status() -> None:
"""Report whether a PRISMA VPN droplet is currently tracked."""
from prisma import vpn as _vpn
for k, v in _vpn.status().items():
typer.echo(f" {k}: {v}")
@vpn_app.command(name="attach-zotero")
def vpn_attach_zotero() -> None:
"""Start the sidecar tunnel + rewire Zotero prefs.
Use when the droplet is already up but Zotero was never attached,
or to re-attach after a zotero container rebuild wiped prefs.
"""
import json
from prisma import vpn as _vpn
if not _vpn._DROPLET_JSON.is_file(): # noqa: SLF001
raise typer.BadParameter(
"no droplet tracked — run `stack prisma vpn up` first."
)
meta = json.loads(_vpn._DROPLET_JSON.read_text()) # noqa: SLF001
_vpn.attach_zotero_proxy(meta["public_ip"])
typer.echo(" sidecar started on gateway network as `fetch-proxy`")
typer.echo(" zotero prefs patched → socks5://fetch-proxy:1080")
typer.echo(" zotero restarted")
@vpn_app.command(name="detach-zotero")
def vpn_detach_zotero() -> None:
"""Stop the sidecar and restore Zotero's pre-proxy prefs."""
from prisma import vpn as _vpn
_vpn.detach_zotero_proxy()
typer.echo(" sidecar removed; prefs restored; zotero restarted.")
@vpn_app.command(name="verify")
def vpn_verify() -> None:
"""Open the tunnel, curl ifconfig.co, print egress, close.
Confirms the droplet is routable and reports its public IP /
country — the touchless way to sanity-check before a long fetch
run.
"""
from prisma import vpn as _vpn
with _vpn.active() as url:
info = _vpn.verify_egress(url)
typer.echo(f" proxy: {url}")
typer.echo(f" egress: {info.get('ip', '?')}")
typer.echo(
f" country: {info.get('country', '?')} ({info.get('country_iso', '?')})"
)
typer.echo(f" asn: {info.get('asn_org', '?')}")
# ── Orchestrated run ────────────────────────────────────────────
@app.command(name="run")
def run_all(
name: str = typer.Argument(..., help="Project slug."),
limit: int = typer.Option(
0, "--limit", "-n", help="Same limit applied to every stage."
),
skip_fetch: bool = typer.Option(False, "--skip-fetch"),
skip_extract: bool = typer.Option(False, "--skip-extract"),
db: Path = typer.Option(_DEFAULT_ZOT, "--db"),
storage: Path = typer.Option(_DEFAULT_STORAGE, "--storage"),
hold: bool = typer.Option(True, "--hold/--no-hold"),
) -> None:
"""Orchestrate all PRISMA stages in order.
screen → fetch → eligible → extract → flow, all inside one zotero
hold. Fetch only pulls PDFs for non-excluded items; stage 2
exclusions are terminal by design.
"""
import contextlib
import os
from prisma import (
eligibility,
)
from prisma import (
extract as _extract,
)
from prisma import (
fetch as _fetch,
)
from prisma import (
flow as _flow,
)
from prisma import (
screen as _screen,
)
from prisma import (
vpn as _vpn,
)
from prisma.llm import make_provider
from prisma.project import load
from zot.db import Db
def _proxy_cm():
explicit = os.environ.get("PRISMA_FETCH_PROXY")
if explicit:
return contextlib.nullcontext(explicit)
if _vpn.status().get("status") == "up":
return _vpn.active()
return contextlib.nullcontext(None)
def _go() -> None:
provider = make_provider()
with Db(str(db)) as zdb:
project = load(zdb, name)
typer.echo("==> stage 2 (title/abstract screen)")
s2 = _screen.run(
zdb, provider, project, storage_dir=storage, limit=limit or None
)
typer.echo(f" {s2}")
if not skip_fetch:
typer.echo("==> fetch (PDFs for non-excluded items)")
with _proxy_cm() as fetch_proxy:
fetch_stats = _fetch.run(
zdb,
project=name,
storage_dir=storage,
scratch_dir=Path(".state/prisma-fetch"),
email=os.environ.get("UNPAYWALL_EMAIL", "dev@fhirworx.io"),
fetch_proxy=fetch_proxy,
limit=limit or None,
)
typer.echo(f" {fetch_stats}")
typer.echo("==> stage 3 (full-text eligibility)")
s3 = eligibility.run(
zdb, provider, project, storage_dir=storage, limit=limit or None
)
typer.echo(f" {s3}")
if not skip_extract:
typer.echo("==> stage 3+ (data extraction)")
se = _extract.run(
zdb, provider, project, storage_dir=storage, limit=limit or None
)
typer.echo(f" {se}")
typer.echo("==> flow diagram")
counts = _flow.count(zdb, name)
typer.echo(_flow.text_summary(counts))
_hold(_go, hold=hold)