models.py

"""Typed models for the MolForge OpenEnv environment."""

from __future__ import annotations

from typing import Any, Dict, List, Literal, Optional

from openenv.core.env_server.types import Action, Observation, State
from pydantic import BaseModel, Field


EDIT_TYPES = Literal["add_fragment", "substitute", "remove", "undo_last_edit"]
ACTION_TYPES = Literal["edit", "run_assay", "submit", "restart", "defer"]
TOOL_TYPES = Literal[
    "evaluate_properties",
    "dock_target",
    "assay_toxicity",
    "estimate_synthesizability",
    "evaluate_novelty",
    "search_literature",
    "run_md_simulation",
]
SLOT_TYPES = Literal["warhead", "hinge", "solvent_tail", "back_pocket"]
ROLE_TYPES = Literal[
    "lead_chemist",
    "toxicologist",
    "assay_planner",
    "process_chemist",
    "team",
]
MESSAGE_TYPES = Literal[
    "proposal",
    "objection",
    "risk_flag",
    "assay_request",
    "approval",
    "rejection",
    "revision_request",
    "submission_recommendation",
]
SEVERITY_TYPES = Literal["low", "medium", "high", "critical"]
EFFECT_TYPES = Literal["up", "down", "neutral", "unknown", "not_applicable"]
COORDINATION_MODES = Literal["single_agent", "multi_agent"]
GOVERNANCE_STATES = Literal["ready", "executed", "needs_revision", "policy_veto"]


class MoleculeSlot(BaseModel):
    """Visible fragment assignment for a molecule slot."""

    slot: SLOT_TYPES
    fragment: str = Field(..., description="Selected fragment for the slot")
    editable: bool = Field(default=True, description="Whether the slot is editable")


class AssayReading(BaseModel):
    """Structured oracle result surfaced to the agent."""

    tool_name: str
    property_name: str
    estimate: float = Field(..., ge=0.0, le=1.0)
    confidence_low: float = Field(..., ge=0.0, le=1.0)
    confidence_high: float = Field(..., ge=0.0, le=1.0)
    runs: int = Field(default=1, ge=1)
    molecule_signature: str
    summary: str = ""


class RewardComponent(BaseModel):
    """Named reward component used in report cards and debugging."""

    name: str
    value: float
    explanation: str


class ConstraintCheck(BaseModel):
    """Constraint status based only on currently visible evidence."""

    name: str
    target: str
    satisfied: Optional[bool] = None
    actual: Optional[float] = None
    evidence_status: Literal["known", "unknown"] = "unknown"


class AgentMessage(BaseModel):
    """Structured inter-agent communication message."""

    message_id: str = ""
    sender: ROLE_TYPES
    receiver: str = "team"
    message_type: MESSAGE_TYPES
    severity: SEVERITY_TYPES = "low"
    reference_action_type: Optional[ACTION_TYPES] = None
    summary: str = Field(default="", max_length=240)
    payload: Dict[str, Any] = Field(default_factory=dict)


class RoleObservation(BaseModel):
    """Role-specific structured observation slice."""

    role: ROLE_TYPES
    local_objective: str
    permissions: List[str] = Field(default_factory=list)
    observation: Dict[str, Any] = Field(default_factory=dict)


class GovernanceStatus(BaseModel):
    """Outcome of the multi-agent review process for the last turn."""

    status: GOVERNANCE_STATES = "ready"
    explanation: str = ""
    required_roles: List[str] = Field(default_factory=list)
    approvals: List[str] = Field(default_factory=list)
    objections: List[str] = Field(default_factory=list)
    vetoes: List[str] = Field(default_factory=list)
    executable: bool = True


class MolForgeAction(Action):
    """Single team turn action spanning edits, assays, messages, and submission."""

    action_type: ACTION_TYPES = Field(
        ..., description="High-level action type to execute this turn"
    )
    acting_role: ROLE_TYPES = Field(
        default="lead_chemist",
        description="Role claiming ownership of the executable team decision",
    )
    edit_type: Optional[EDIT_TYPES] = Field(
        default=None, description="Edit subtype when action_type is edit"
    )
    slot: Optional[SLOT_TYPES] = Field(
        default=None, description="Editable molecular slot when performing edits"
    )
    fragment: Optional[str] = Field(
        default=None, description="Fragment identifier for edit actions"
    )
    tool_name: Optional[TOOL_TYPES] = Field(
        default=None, description="Oracle or tool name for run_assay actions"
    )
    messages: List[AgentMessage] = Field(
        default_factory=list,
        description="Structured multi-agent communication bundle for this decision turn",
    )
    rationale: str = Field(
        default="",
        description="Short explanation of why the final decision should help",
        max_length=400,
    )
    evidence: List[str] = Field(
        default_factory=list,
        description="Visible observation facts supporting the action; do not include hidden state.",
        max_length=5,
    )
    expected_effects: Dict[str, EFFECT_TYPES] = Field(
        default_factory=dict,
        description="Directional public prediction for potency, toxicity, synth, novelty, or budget.",
    )


class MolForgeObservation(Observation):
    """Observation emitted after reset and each step."""

    scenario_id: str
    difficulty: str
    state_label: str = "[start]"
    state_path: List[str] = Field(default_factory=list)
    coordination_mode: COORDINATION_MODES = "multi_agent"
    enabled_roles: List[str] = Field(default_factory=list)
    task_brief: str
    target_name: str
    current_molecule: str
    molecule_slots: List[MoleculeSlot] = Field(default_factory=list)
    editable_slots: List[str] = Field(default_factory=list)
    step_index: int = Field(default=0, ge=0)
    max_steps: int = Field(default=0, ge=1)
    remaining_budget: int = Field(default=0, ge=0)
    budget_used: int = Field(default=0, ge=0)
    max_budget: int = Field(default=0, ge=1)
    known_assays: List[AssayReading] = Field(default_factory=list)
    role_observations: List[RoleObservation] = Field(default_factory=list)
    message_log: List[AgentMessage] = Field(default_factory=list)
    governance: GovernanceStatus = Field(default_factory=GovernanceStatus)
    last_transition_summary: str = ""
    visible_metrics: Dict[str, float] = Field(default_factory=dict)
    constraint_status: List[ConstraintCheck] = Field(default_factory=list)
    reward_breakdown: List[RewardComponent] = Field(default_factory=list)
    allowed_actions: List[str] = Field(default_factory=list)
    report_card: str = ""


class MolForgeState(State):
    """Internal environment state surfaced through the state() API."""

    scenario_id: str = ""
    difficulty: str = ""
    state_label: str = "[start]"
    state_path: List[str] = Field(default_factory=list)
    coordination_mode: COORDINATION_MODES = "multi_agent"
    enabled_roles: List[str] = Field(default_factory=list)
    target_name: str = ""
    current_molecule: str = ""
    remaining_budget: int = 0
    budget_used: int = 0
    max_budget: int = 0
    visited_states: int = 0
    known_assay_count: int = 0
    invalid_action_count: int = 0
    objection_count: int = 0
    oracle_call_count: int = 0
    message_count: int = 0
    decision_count: int = 0
    submitted: bool = False
    last_error_code: str = ""
    reward_total: float = 0.0
    metadata: Dict[str, Any] = Field(default_factory=dict)