| from pydantic import BaseModel, Field |
| from typing import Optional, Dict, Literal, List, Any |
|
|
|
|
| class Action(BaseModel): |
| optimized_code: str |
| strategy: Optional[str] = None |
| expected_speedup: Optional[float] = None |
|
|
| class Reward(BaseModel): |
| value: float = Field(ge=0.0, le=1.0) |
| components: Dict[str, float] |
|
|
| class Observation(BaseModel): |
| task_id: str |
| task_name: str |
| difficulty: Literal["easy", "medium", "hard"] |
| baseline_code: str |
| current_best_code: str |
| current_best_speedup: float |
| step_count: int |
| max_steps: int |
| pending_checks: List[str] |
| completed_checks: List[str] |
| done: bool |
|
|
| class EnvState(BaseModel): |
| initialized: bool |
| task_id: Optional[str] =None |
| step_count: int = 0 |
| max_steps: int = 0 |
| total_reward: float = 0.0 |
| best_code: str = "" |
| best_speedup: float = 1.0 |
| completed_checks: List[str] = Field(default_factory=list) |
| action_history: List[Dict[str, Any]] = Field(default_factory=list) |
|
|
| class ResetRequest(BaseModel): |
| task_id: Optional[str] = None |
|
|
| class StepRequest(BaseModel): |
| action: Action |
|
|
| class StepResult(BaseModel): |
| observation:Observation |
| reward: Reward |
| done: bool |
| info: Dict[str, Any] |
|
|
|
|
