File size: 2,610 Bytes
37a52ed 4888b97 37a52ed 4888b97 37a52ed 0f6657d 4888b97 37a52ed 0f6657d 37a52ed 4888b97 37a52ed 4888b97 0f6657d 37a52ed 4888b97 37a52ed 0f6657d 4888b97 37a52ed 0f6657d 37a52ed 4888b97 42d0b80 4888b97 0f6657d 37a52ed 4888b97 37a52ed 4888b97 37a52ed 0f6657d 4888b97 37a52ed 4888b97 fe5a027 4888b97 | 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 | from __future__ import annotations
from typing import Any, Dict, List, Literal, Optional
from pydantic import BaseModel, Field, validator
Difficulty = Literal["easy", "medium", "hard"]
DoneReason = Optional[Literal["exfil_detected", "injection_flagged", "max_steps"]]
class Action(BaseModel):
type: Literal["ReadDoc", "WriteDB", "SendEmail", "FlagInjection", "NoOp"]
to: Optional[str] = None
body: Optional[str] = None
content: Optional[str] = None
reason: Optional[str] = None
@validator("to", always=True)
def validate_send_email(cls, value, values):
if values.get("type") == "SendEmail" and not value:
raise ValueError("SendEmail requires a 'to' address")
return value
class Observation(BaseModel):
subject: str
body: str
last_action: Optional[str] = None
last_action_result: str
step_count: int
task_difficulty: Difficulty
summary_required: bool = False
current_summary: Optional[str] = None
class Reward(BaseModel):
value: float
reason: str
breakdown: Dict[str, float]
@validator("value")
def clamp_reward(cls, value):
return max(0.0, min(1.0, value))
class State(BaseModel):
current_email: Dict[str, Any]
actions_taken: List[Dict[str, Any]]
cumulative_reward: float
step_count: int
done: bool
done_reason: DoneReason = None
task_difficulty: Difficulty
current_summary: Optional[str] = None
class EpisodeResult(BaseModel):
final_score: float
breakdown: Dict[str, float]
exfil_occurred: bool
injection_flagged: bool
task_completed: bool
summary_score: Optional[float] = None
total_steps: int
trajectory: List[Dict[str, Any]]
@validator("final_score")
def clamp_final_score_to_open_interval(cls, value):
epsilon = 0.0001
return min(1.0 - epsilon, max(epsilon, value))
class ResetRequest(BaseModel):
difficulty: Difficulty = "easy"
seed: Optional[int] = Field(default=None, description="Optional deterministic seed for scenario selection")
class StepRequest(BaseModel):
action: Action
class StepResponse(BaseModel):
observation: Observation
reward: Reward
done: bool
info: Dict[str, Any]
class HealthResponse(BaseModel):
status: str
env: str
version: str
class MetadataResponse(BaseModel):
name: str
description: str
version: str
tasks: List[str]
class SchemaResponse(BaseModel):
action: Dict[str, Any]
observation: Dict[str, Any]
state: Dict[str, Any]
reward: Dict[str, Any]
result: Dict[str, Any]
|