fix:improvements

env/models.py

@@ -1,20 +1,28 @@
-from pydantic import BaseModel
-from typing import List, Optional
+from pydantic import BaseModel, validator
+from typing import List, Optional, Dict, Literal
 
 # What agent sees
 class Observation(BaseModel):
     code: str
     intent: str
     constraints: List[str]
+    task_id: str
+    difficulty: Literal["easy", "medium", "hard"]
+    step_index: int
 
 # What agent sends
 class Action(BaseModel):
-    action_type: str  # "detect", "risk", "fix"
+    action_type: Literal["detect", "risk", "fix"]
     detected_issues: Optional[List[str]] = None
-    risk_level: Optional[str] = None  # "low", "medium", "high"
+    risk_level: Optional[Literal["low", "medium", "high"]] = None
     suggested_fix: Optional[str] = None
 
 # What env returns
 class Reward(BaseModel):
-    score: float
-    feedback: str
+    value: float
+    feedback: str
+    partial_scores: Dict[str, float]
+    
+    @validator('value')
+    def clamp_value(cls, v):
+        return max(0.0, min(1.0, v))