damaros 0.1.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- damaros/__init__.py +27 -0
- damaros/_version.py +1 -0
- damaros/adapters/__init__.py +6 -0
- damaros/adapters/__main__.py +169 -0
- damaros/adapters/inspect.py +148 -0
- damaros/adapters/nemo_gym.py +100 -0
- damaros/adapters/openenv.py +50 -0
- damaros/adapters/openenv_mcp.py +62 -0
- damaros/adapters/seat.py +100 -0
- damaros/adapters/verifiers.py +48 -0
- damaros/adapters/verifiers_v1.py +95 -0
- damaros/client.py +87 -0
- damaros/env.py +132 -0
- damaros/first_run.py +92 -0
- damaros/host.py +64 -0
- damaros/host_stub.py +64 -0
- damaros/login.py +70 -0
- damaros/mcp.py +85 -0
- damaros/protocol.py +41 -0
- damaros/registry.json +2534 -0
- damaros/resident_env.py +141 -0
- damaros/server.py +693 -0
- damaros/spaces.py +143 -0
- damaros/test_client.py +39 -0
- damaros/test_env.py +576 -0
- damaros/test_resident.py +294 -0
- damaros/test_server.py +330 -0
- damaros/test_trauma.py +220 -0
- damaros/test_trauma_env.py +147 -0
- damaros/trauma.py +188 -0
- damaros/trauma_baselines.py +94 -0
- damaros/trauma_env.py +210 -0
- damaros/verify_record.py +183 -0
- damaros-0.1.0.dist-info/METADATA +817 -0
- damaros-0.1.0.dist-info/RECORD +40 -0
- damaros-0.1.0.dist-info/WHEEL +5 -0
- damaros-0.1.0.dist-info/entry_points.txt +2 -0
- damaros-0.1.0.dist-info/licenses/LICENSE +202 -0
- damaros-0.1.0.dist-info/top_level.txt +2 -0
- damaros_verifiers.py +4 -0
damaros/__init__.py
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
"""The ICU's seats from Python, over the local servers the CLI already runs. Nothing is reimplemented: the tools, the clock, the reward and the cost are the TypeScript
|
|
2
|
+
harness's. No runtime dependencies; Gymnasium is used when installed.
|
|
3
|
+
|
|
4
|
+
Phase 2 U4: `from damaros import make` — one SDK shape over both worlds.
|
|
5
|
+
"""
|
|
6
|
+
from ._version import __version__
|
|
7
|
+
from .client import WorldClient, make
|
|
8
|
+
from .env import BedsideEnv, UnitEnv
|
|
9
|
+
from .host import ResidentHost, serve_host
|
|
10
|
+
from .resident_env import HostLost, ResidentEnv, ResidentLanes
|
|
11
|
+
from .spaces import EventsSpace, SceneSpace, ToolCallSpace
|
|
12
|
+
|
|
13
|
+
__all__ = [
|
|
14
|
+
"__version__",
|
|
15
|
+
"make",
|
|
16
|
+
"WorldClient",
|
|
17
|
+
"BedsideEnv",
|
|
18
|
+
"UnitEnv",
|
|
19
|
+
"ResidentEnv",
|
|
20
|
+
"ResidentLanes",
|
|
21
|
+
"ResidentHost",
|
|
22
|
+
"serve_host",
|
|
23
|
+
"HostLost",
|
|
24
|
+
"ToolCallSpace",
|
|
25
|
+
"EventsSpace",
|
|
26
|
+
"SceneSpace",
|
|
27
|
+
]
|
damaros/_version.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
1
|
+
__version__ = "0.1.0"
|
|
@@ -0,0 +1,6 @@
|
|
|
1
|
+
"""Thin adapters from the ICU's physician seat to the three lab stacks (OpenEnv, verifiers, Inspect). Each is under 200 lines, wraps one `Seat`, and takes its reward
|
|
2
|
+
from the server (`end_case`): none computes a reward, none can ask for a test or evaluation case (the local server it wraps serves the sample and train splits only).
|
|
3
|
+
The reward and the harm cost travel as separate fields in every adapter and are never summed. `python3 -m damaros.adapters smoke <name>` plays one episode of one."""
|
|
4
|
+
from .seat import Seat
|
|
5
|
+
|
|
6
|
+
__all__ = ["Seat"]
|
|
@@ -0,0 +1,169 @@
|
|
|
1
|
+
"""python3 -m damaros.adapters smoke <openenv|openenv_mcp|verifiers|verifiers_v1|nemo_gym|inspect>: plays one patient through that adapter with a scripted policy (no model, no GPU, no network)
|
|
2
|
+
and prints the server's reward and cost, apart."""
|
|
3
|
+
import asyncio
|
|
4
|
+
import json
|
|
5
|
+
import sys
|
|
6
|
+
|
|
7
|
+
def summary(credit):
|
|
8
|
+
"""What the credit comes to, in two numbers every adapter can carry: how many acting calls were credited, and what they did to the chance of death in all."""
|
|
9
|
+
credit = credit or []
|
|
10
|
+
return {"calls": len(credit), "net_p_death": round(sum(k["effect"]["p_death"] for k in credit), 6)}
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
SMOKE_SEED = "smoke-182" # a patient the scripted policy harms twice: a reward of -0.1 and a cost with two kinds in it, so an adapter that dropped either could not pass
|
|
14
|
+
POLICY = [("order_lab", {"test": "blood_culture_x2", "priority": "stat"}), ("order_med", {"drug": "cefepime", "priority": "stat"})] + [("wait", {"minutes": 240})] * 200 # waits end the patient at the horizon
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
def openenv(seed):
|
|
18
|
+
from .openenv import IcuEnvironment, ToolCall
|
|
19
|
+
|
|
20
|
+
env = IcuEnvironment(size=14)
|
|
21
|
+
env.reset(seed=seed)
|
|
22
|
+
for name, args in POLICY:
|
|
23
|
+
obs = env.step(ToolCall(name=name, arguments=args))
|
|
24
|
+
if obs.done:
|
|
25
|
+
break
|
|
26
|
+
env.close()
|
|
27
|
+
return obs.reward, obs.cost, summary(obs.credit)
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def verifiers(seed):
|
|
31
|
+
from .verifiers import IcuEnv
|
|
32
|
+
|
|
33
|
+
env = IcuEnv(seeds=[seed], size=14)
|
|
34
|
+
state = {"info": {"seed": str(seed)}}
|
|
35
|
+
asyncio.run(env.setup_state(state))
|
|
36
|
+
for i, (name, args) in enumerate(POLICY):
|
|
37
|
+
msg = {"role": "assistant", "content": "", "tool_calls": [{"id": f"c{i}", "type": "function", "function": {"name": name, "arguments": json.dumps(args)}}]}
|
|
38
|
+
asyncio.run(env.env_response([msg], state))
|
|
39
|
+
if state["server_reward"] is not None:
|
|
40
|
+
break
|
|
41
|
+
state["seat"].close()
|
|
42
|
+
return state["server_reward"], state["cost"], summary(state["credit"])
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
def inspect(seed):
|
|
46
|
+
from inspect_ai.solver import TaskState
|
|
47
|
+
from .inspect import KEY, tools_for
|
|
48
|
+
from .seat import Seat
|
|
49
|
+
|
|
50
|
+
seat = Seat(size=14)
|
|
51
|
+
seat.reset(seed)
|
|
52
|
+
|
|
53
|
+
class Store:
|
|
54
|
+
data = {}
|
|
55
|
+
set = lambda self, k, v: self.data.__setitem__(k, v) # noqa: E731
|
|
56
|
+
|
|
57
|
+
state = type("S", (), {"store": Store()})()
|
|
58
|
+
tools = {d.name: d.tool for d in tools_for(seat, state)}
|
|
59
|
+
for name, args in POLICY:
|
|
60
|
+
asyncio.run(tools[name](**args))
|
|
61
|
+
if KEY in state.store.data:
|
|
62
|
+
break
|
|
63
|
+
seat.close()
|
|
64
|
+
r = state.store.data.get(KEY, {})
|
|
65
|
+
return r.get("reward"), r.get("cost"), summary(r.get("credit"))
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
def openenv_mcp(seed):
|
|
69
|
+
from openenv.core.env_server.mcp_types import CallToolAction, ListToolsAction
|
|
70
|
+
|
|
71
|
+
from .openenv_mcp import IcuMCPEnvironment
|
|
72
|
+
|
|
73
|
+
env = IcuMCPEnvironment(size=14)
|
|
74
|
+
env.reset(seed=seed)
|
|
75
|
+
assert env.step(ListToolsAction()).tools, "the MCP form lists the seat's tools"
|
|
76
|
+
for name, args in POLICY:
|
|
77
|
+
obs = env.step(CallToolAction(tool_name=name, arguments=args))
|
|
78
|
+
if obs.done:
|
|
79
|
+
break
|
|
80
|
+
env.close()
|
|
81
|
+
return obs.reward, (obs.metadata or {}).get("cost"), summary((obs.metadata or {}).get("credit"))
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def verifiers_v1(seed):
|
|
85
|
+
from damaros.adapters.verifiers_v1 import IcuEnv, IcuEnvConfig
|
|
86
|
+
|
|
87
|
+
env = IcuEnv(IcuEnvConfig(taskset={"id": "damaros_verifiers", "seeds": [seed], "size": 14}, max_turns=len(POLICY)))
|
|
88
|
+
(task,) = list(env.taskset)
|
|
89
|
+
agents, trace = scripted_agents(POLICY)
|
|
90
|
+
asyncio.run(env.run(task, agents))
|
|
91
|
+
m = trace.metrics
|
|
92
|
+
return trace.rewards["damaros_reward"], {"severe": m["damaros_cost_severe"], "other": m["damaros_cost_other"], "by_kind": {k[len("damaros_cost_"):]: v for k, v in m.items() if k.startswith("damaros_cost_") and k not in ("damaros_cost_severe", "damaros_cost_other")}}, {"calls": int(m["damaros_credit_calls"]), "net_p_death": round(m["damaros_credit_net_p_death"], 6)}
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
def scripted_agents(policy):
|
|
96
|
+
"""A verifiers v1 `agents` whose player replies with one JSON tool call a turn from `policy` ([(name, arguments)]) and then stops, with the trace it records into: the model's place, taken by a list,
|
|
97
|
+
so `IcuEnv.run` is played with no model. Returns (agents, trace); `trace.rewards` and `trace.metrics` hold what the env recorded."""
|
|
98
|
+
|
|
99
|
+
class Trace:
|
|
100
|
+
def __init__(self):
|
|
101
|
+
self.rewards, self.metrics = {}, {}
|
|
102
|
+
|
|
103
|
+
def record_reward(self, k, v, w=1.0):
|
|
104
|
+
self.rewards[k] = v
|
|
105
|
+
|
|
106
|
+
def record_metric(self, k, v):
|
|
107
|
+
self.metrics[k] = v
|
|
108
|
+
|
|
109
|
+
class Reply:
|
|
110
|
+
def __init__(self, text, terminated):
|
|
111
|
+
self.last_reply, self.terminated = text, terminated
|
|
112
|
+
|
|
113
|
+
class Player:
|
|
114
|
+
def __init__(self):
|
|
115
|
+
self.trace, self.calls = Trace(), iter(policy)
|
|
116
|
+
|
|
117
|
+
async def turn(self, _text):
|
|
118
|
+
call = next(self.calls, None)
|
|
119
|
+
return Reply(json.dumps({"name": call[0], "arguments": call[1]}) if call else "", call is None)
|
|
120
|
+
|
|
121
|
+
async def __aenter__(self):
|
|
122
|
+
return self
|
|
123
|
+
|
|
124
|
+
async def __aexit__(self, *_):
|
|
125
|
+
return False
|
|
126
|
+
|
|
127
|
+
player = Player()
|
|
128
|
+
agents = type("Agents", (), {"player": type("P", (), {"interaction": lambda self, _task: player})()})()
|
|
129
|
+
return agents, player.trace
|
|
130
|
+
|
|
131
|
+
|
|
132
|
+
def nemo_gym(seed):
|
|
133
|
+
from unittest.mock import MagicMock
|
|
134
|
+
|
|
135
|
+
from fastapi.testclient import TestClient
|
|
136
|
+
from nemo_gym.server_utils import ServerClient
|
|
137
|
+
|
|
138
|
+
from .nemo_gym import IcuResourcesServer, IcuResourcesServerConfig
|
|
139
|
+
|
|
140
|
+
server = IcuResourcesServer(config=IcuResourcesServerConfig(host="0.0.0.0", port=8080, entrypoint="", name="", size=14), server_client=MagicMock(spec=ServerClient))
|
|
141
|
+
client = TestClient(server.setup_webserver())
|
|
142
|
+
client.post("/seed_session", json={"seed": seed})
|
|
143
|
+
for name, args in POLICY:
|
|
144
|
+
if client.post(f"/{name}", json=args).json()["done"]:
|
|
145
|
+
break
|
|
146
|
+
response = {"id": "smoke", "created_at": 0.0, "model": "scripted", "object": "response", "output": [], "parallel_tool_calls": False, "tool_choice": "auto", "tools": []}
|
|
147
|
+
verified = client.post("/verify", json={"responses_create_params": {"input": []}, "response": response}).json()
|
|
148
|
+
return verified["reward"], verified["cost"], summary(verified.get("credit"))
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
ADAPTERS = ("openenv", "openenv_mcp", "verifiers", "verifiers_v1", "nemo_gym", "inspect")
|
|
152
|
+
EXTRA = {"openenv_mcp": "openenv"} # the MCP form shares the openenv extra
|
|
153
|
+
|
|
154
|
+
|
|
155
|
+
def main(argv):
|
|
156
|
+
if len(argv) != 2 or argv[0] != "smoke" or argv[1] not in ADAPTERS:
|
|
157
|
+
print(f"usage: python3 -m damaros.adapters smoke <{'|'.join(ADAPTERS)}>", file=sys.stderr)
|
|
158
|
+
return 2
|
|
159
|
+
try:
|
|
160
|
+
reward, cost, credit = globals()[argv[1]](SMOKE_SEED)
|
|
161
|
+
except ImportError as e:
|
|
162
|
+
print(f"{argv[1]} is not installed ({e}). pip install 'damaros[{EXTRA.get(argv[1], argv[1])}]'", file=sys.stderr)
|
|
163
|
+
return 2
|
|
164
|
+
print(json.dumps({"adapter": argv[1], "reward": reward, "cost": cost, "credit": credit}))
|
|
165
|
+
return 0 if reward is not None else 1
|
|
166
|
+
|
|
167
|
+
|
|
168
|
+
if __name__ == "__main__":
|
|
169
|
+
sys.exit(main(sys.argv[1:]))
|
|
@@ -0,0 +1,148 @@
|
|
|
1
|
+
"""Inspect adapter: a `Task` over the sample and train splits with a `Scorer` that reads the server's score. The solver gives the model the physician's tools and plays them on the
|
|
2
|
+
seat; the scorer returns the server's reward as its value and the harm cost in its metadata, apart, and never sums them. Needs `pip install damaros[inspect]`."""
|
|
3
|
+
import json
|
|
4
|
+
|
|
5
|
+
from inspect_ai import Task, task
|
|
6
|
+
from inspect_ai.dataset import MemoryDataset, Sample
|
|
7
|
+
from inspect_ai.model import ChatMessageUser
|
|
8
|
+
from inspect_ai.scorer import Score, Target, mean, scorer
|
|
9
|
+
from inspect_ai.solver import Generate, TaskState, solver
|
|
10
|
+
from inspect_ai.tool import ToolDef, ToolParam, ToolParams
|
|
11
|
+
|
|
12
|
+
from .seat import Seat
|
|
13
|
+
|
|
14
|
+
KEY = "damaros"
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
def params_of(schema):
|
|
18
|
+
"""Inspect wants a description for every parameter; the server's schema has one only where the tool says something about it, so the others are named by the parameter."""
|
|
19
|
+
props = {k: ToolParam.model_validate({"description": k.replace("_", " "), **v}) for k, v in schema.get("properties", {}).items()}
|
|
20
|
+
return ToolParams(properties=props, required=schema.get("required", []))
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def tools_for(seat, state):
|
|
24
|
+
"""One Inspect tool per tool the server lists; a call runs on the seat, and the patient's end writes the server's reward and cost into the sample's store."""
|
|
25
|
+
defs = []
|
|
26
|
+
for spec in seat.tool_list:
|
|
27
|
+
schema = spec["input_schema"]
|
|
28
|
+
|
|
29
|
+
def make(name):
|
|
30
|
+
async def run(**arguments):
|
|
31
|
+
text, reward, done = seat.call(name, arguments)
|
|
32
|
+
if done:
|
|
33
|
+
state.store.set(KEY, {"reward": reward, "cost": seat.cost, "credit": seat.credit, "budget": seat.budget})
|
|
34
|
+
return text
|
|
35
|
+
|
|
36
|
+
return run
|
|
37
|
+
|
|
38
|
+
defs.append(ToolDef(make(spec["name"]), name=spec["name"], description=spec.get("description") or spec["name"], parameters=params_of(schema)))
|
|
39
|
+
return defs
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
@solver
|
|
43
|
+
def physician(size=14, harm_budget=None, server=None):
|
|
44
|
+
async def solve(state: TaskState, generate: Generate) -> TaskState:
|
|
45
|
+
seat = Seat(size=size, harm_budget=harm_budget, server=server)
|
|
46
|
+
text, _ = seat.reset(state.metadata["seed"])
|
|
47
|
+
state.tools = [d.as_tool() for d in tools_for(seat, state)]
|
|
48
|
+
state.messages.append(ChatMessageUser(content=text))
|
|
49
|
+
try:
|
|
50
|
+
return await generate(state)
|
|
51
|
+
finally:
|
|
52
|
+
seat.close()
|
|
53
|
+
|
|
54
|
+
return solve
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
@scorer(metrics=[mean()])
|
|
58
|
+
def server_reward():
|
|
59
|
+
async def score(state: TaskState, target: Target) -> Score:
|
|
60
|
+
result = state.store.get(KEY) or {"reward": 0.0, "cost": None}
|
|
61
|
+
return Score(value=float(result["reward"] or 0.0), answer=json.dumps(result["cost"]), metadata={"cost": result["cost"], "credit": result.get("credit"), "budget": result.get("budget")})
|
|
62
|
+
|
|
63
|
+
return score
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
@task
|
|
67
|
+
def icu(seeds=range(8), size=14, harm_budget=None, server=None):
|
|
68
|
+
"""`server` (`-T server=https://...`, with the run's token in DAMAROS_RUN_TOKEN, never a -T argument): each sample plays the hosted run's next patient; `seeds` then only counts the samples. Run it with `--max-samples 1`: a run has one patient in hand at a time."""
|
|
69
|
+
return Task(dataset=MemoryDataset([Sample(input=f"Care for the patient on seed {s}.", metadata={"seed": str(s)}) for s in seeds]), solver=physician(size=size, harm_budget=harm_budget, server=server), scorer=server_reward())
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
@task
|
|
73
|
+
def icu_server(server, token, agent, is_=None, analysis_out=None):
|
|
74
|
+
"""The `--server` form (D4): submits an EVALUATION run through the hosted server and reads the run analysis back. The cases stay on the server: each sample is "the next patient",
|
|
75
|
+
nothing here can name one, and the sample's score is 0 (an evaluation patient has no outcome until the run is scored); the task's metric is the server's verdict, read from the customer's
|
|
76
|
+
analysis, which is also written to `analysis_out` when given. Run it with one sample at a time: `inspect eval ... --max-samples 1`.
|
|
77
|
+
|
|
78
|
+
inspect eval damaros.adapters.inspect@icu_server -T server=https://app.damaros.ai -T token=... -T agent=my-agent@1.2 --max-samples 1
|
|
79
|
+
"""
|
|
80
|
+
from ..server import Server
|
|
81
|
+
|
|
82
|
+
run = Server(server, token).open_evaluation_run(agent, is_)
|
|
83
|
+
seat = run.seat()
|
|
84
|
+
tools = [{"name": t["name"], "description": t.get("description", ""), "input_schema": t.get("inputSchema") or t.get("input_schema")} for t in seat.tool_list]
|
|
85
|
+
|
|
86
|
+
@solver
|
|
87
|
+
def play():
|
|
88
|
+
async def solve(state: TaskState, generate: Generate) -> TaskState:
|
|
89
|
+
patient = seat.next_case()
|
|
90
|
+
if patient is None:
|
|
91
|
+
return state
|
|
92
|
+
defs = []
|
|
93
|
+
for spec in tools:
|
|
94
|
+
schema = spec["input_schema"]
|
|
95
|
+
|
|
96
|
+
def make(name):
|
|
97
|
+
async def call(**arguments):
|
|
98
|
+
out = seat.call(name, arguments)
|
|
99
|
+
if out.get("ended"):
|
|
100
|
+
seat.end_case()
|
|
101
|
+
return json.dumps(out)
|
|
102
|
+
|
|
103
|
+
return call
|
|
104
|
+
|
|
105
|
+
defs.append(ToolDef(make(spec["name"]), name=spec["name"], description=spec["description"] or spec["name"], parameters=params_of(schema)))
|
|
106
|
+
state.tools = [d.as_tool() for d in defs]
|
|
107
|
+
state.messages.append(ChatMessageUser(content=json.dumps(patient)))
|
|
108
|
+
try:
|
|
109
|
+
return await generate(state)
|
|
110
|
+
finally:
|
|
111
|
+
seat.end_case() # hand the bed back whether or not the model did; a test patient shows no outcome
|
|
112
|
+
run.handed_back += 1
|
|
113
|
+
if run.handed_back >= (run.cases_planned or 1):
|
|
114
|
+
seat.next_case() # the panel is played: asking for another patient is what closes the run, and the server scores it
|
|
115
|
+
run.analysis() # wait here, once, for the score
|
|
116
|
+
|
|
117
|
+
return solve
|
|
118
|
+
|
|
119
|
+
@scorer(metrics=[server_verdict(run, analysis_out)])
|
|
120
|
+
def pending():
|
|
121
|
+
async def score(state: TaskState, target: Target) -> Score:
|
|
122
|
+
return Score(value=0.0, explanation="an evaluation patient has no outcome until the server scores the run")
|
|
123
|
+
|
|
124
|
+
return score
|
|
125
|
+
|
|
126
|
+
n = run.cases_planned or 1
|
|
127
|
+
return Task(dataset=MemoryDataset([Sample(input="Care for the next patient.", metadata={"n": i}) for i in range(n)]), solver=play(), scorer=pending())
|
|
128
|
+
|
|
129
|
+
|
|
130
|
+
def server_verdict(run, analysis_out=None):
|
|
131
|
+
"""A metric: 1.0 if the run met the account's thresholds and 0.0 if it did not, read from the customer's analysis once the server has scored the run (the solver of the last sample waits
|
|
132
|
+
for it; Inspect asks a metric after every sample, so this never waits). The analysis is written to `analysis_out` when given."""
|
|
133
|
+
from inspect_ai.scorer import metric
|
|
134
|
+
|
|
135
|
+
@metric
|
|
136
|
+
def damaros_verdict():
|
|
137
|
+
def compute(scores):
|
|
138
|
+
analysis = run.scored
|
|
139
|
+
if not scores or analysis is None:
|
|
140
|
+
return 0.0
|
|
141
|
+
if analysis_out:
|
|
142
|
+
with open(analysis_out, "w") as f:
|
|
143
|
+
json.dump(analysis, f, indent=2)
|
|
144
|
+
return 1.0 if analysis["verdict"]["result"] == "met" else 0.0
|
|
145
|
+
|
|
146
|
+
return compute
|
|
147
|
+
|
|
148
|
+
return damaros_verdict()
|
|
@@ -0,0 +1,100 @@
|
|
|
1
|
+
"""NeMo Gym adapter: a resources server whose tool routes run on the physician's seat and whose `/verify` returns the server's reward, so a NeMo Gym rollout collects Damaros episodes with no Damaros
|
|
2
|
+
code on the lab's side. One route per tool (`/get_chart`, `/order_med`, ...), each taking the tool's arguments as the body; `/seed_session` opens a patient from a seed; `/verify` returns the reward
|
|
3
|
+
as `end_case` gave it, with the harm cost beside it in `cost` (never added to `reward`). `tool_definitions()` is the Responses-API tool list to put in a dataset row's `responses_create_params`.
|
|
4
|
+
Written against nemo-gym 0.4.0 (2026-07-02, Apache-2.0, Python >=3.12); 0.5.0 and later need Python >=3.13.14, so a 3.12 environment stays on 0.4.0. Needs `pip install damaros[nemo_gym]`."""
|
|
5
|
+
from typing import Any, Dict, List, Optional
|
|
6
|
+
|
|
7
|
+
from fastapi import FastAPI, Request
|
|
8
|
+
from pydantic import BaseModel, Field
|
|
9
|
+
|
|
10
|
+
from nemo_gym.base_resources_server import (
|
|
11
|
+
BaseResourcesServerConfig,
|
|
12
|
+
BaseSeedSessionRequest,
|
|
13
|
+
BaseSeedSessionResponse,
|
|
14
|
+
BaseVerifyRequest,
|
|
15
|
+
BaseVerifyResponse,
|
|
16
|
+
SimpleResourcesServer,
|
|
17
|
+
)
|
|
18
|
+
from nemo_gym.server_utils import SESSION_ID_KEY
|
|
19
|
+
|
|
20
|
+
from .seat import Seat
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class IcuResourcesServerConfig(BaseResourcesServerConfig):
|
|
24
|
+
size: int = 14
|
|
25
|
+
clusters: Optional[List[str]] = None
|
|
26
|
+
harm_budget: Optional[str] = None
|
|
27
|
+
server: Optional[str] = None # the hosted server's address: each session plays the next patient of the run whose token is DAMAROS_RUN_TOKEN (the environment, never this config)
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
class IcuSeedSessionRequest(BaseSeedSessionRequest):
|
|
31
|
+
seed: str = "nemo_gym"
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
class ToolResult(BaseModel):
|
|
35
|
+
text: str
|
|
36
|
+
done: bool = False
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
class IcuVerifyResponse(BaseVerifyResponse):
|
|
40
|
+
cost: Optional[Dict[str, Any]] = None # the harm cost, apart from `reward`; None while the patient is still open
|
|
41
|
+
credit: Optional[list] = None # what each acting call did to the patient (the server's); None while the patient is still open
|
|
42
|
+
budget: Optional[Dict[str, Any]] = None
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
class IcuResourcesServer(SimpleResourcesServer):
|
|
46
|
+
config: IcuResourcesServerConfig
|
|
47
|
+
seats: Dict[str, Any] = Field(default_factory=dict) # session id -> Seat
|
|
48
|
+
rewards: Dict[str, Optional[float]] = Field(default_factory=dict)
|
|
49
|
+
|
|
50
|
+
def _seat_args(self) -> dict:
|
|
51
|
+
return {"size": self.config.size, "clusters": self.config.clusters, "harm_budget": self.config.harm_budget, "server": self.config.server}
|
|
52
|
+
|
|
53
|
+
def tool_definitions(self) -> List[dict]:
|
|
54
|
+
"""The seat's tools as Responses-API function tools, for a dataset row's `responses_create_params.tools`."""
|
|
55
|
+
probe = Seat(**self._seat_args())
|
|
56
|
+
tools = [{"type": "function", "name": t["name"], "description": t.get("description", ""), "parameters": t["input_schema"], "strict": False} for t in probe.tools()]
|
|
57
|
+
probe.close()
|
|
58
|
+
return tools
|
|
59
|
+
|
|
60
|
+
def setup_webserver(self) -> FastAPI:
|
|
61
|
+
app = super().setup_webserver()
|
|
62
|
+
for tool in self.tool_definitions():
|
|
63
|
+
app.post(f"/{tool['name']}")(self._route(tool["name"]))
|
|
64
|
+
return app
|
|
65
|
+
|
|
66
|
+
def _route(self, name: str):
|
|
67
|
+
async def call(request: Request, body: Dict[str, Any]) -> ToolResult:
|
|
68
|
+
sid = request.session[SESSION_ID_KEY]
|
|
69
|
+
if sid not in self.seats:
|
|
70
|
+
return ToolResult(text="No patient is open. Seed the session first.", done=False)
|
|
71
|
+
text, reward, done = self.seats[sid].call(name, body)
|
|
72
|
+
if done:
|
|
73
|
+
self.rewards[sid] = reward
|
|
74
|
+
return ToolResult(text=text, done=done)
|
|
75
|
+
|
|
76
|
+
call.__name__ = name
|
|
77
|
+
return call
|
|
78
|
+
|
|
79
|
+
async def seed_session(self, request: Request, body: IcuSeedSessionRequest) -> BaseSeedSessionResponse:
|
|
80
|
+
sid = request.session[SESSION_ID_KEY]
|
|
81
|
+
if sid in self.seats:
|
|
82
|
+
self.seats.pop(sid).close()
|
|
83
|
+
self.seats[sid] = Seat(**self._seat_args())
|
|
84
|
+
self.seats[sid].reset(body.seed)
|
|
85
|
+
self.rewards[sid] = None
|
|
86
|
+
return BaseSeedSessionResponse()
|
|
87
|
+
|
|
88
|
+
async def verify(self, request: Request, body: BaseVerifyRequest) -> IcuVerifyResponse:
|
|
89
|
+
"""The server's reward for the patient, as `end_case` gave it: zero for one never handed back. The seat closes here."""
|
|
90
|
+
sid = request.session[SESSION_ID_KEY]
|
|
91
|
+
seat = self.seats.pop(sid, None)
|
|
92
|
+
reward = self.rewards.pop(sid, None)
|
|
93
|
+
cost, credit, budget = (seat.cost, seat.credit, seat.budget) if seat is not None and reward is not None else (None, None, None)
|
|
94
|
+
if seat is not None:
|
|
95
|
+
seat.close()
|
|
96
|
+
return IcuVerifyResponse(**body.model_dump(), reward=float(reward or 0.0), cost=cost, credit=credit, budget=budget)
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
if __name__ == "__main__":
|
|
100
|
+
IcuResourcesServer.run_webserver()
|
|
@@ -0,0 +1,50 @@
|
|
|
1
|
+
"""OpenEnv adapter: an `Environment` with reset, step and state over the physician's seat. The observation is the seat's text plus the structured chart; the reward is
|
|
2
|
+
the server's, given on the step that ends the patient, and the harm cost is its own field. Needs `pip install damaros[openenv]`."""
|
|
3
|
+
from typing import Any, Optional
|
|
4
|
+
|
|
5
|
+
from openenv.core.env_server import Action, Environment, Observation, State
|
|
6
|
+
|
|
7
|
+
from .seat import Seat
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class ToolCall(Action):
|
|
11
|
+
name: str
|
|
12
|
+
arguments: dict[str, Any] = {}
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class IcuObservation(Observation):
|
|
16
|
+
text: str = ""
|
|
17
|
+
chart: Optional[dict[str, Any]] = None
|
|
18
|
+
cost: Optional[dict[str, Any]] = None # the harm cost, apart from `reward`; only when the patient ends
|
|
19
|
+
credit: Optional[list[dict[str, Any]]] = None # what each acting call did to the patient (the server's): the run with it minus the run without it, with the same case and seed; only when the patient ends
|
|
20
|
+
budget: Optional[dict[str, Any]] = None # with a harm budget set: {"file_hash", "breached": the rule that ended the patient, or null}
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class IcuState(State):
|
|
24
|
+
seed: str = ""
|
|
25
|
+
t: Optional[int] = None
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
class IcuEnvironment(Environment[ToolCall, IcuObservation, IcuState]):
|
|
29
|
+
def __init__(self, size: int = 14, clusters: Optional[list[str]] = None, harm_budget: Optional[str] = None, server: Optional[str] = None, run_token: Optional[str] = None):
|
|
30
|
+
"""`server` (and `run_token`, or DAMAROS_RUN_TOKEN): play a run on the hosted server, one of its patients per reset."""
|
|
31
|
+
super().__init__()
|
|
32
|
+
self.seat = Seat(size=size, clusters=clusters, harm_budget=harm_budget, server=server, run_token=run_token)
|
|
33
|
+
self._state = IcuState()
|
|
34
|
+
|
|
35
|
+
def reset(self, seed: Optional[int] = None, episode_id: Optional[str] = None, **kwargs: Any) -> IcuObservation:
|
|
36
|
+
text, _ = self.seat.reset(seed if seed is not None else "openenv")
|
|
37
|
+
self._state = IcuState(episode_id=episode_id, step_count=0, seed=str(seed))
|
|
38
|
+
return IcuObservation(text=text, chart=self.seat.chart(text), done=False, reward=None)
|
|
39
|
+
|
|
40
|
+
def step(self, action: ToolCall, timeout_s: Optional[float] = None, **kwargs: Any) -> IcuObservation:
|
|
41
|
+
text, reward, done = self.seat.call(action.name, action.arguments)
|
|
42
|
+
self._state.step_count += 1
|
|
43
|
+
return IcuObservation(text=text, chart=self.seat.chart(text), done=done, reward=reward, cost=self.seat.cost if done else None, credit=self.seat.credit if done else None, budget=self.seat.budget if done else None)
|
|
44
|
+
|
|
45
|
+
@property
|
|
46
|
+
def state(self) -> IcuState:
|
|
47
|
+
return self._state
|
|
48
|
+
|
|
49
|
+
def close(self) -> None:
|
|
50
|
+
self.seat.close()
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
"""OpenEnv's MCP form (RFC 003, "MCP Support: Traditional Tool Calling"): an `MCPEnvironment` whose FastMCP server carries the physician's tools, so a policy lists them with `ListToolsAction` and calls
|
|
2
|
+
them with `CallToolAction`, the same as at inference. The RFC file is on OpenEnv's main (merged with #169, 2025-11-24) although its own header still reads "In Review"; the classes it describes ship in
|
|
3
|
+
openenv-core 0.3.0 (2026-05-11), which this is written against. Each tool is the seat's, with the seat's own JSON schema. The call that ends the patient comes back with `done` true, the server's reward
|
|
4
|
+
in `reward` and the harm cost beside it in `metadata["cost"]`, never added to the reward. It sits beside the plain `Action` form in `openenv.py`. Needs `pip install damaros[openenv]`."""
|
|
5
|
+
from typing import Any, Optional
|
|
6
|
+
|
|
7
|
+
from fastmcp import FastMCP
|
|
8
|
+
from fastmcp.tools import Tool
|
|
9
|
+
from fastmcp.tools.base import ToolResult
|
|
10
|
+
from openenv.core.env_server.mcp_environment import MCPEnvironment
|
|
11
|
+
from openenv.core.env_server.mcp_types import CallToolAction
|
|
12
|
+
from openenv.core.env_server.types import Observation, State
|
|
13
|
+
|
|
14
|
+
from .seat import Seat
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class IcuMCPEnvironment(MCPEnvironment):
|
|
18
|
+
def __init__(self, size: int = 14, clusters: Optional[list[str]] = None, harm_budget: Optional[str] = None, server: Optional[str] = None, run_token: Optional[str] = None):
|
|
19
|
+
"""`server` (and `run_token`, or DAMAROS_RUN_TOKEN): play a run on the hosted server, one of its patients per reset."""
|
|
20
|
+
self.seat = Seat(size=size, clusters=clusters, harm_budget=harm_budget, server=server, run_token=run_token)
|
|
21
|
+
mcp = FastMCP("damaros-icu")
|
|
22
|
+
for spec in self.seat.tools(): # the tool list does not depend on the seed
|
|
23
|
+
mcp.add_tool(self._tool(spec))
|
|
24
|
+
super().__init__(mcp)
|
|
25
|
+
self._state = State(episode_id=None, step_count=0)
|
|
26
|
+
self._reward: Optional[float] = None
|
|
27
|
+
|
|
28
|
+
def _tool(self, spec: dict) -> Tool:
|
|
29
|
+
env = self
|
|
30
|
+
|
|
31
|
+
class SeatTool(Tool):
|
|
32
|
+
async def run(self, arguments: dict[str, Any]) -> ToolResult:
|
|
33
|
+
text, reward, done = env.seat.call(spec["name"], arguments)
|
|
34
|
+
if done:
|
|
35
|
+
env._reward = reward
|
|
36
|
+
return ToolResult(content=text)
|
|
37
|
+
|
|
38
|
+
return SeatTool(name=spec["name"], description=spec.get("description", ""), parameters=spec["input_schema"])
|
|
39
|
+
|
|
40
|
+
def reset(self, seed: Optional[int] = None, episode_id: Optional[str] = None, **kwargs: Any) -> Observation:
|
|
41
|
+
text, _ = self.seat.reset(seed if seed is not None else "openenv")
|
|
42
|
+
self._state, self._reward = State(episode_id=episode_id, step_count=0), None
|
|
43
|
+
return Observation(done=False, reward=None, metadata={"chart": text})
|
|
44
|
+
|
|
45
|
+
def step(self, action, timeout_s: Optional[float] = None, **kwargs: Any) -> Observation:
|
|
46
|
+
obs = super().step(action, timeout_s=timeout_s, **kwargs)
|
|
47
|
+
self._state.step_count += 1
|
|
48
|
+
if isinstance(action, CallToolAction) and self._reward is not None:
|
|
49
|
+
obs.done, obs.reward = True, self._reward
|
|
50
|
+
obs.metadata = {**(obs.metadata or {}), "cost": self.seat.cost, "credit": self.seat.credit, "budget": self.seat.budget}
|
|
51
|
+
return obs
|
|
52
|
+
|
|
53
|
+
def _step_impl(self, action, timeout_s: Optional[float] = None, **kwargs: Any) -> Observation:
|
|
54
|
+
return Observation(done=False, reward=None, metadata={"error": f"Unknown action {type(action).__name__}; use ListToolsAction or CallToolAction."})
|
|
55
|
+
|
|
56
|
+
@property
|
|
57
|
+
def state(self) -> State:
|
|
58
|
+
return self._state
|
|
59
|
+
|
|
60
|
+
def close(self) -> None:
|
|
61
|
+
self.seat.close()
|
|
62
|
+
super().close()
|