File size: 6,673 Bytes
406a5e6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | """File-reference entry point that runs the PFB-MAX stack under the OFFICIAL
astabench harness (so submissions and live gpt-4o grading use the real path).
inspect eval astabench/paper_finder_validation \
--solver "<abs>/pfbmax/inspect_entry.py@pfbmax_solver" \
--model openai/gpt-4o-mini
Why a file shim (same reasoning as iris_asta/inspect_solver.py): inspect
discovers solvers by AST-scanning for a LITERAL top-level @solver function,
and its loader execs the file without registering it in sys.modules — so this
module stays dataclass-free at top level and pulls the implementation in
through ordinary imports.
Corpus access: the harness hands the solver its own date-wrapped MCP tools,
but those are async and PFB is search-heavy (the IRIS project measured the
async tool bridge deadlocking on this task). We therefore use a direct
AstaClient like IRIS does, and enforce the sample's snapshot date ourselves
from state.metadata["insertion_date"] — the same contract the task tools
enforce, so results stay date-legal.
"""
import json
import os
import sys
from inspect_ai.model import ChatMessageAssistant
from inspect_ai.solver import solver
_HERE = os.path.dirname(os.path.abspath(__file__))
_BUNDLE = os.path.dirname(_HERE)
for _p in (_HERE, os.path.join(_BUNDLE, "iris_asta")):
if _p not in sys.path:
sys.path.insert(0, _p)
def _load_env() -> None:
env_path = os.path.join(_BUNDLE, "iris_asta", ".env")
if not os.path.exists(env_path):
return
for line in open(env_path, encoding="utf-8"):
line = line.strip()
if "=" in line and not line.startswith("#"):
k, v = line.split("=", 1)
os.environ.setdefault(k, v.strip().strip('"').strip("'"))
def _extract_query(state) -> str:
meta = getattr(state, "metadata", None) or {}
raw = meta.get("raw_query")
if isinstance(raw, str) and raw.strip():
return raw.strip()
text = str(getattr(state, "input_text", "") or "")
marker = "find papers relevant to the following query"
low = text.lower()
idx = low.find(marker)
if idx >= 0:
text = text[idx + len(marker):]
return text.strip().strip(":").strip()
@solver
def pfbmax_solver(**kwargs):
"""PFB-MAX: compile -> route -> (metadata | specific | semantic) -> emit."""
_load_env()
# Fail fast on missing credentials. Every stage below soft-fails by design
# (an exception must never cost a sample), which means a missing key would
# otherwise produce a complete, hours-long, all-zero run with no diagnostic.
if not (os.environ.get("OPENAI_API_KEY") or "").strip():
raise RuntimeError(
"OPENAI_API_KEY is not set. Export it, or copy "
"iris_asta/.env.example to iris_asta/.env and fill it in.")
if not (os.environ.get("ASTA_TOOL_KEY") or "").strip():
raise RuntimeError(
"ASTA_TOOL_KEY is not set: corpus search would run against the "
"anonymous rate-limit pool and score near zero. Request a free key "
"and put it in iris_asta/.env (see iris_asta/.env.example).")
os.environ.setdefault("IRIS_ASTA_MCP_DEADLINE_S", "15")
os.environ.setdefault("IRIS_ASTA_TIMEOUT_S", "45")
os.environ.setdefault("IRIS_ASTA_RATE_LIMIT_RPS", "2.0")
async def solve(state, generate):
import asyncio
def _run():
# inspect's loader execs this file without a stable __file__/sys.path
# contract, so re-assert the import roots inside the worker.
here = os.environ.get("PFBMAX_DIR") or _HERE
bundle = os.path.dirname(here)
for p in (here, os.path.join(bundle, "iris_asta")):
if p not in sys.path:
sys.path.insert(0, p)
from iris_asta.asta_client import AstaClient
from iris_asta.config import load_config
import router
from costmeter import CostMeter
from llm import LLM
try:
from corpus_cache import CachedClient
except Exception:
CachedClient = None
meta = getattr(state, "metadata", None) or {}
inserted_before = meta.get("insertion_date") or "2025-06-01"
query = _extract_query(state)
# PFBMAX_USE_MCP=1: force the Asta MCP gateway. Measured
# 2026-08-18: api.semanticscholar.org hard-429s this key for
# >24h while asta-tools.allen.ai/mcp answers in 0.8s.
_mcp = bool((os.environ.get("PFBMAX_USE_MCP") or "").strip())
if _mcp:
# MCP truncates responses over ~400KB; giant citation pages
# (limit 1000) always truncate -> exception -> bisect ->
# truncate again, burning the whole 48-call budget at ~60s a
# call. Smaller pages never truncate. In-process override
# only; iris_asta itself stays untouched on disk.
try:
import iris_asta.solvers.pfb as _pfb
_pfb._CITATION_LIMIT = min(
getattr(_pfb, "_CITATION_LIMIT", 1000), 200)
except Exception:
pass
client = AstaClient(load_config(),
use_mcp=True if _mcp else None)
if CachedClient is not None:
client = CachedClient(client)
meter = CostMeter()
llm = LLM(meter=meter)
trace = {}
try:
results = router.solve(query, client, llm, inserted_before,
trace=trace)
except Exception as exc: # never fail the sample
trace["fatal"] = repr(exc)[:300]
results = []
return results, meter.total_usd(), trace
results, usd, trace = await asyncio.to_thread(_run)
payload = {"output": {"results": [
{"paper_id": str(pid), "markdown_evidence": ev}
for pid, ev in results]}}
completion = json.dumps(payload, ensure_ascii=True)
state.output.completion = completion
state.messages.append(ChatMessageAssistant(content=completion))
try:
state.metadata["pfbmax_usd"] = round(usd, 6)
state.metadata["pfbmax_route"] = trace.get("route")
state.metadata["pfbmax_n"] = len(results)
except Exception:
pass
return state
return solve
|