File size: 5,536 Bytes
9b1756a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
157
158
159
160
161
162
"""Action, observation, and tool-result models for the Pulse environment."""

from __future__ import annotations

from typing import Any

from pydantic import BaseModel, ConfigDict, Field

try:
    from openenv.core.env_server.types import Action, Observation
except ImportError:  # pragma: no cover - allows local mock development without openenv
    class Action(BaseModel):
        """Fallback action base model for local development."""

    class Observation(BaseModel):
        """Fallback observation base model for local development."""


from .patient_state import (
    ArterialBloodGasResult,
    BasicMetabolicPanelResult,
    CompleteBloodCountResult,
    LactateTrend,
    MentalStatus,
    PatientState,
    ScenarioDifficulty,
)
from .tool_catalog import EXTENDED_TOOL_NAMES, INITIAL_TOOL_NAMES, KNOWN_TOOL_NAMES


class PulsePhysiologyAction(Action):
    """Tool invocation sent from the agent or client."""

    tool_name: str = Field(..., description="Name of the tool to execute.")
    arguments: dict[str, Any] = Field(
        default_factory=dict,
        description="Structured arguments for the named tool.",
    )
    reasoning: str | None = Field(
        default=None,
        description="Optional human-readable rationale for the action.",
    )


ToolAction = PulsePhysiologyAction


class ToolResult(BaseModel):
    """Structured result returned for every handled tool call."""

    model_config = ConfigDict(extra="forbid")

    tool_name: str
    success: bool
    message: str
    state_changed: bool
    changed_fields: list[str] = Field(default_factory=list)


class ToolError(BaseModel):
    """Structured error attached to failed tool calls."""

    model_config = ConfigDict(extra="forbid")

    code: str
    message: str
    retryable: bool


class ObservationMetadata(BaseModel):
    """Non-state metadata returned alongside each step response."""

    step_count: int = Field(default=0, ge=0, description="Episode step count")
    available_tools: list[str] = Field(
        default_factory=lambda: list(KNOWN_TOOL_NAMES),
        description="Tools the environment currently exposes",
    )


class PulsePhysiologyObservation(Observation):
    """Observation returned by the Pulse physiology environment."""

    scenario_id: str = Field(default="baseline")
    scenario_difficulty: ScenarioDifficulty = Field(default="medium")
    patient_id: str = Field(default="standard_male")
    sim_time_s: float = Field(default=0.0)

    heart_rate_bpm: float | None = Field(default=None)
    systolic_bp_mmhg: float | None = Field(default=None)
    diastolic_bp_mmhg: float | None = Field(default=None)
    mean_arterial_pressure_mmhg: float | None = Field(default=None)
    cardiac_output_l_per_min: float | None = Field(default=None)

    spo2: float | None = Field(default=None)
    respiration_rate_bpm: float | None = Field(default=None)
    blood_volume_ml: float | None = Field(default=None)
    mental_status: MentalStatus = Field(default="alert")
    active_alerts: list[str] = Field(default_factory=list)

    etco2_mmhg: float | None = Field(default=None)
    tidal_volume_ml: float | None = Field(default=None)
    breath_sounds: str = Field(default="present bilateral")
    core_temperature_c: float | None = Field(default=None)

    shock_index: float | None = Field(default=None)
    lactate_trend: LactateTrend = Field(default="stable")

    position: str = Field(default="supine")
    oxygen_device: str | None = Field(default=None)
    oxygen_flow_lpm: float | None = Field(default=None)
    airway_support: str | None = Field(default=None)
    intubated: bool = Field(default=False)

    abg_result: ArterialBloodGasResult = Field(default_factory=ArterialBloodGasResult)
    cbc_result: CompleteBloodCountResult = Field(default_factory=CompleteBloodCountResult)
    bmp_result: BasicMetabolicPanelResult = Field(default_factory=BasicMetabolicPanelResult)

    pending_diagnostics: dict[str, int] = Field(default_factory=dict)
    ready_diagnostics: list[str] = Field(default_factory=list)
    active_infusions: dict[str, float] = Field(default_factory=dict)
    active_hemorrhages: dict[str, float] = Field(default_factory=dict)

    available_tools: list[str] = Field(default_factory=list)
    tool_result: ToolResult | None = Field(default=None)
    error: ToolError | None = Field(default=None)
    reward: float | None = Field(default=None)
    done: bool = Field(default=False)
    metadata: dict[str, Any] = Field(default_factory=dict)

    @classmethod
    def from_patient_state(
        cls,
        state: PatientState,
        *,
        reward: float | None = None,
        available_tools: list[str] | None = None,
        tool_result: ToolResult | None = None,
        error: ToolError | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> "PulsePhysiologyObservation":
        payload = state.model_dump()
        payload.update(
            reward=reward,
            done=state.done,
            available_tools=available_tools or [],
            tool_result=tool_result,
            error=error,
            metadata=metadata or {},
        )
        return cls(**payload)


class EnvironmentResponse(BaseModel):
    """Canonical step envelope for the mock and real adapters."""

    observation: PulsePhysiologyObservation
    reward: float
    done: bool
    metadata: ObservationMetadata = Field(default_factory=ObservationMetadata)
    tool_result: ToolResult | None = Field(default=None)
    error: ToolError | None = Field(default=None)