Ramkan7/Patch_Hawk

1"""2PatchHawk typed models — OpenEnv-compliant Pydantic models.3 4Extends openenv.core base types (Action, Observation, State) so the5environment is fully compatible with the OpenEnv framework.6"""7 8from typing import Optional, List9 10from pydantic import BaseModel, ConfigDict, Field11 12from openenv.core import Action, Observation, State13 14 15# ── Observation ──────────────────────────────────────────────────────16 17 18class PatchHawkObservation(Observation):19    """Observation returned by PatchHawkEnv after reset() and step()."""20 21    code_snippet: str = Field(default="", description="Python source code to analyse")22    static_flags: List[int] = Field(23        default_factory=list,24        description="Binary flags indicating static risk patterns",25    )26    risk_score: float = Field(27        default=0.0, description="Precomputed heuristic risk score 0-1"28    )29    sandbox_telemetry: Optional[str] = Field(30        None, description="Output from previous sandbox execution"31    )32 33 34# ── Action ───────────────────────────────────────────────────────────35 36 37class PatchHawkAction(Action):38    """Action submitted to PatchHawkEnv.step().39 40    action_type values:41        0 = ANALYZE42        1 = EXECUTE_SANDBOX43        2 = BLOCK_PR44        3 = SUBMIT_PATCH45        4 = REQUEST_REVIEW46    """47 48    action_type: int = Field(49        ...,50        description="0: ANALYZE, 1: EXECUTE_SANDBOX, 2: BLOCK_PR, "51        "3: SUBMIT_PATCH, 4: REQUEST_REVIEW",52    )53    patch_content: Optional[str] = Field(54        None, description="The unified context patch if action is SUBMIT_PATCH"55    )56    reasoning: Optional[str] = Field(57        None, description="Explanation of the vulnerability and chosen action"58    )59    predicted_risk: Optional[float] = Field(60        None, description="LLM predicted risk score (0.0 to 1.0)"61    )62 63 64# ── State ────────────────────────────────────────────────────────────65 66 67class PatchHawkState(State):68    """Internal state of a PatchHawkEnv episode."""69 70    scenario_id: str = Field(default="", description="Current scenario ID")71    current_task: Optional[str] = Field(72        None,73        description="Active task_id (easy_typosquat, medium_obfuscated, hard_patch)",74    )75    last_action_type: Optional[int] = Field(None, description="Last action type taken")76    patch_validated: bool = Field(77        default=False, description="Whether the last patch was validated"78    )79    sandbox_log: Optional[str] = Field(80        None, description="Most recent sandbox execution log"81    )82 83 84# ── Reward (standalone — no OpenEnv base type for rewards) ───────────85 86 87class PatchHawkReward(BaseModel):88    """Pydantic reward model used by graders and inference logging."""89 90    model_config = ConfigDict(91        extra="forbid",92        validate_assignment=True,93    )94 95    value: float = Field(default=0.0, description="Numeric reward signal")96    reason: str = Field(default="", description="Human-readable reward reason")97