49a2f6f723
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.
247 lines
8.8 KiB
Python
247 lines
8.8 KiB
Python
"""Stage 2 — title/abstract screening loop.
|
|
|
|
Walks the project's identified-but-not-yet-screened items, builds a
|
|
cache-friendly prompt (system + criteria + reasons cached, per-item
|
|
user message uncached), calls the LLM with ``SCREEN_TOOL`` forced as
|
|
the response format, and writes the decision back through
|
|
:mod:`prisma.ingest`.
|
|
|
|
Prompt caching strategy: Anthropic's breakpoint system lets us mark
|
|
the *last* message of a stable prefix with ``cache=True``. Subsequent
|
|
calls inside the cache TTL pay ~10% of input cost on that prefix. Our
|
|
prefix is the system prompt + the criteria + the reasons codebook —
|
|
all stable for the whole run. The per-item message is the only
|
|
tail-variable piece.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING
|
|
|
|
from prisma.export import load_item, to_markdown
|
|
from prisma.ingest import apply_screen_decision
|
|
from prisma.llm import LLMCall, LLMMessage, Provider
|
|
from prisma.project import Project
|
|
from prisma.tools import SCREEN_TOOL
|
|
|
|
if TYPE_CHECKING:
|
|
from zot.db import Db
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
# Kept to a single sentence deliberately: long, bulleted system
|
|
# prompts make qwen-family models served through Ollama's chat
|
|
# template emit an empty response (no text, no tool call) — the
|
|
# guidance bullets live in the first user message instead, which
|
|
# every provider handles.
|
|
_SYSTEM = (
|
|
"You are a systematic-review assistant running PRISMA screening. "
|
|
"Record each decision by calling the record_screening_decision tool."
|
|
)
|
|
|
|
_GUIDANCE = """\
|
|
## Screening rules
|
|
|
|
- Be rigorous, conservative, and explicit; cite the specific criterion
|
|
that drives each decision in the rationale field.
|
|
- Call the `record_screening_decision` tool exactly once per item.
|
|
- Favor `uncertain` when the abstract is ambiguous on population,
|
|
design, or outcomes — the full-text stage will resolve it. If no
|
|
abstract is available, judge from the title and default to
|
|
`uncertain` unless the title alone cleanly matches an exclusion
|
|
code.
|
|
- Use `exclude` only when a reason code cleanly applies.
|
|
- Reason codes must come from the codebook; do not invent codes.
|
|
- Themes are optional; 0-4 short kebab-case slugs.
|
|
"""
|
|
|
|
|
|
def _build_call(project: Project, item_markdown: str) -> LLMCall:
|
|
"""Assemble messages with the stable prefix cached."""
|
|
return LLMCall(
|
|
messages=[
|
|
LLMMessage(role="system", content=_SYSTEM),
|
|
LLMMessage(
|
|
role="user",
|
|
content=(
|
|
_GUIDANCE
|
|
+ "\n## Inclusion / exclusion criteria\n\n"
|
|
+ project.criteria
|
|
+ "\n\n## Exclusion reason codebook\n\n"
|
|
+ project.reasons
|
|
),
|
|
cache=True, # <-- cache breakpoint lands here
|
|
),
|
|
LLMMessage(role="user", content=item_markdown),
|
|
],
|
|
tools=[SCREEN_TOOL],
|
|
force_tool=SCREEN_TOOL.name,
|
|
max_tokens=1024,
|
|
temperature=0.0,
|
|
)
|
|
|
|
|
|
# ── Queue ────────────────────────────────────────────────────────
|
|
|
|
|
|
def _queue(db: Db, project: str, limit: int | None) -> list[int]:
|
|
"""IDs of items in the project that have NOT been screened yet.
|
|
|
|
Definition: tagged ``project:<name>``, NOT a PRISMA anchor
|
|
(``module:prisma``), and not already tagged with any
|
|
``stage:2-…`` / ``stage:3-…`` / ``stage:4-…`` marker.
|
|
"""
|
|
rows = db.con.execute(
|
|
"""
|
|
SELECT DISTINCT i.itemID
|
|
FROM items i
|
|
JOIN itemTags it ON i.itemID = it.itemID
|
|
JOIN tags t ON it.tagID = t.tagID
|
|
WHERE t.name = ?
|
|
AND i.itemID NOT IN (
|
|
SELECT it2.itemID
|
|
FROM itemTags it2
|
|
JOIN tags t2 ON it2.tagID = t2.tagID
|
|
WHERE t2.name LIKE 'stage:2-%'
|
|
OR t2.name LIKE 'stage:3-%'
|
|
OR t2.name LIKE 'stage:4-%'
|
|
OR t2.name = 'module:prisma'
|
|
)
|
|
ORDER BY i.itemID
|
|
"""
|
|
+ (f" LIMIT {int(limit)}" if limit else ""),
|
|
(f"project:{project}",),
|
|
).fetchall()
|
|
return [r[0] for r in rows]
|
|
|
|
|
|
# ── Runner ───────────────────────────────────────────────────────
|
|
|
|
|
|
def run(
|
|
db: Db,
|
|
provider: Provider,
|
|
project: Project,
|
|
*,
|
|
storage_dir: Path,
|
|
limit: int | None = None,
|
|
progress: bool = True,
|
|
) -> dict[str, int]:
|
|
"""Screen every unscreened item, commit once per decision.
|
|
|
|
Commits per-item (not per-batch) so crashes only lose the
|
|
currently-in-flight call — previously-screened items stay tagged.
|
|
"""
|
|
ids = _queue(db, project.name, limit)
|
|
stats = {"screened": 0, "include": 0, "exclude": 0, "uncertain": 0, "errors": 0}
|
|
|
|
for i, zot_id in enumerate(ids, 1):
|
|
try:
|
|
snap = load_item(db, zot_id, storage_dir)
|
|
md = to_markdown(snap, include_fulltext=False)
|
|
call = _build_call(project, md)
|
|
result = provider.complete(call)
|
|
if not result.tool_calls:
|
|
log.warning("no tool_call for item %s — skipping", zot_id)
|
|
stats["errors"] += 1
|
|
continue
|
|
payload = result.tool_calls[0].get("input") or {}
|
|
apply_screen_decision(db, zot_id, project.name, payload)
|
|
db.commit()
|
|
decision = (payload.get("decision") or "uncertain").lower()
|
|
stats[decision] = stats.get(decision, 0) + 1
|
|
stats["screened"] += 1
|
|
if progress and i % 10 == 0:
|
|
print(
|
|
f" screened {i}/{len(ids)} "
|
|
f"(inc={stats['include']} exc={stats['exclude']} "
|
|
f"unc={stats['uncertain']})"
|
|
)
|
|
except Exception as e: # noqa: BLE001 — one item failing doesn't kill the run
|
|
log.warning("screen failed on item %s: %s", zot_id, e)
|
|
stats["errors"] += 1
|
|
continue
|
|
|
|
return stats
|
|
|
|
|
|
def run_fleet(
|
|
db: Db,
|
|
project: Project,
|
|
*,
|
|
hosts: list[str],
|
|
model: str,
|
|
storage_dir: Path,
|
|
num_ctx: int = 8192,
|
|
per_host: int = 1,
|
|
limit: int | None = None,
|
|
progress: bool = True,
|
|
) -> dict[str, int]:
|
|
"""Screen in parallel across a fleet of Ollama hosts (#655).
|
|
|
|
LLM calls fan out on a thread pool with at most ``per_host``
|
|
in-flight requests per host (Ollama serializes per model anyway);
|
|
every SQLite touch — the item loads and the decision writes —
|
|
stays on the calling thread, so ``db``'s single connection is
|
|
never shared across threads. Per-item commits keep the crash
|
|
story identical to :func:`run`: a dead host or dead run loses at
|
|
most its in-flight items, which stay queued for the next pass.
|
|
|
|
The pool driver itself lives in :mod:`prisma.fleet` — shared with
|
|
:func:`prisma.eligibility.run_fleet` and
|
|
:func:`prisma.extract.run_fleet` — this function only supplies the
|
|
per-item prompt and decision-apply glue.
|
|
"""
|
|
from prisma.fleet import run_pool
|
|
from prisma.llm import LLMResult, OllamaNativeProvider
|
|
|
|
providers = {
|
|
h.rstrip("/"): OllamaNativeProvider(host=h, model=model, num_ctx=num_ctx)
|
|
for h in hosts
|
|
}
|
|
ids = _queue(db, project.name, limit)
|
|
stats = {"screened": 0, "include": 0, "exclude": 0, "uncertain": 0, "errors": 0}
|
|
|
|
def prepare(zot_id: int):
|
|
snap = load_item(db, zot_id, storage_dir)
|
|
return _build_call(project, to_markdown(snap, include_fulltext=False))
|
|
|
|
def apply(zot_id: int, result: LLMResult) -> None:
|
|
if not result.tool_calls:
|
|
log.warning("no tool_call for item %s — skipping", zot_id)
|
|
stats["errors"] += 1
|
|
return
|
|
payload = result.tool_calls[0].get("input") or {}
|
|
apply_screen_decision(db, zot_id, project.name, payload)
|
|
db.commit()
|
|
decision = (payload.get("decision") or "uncertain").lower()
|
|
stats[decision] = stats.get(decision, 0) + 1
|
|
stats["screened"] += 1
|
|
if progress and stats["screened"] % 10 == 0:
|
|
print(
|
|
f" screened {stats['screened']}/{len(ids)} "
|
|
f"(inc={stats['include']} exc={stats['exclude']} "
|
|
f"unc={stats['uncertain']})"
|
|
)
|
|
|
|
def on_error(zot_id: int, e: Exception, stage: str) -> None:
|
|
if stage == "prepare":
|
|
log.warning("load failed on item %s: %s", zot_id, e)
|
|
else:
|
|
log.warning("screen failed on item %s: %s", zot_id, e)
|
|
stats["errors"] += 1
|
|
|
|
run_pool(
|
|
ids,
|
|
providers,
|
|
per_host=per_host,
|
|
prepare=prepare,
|
|
apply=apply,
|
|
on_error=on_error,
|
|
)
|
|
|
|
return stats
|