File size: 6,736 Bytes
98bde72 | 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 | """Bounded execution of explicit controller decisions against a tool registry."""
from __future__ import annotations
import json
import math
import copy
import subprocess
import tempfile
from pathlib import Path
from typing import Callable
from .schema import Decision, DesignSpec, ToolResult, canonical
from .trace import Trace, replay
class Tool:
def __init__(self, name: str, function: Callable, cost: dict[str,float],
description: str, argument_schema: dict | None = None):
if any(not math.isfinite(v) or v < 0 for v in cost.values()):
raise ValueError("tool costs must be nonnegative")
self.name,self.function,self.cost,self.description=name,function,cost,description
self.argument_schema=argument_schema or {}
def __call__(self, arguments, state):
if self.argument_schema.get("type"):
from jsonschema import validate
validate(arguments,self.argument_schema)
return ToolResult.model_validate(self.function(arguments,copy.deepcopy(state)))
class CommandTool(Tool):
"""Execute a preconfigured JSON-in/JSON-out worker in its own environment.
argv is supplied by the operator's registry, never by a controller decision.
Workers receive a request file and an output file as their last arguments.
"""
def __init__(self, name, argv, cwd, cost, description, timeout=1800, argument_schema=None):
self.argv=list(argv); self.cwd=Path(cwd).resolve(); self.timeout=timeout
super().__init__(name,self.execute,cost,description,argument_schema)
def execute(self, arguments, state):
with tempfile.TemporaryDirectory(prefix="peppa-job-") as directory:
request=Path(directory)/"request.json";output=Path(directory)/"result.json"
request.write_text(canonical({"arguments":arguments,"state":state}))
completed=subprocess.run(self.argv+[str(request),str(output)],cwd=self.cwd,
capture_output=True,text=True,timeout=self.timeout,check=False)
if completed.returncode:
raise RuntimeError(f"worker exit {completed.returncode}: {completed.stderr[-2000:]}")
if not output.exists():
raise RuntimeError("worker did not write its declared result")
return json.loads(output.read_text())
class Engine:
def __init__(self, spec: DesignSpec, registry: dict[str,Tool], trace_path: str | Path):
self.registry=registry;self.trace=Trace(trace_path)
if not self.trace.events:
self.trace.append("initialize",spec.model_dump(mode="json"))
self.state=replay(self.trace.path)
if self.state["spec"] != spec.model_dump(mode="json"):
raise ValueError("resume specification differs from original episode")
def reserve(self, name: str, cost: dict[str,float]):
limits=self.state["spec"]["budgets"]
for resource,amount in cost.items():
if not math.isfinite(amount) or amount<0 or resource not in limits or self.state["spent"].get(resource,0)+amount>limits[resource]:
raise ValueError(f"budget exceeded or undefined: {resource}")
self.trace.append("reserve",{"tool":name,"cost":cost})
self.state=replay(self.trace.path)
def apply(self, decision: Decision):
if self.state["stopped"]:
raise ValueError("episode already stopped")
unknown=set(decision.evidence_ids)-set(self.state["evidence"])
if unknown:
raise ValueError(f"unknown evidence IDs: {sorted(unknown)}")
self.trace.append("decision",decision.model_dump(mode="json"))
if decision.stop:
self.trace.append("stop",{"summary":decision.decision_summary})
elif decision.tool=="revise_plan":
# Scientific endpoints and budgets are immutable. Planning weights and
# questions may change, with the visible evidence retained in the trace.
allowed={"weights","motifs","questions","next_assays","hypotheses"}
if set(decision.arguments)-allowed:
raise ValueError("revision attempts to change a fixed scientific requirement")
weights=decision.arguments.get("weights")
if weights is not None:
if not isinstance(weights,dict) or not weights or any(float(v)<0 for v in weights.values()) or sum(weights.values())<=0:
raise ValueError("invalid objective weights")
names={r["endpoint"] for r in self.state["spec"]["requirements"]}
if set(weights)-names:
raise ValueError("unknown objective name")
self.reserve("revise_plan",{"tool_calls":1})
self.trace.append("revise",decision.arguments)
else:
if decision.tool not in self.registry:
raise ValueError("tool is not registered")
tool=self.registry[decision.tool]
self.reserve(tool.name,tool.cost)
try:
result=tool(decision.arguments,self.state)
known=set(self.state["candidates"])|{c.id for c in result.candidates}
if any(m.candidate_id not in known for m in result.measurements):
raise ValueError("measurement references an unknown candidate")
self.trace.append("result",result.model_dump(mode="json"))
except Exception as exc:
self.trace.append("error",{"tool":tool.name,"type":type(exc).__name__,"message":str(exc)})
self.state=replay(self.trace.path)
raise
self.state=replay(self.trace.path)
return self.state
def run(self, controller, max_steps=40):
for _ in range(max_steps):
if self.state["stopped"]:break
try:
cost={"controller_calls":1}
if getattr(controller,"token_reservation",0):cost["controller_tokens"]=controller.token_reservation
self.reserve("controller",cost)
except ValueError as exc:
self.trace.append("stop",{"summary":str(exc)});self.state=replay(self.trace.path);break
try:
decision,metadata=controller.next(copy.deepcopy(self.state),self.registry)
self.trace.append("controller",metadata)
self.apply(decision)
except Exception as exc:
self.trace.append("feedback",{"error":str(exc)})
self.state=replay(self.trace.path)
if not self.state["stopped"]:
self.trace.append("stop",{"summary":"controller step limit reached"})
self.state=replay(self.trace.path)
return self.state
|