File size: 4,272 Bytes
f44f429 | 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 | """Pydantic v2 models representing actions, observations, and state payloads."""
from typing import Optional, Any, Dict
from pydantic import BaseModel, Field
# ── Agent Action ──────────────────────────────────────────────────────────────
class CodeReviewAction(BaseModel):
"""Action taken by the agent: a structured code review or a file request."""
request_file: Optional[bool] = Field(None, description="Request the file contents")
bug_identified: Optional[bool] = Field(None, description="Whether a bug was found")
bug_location: Optional[str] = Field(None, description="Location of the bug (function, line, variable)")
bug_type: Optional[str] = Field(None, description="Type: off-by-one | logic-error | security-vulnerability | none")
bug_description: Optional[str] = Field(None, description="Detailed explanation of why this is a bug")
severity: Optional[str] = Field(None, description="Severity: none | low | medium | high | critical")
suggested_fix: Optional[str] = Field(None, description="The corrected code or a description of how to fix it")
# ── Observation ───────────────────────────────────────────────────────────────
class CodeObservation(BaseModel):
"""What the agent sees at each step."""
task_id: str = Field(..., description="Unique task identifier")
language: str = Field(..., description="Programming language")
difficulty: str = Field(..., description="Level: easy | medium | hard")
code_snippet: str = Field(..., description="The code to review")
context: str = Field(..., description="Production context describing what the code does")
pr_title: str = Field(..., description="Pull request title submitted by developer")
file_path: str = Field(..., description="File path of the code in the repository")
# ── Step Result ───────────────────────────────────────────────────────────────
class StepResult(BaseModel):
"""Result returned from env.step()."""
observation: Optional[CodeObservation] = Field(None, description="Observation if not terminal")
reward: float = Field(..., description="Reward generated for the preceding action")
done: bool = Field(..., description="Terminal state flag")
info: Dict[str, Any] = Field(default_factory=dict, description="Metadata dictionary")
# ── State ─────────────────────────────────────────────────────────────────────
class StateResponse(BaseModel):
"""Internal environment state exposed via /state."""
task_id: str = Field(..., description="Current running task")
step: int = Field(..., description="Current evaluation step")
done: bool = Field(..., description="Whether the episode resides in a terminal state")
total_reward: float = Field(..., description="Sum of step rewards over the episode")
# ── API Helpers ───────────────────────────────────────────────────────────────
class ResetResponse(BaseModel):
"""Response wrapper returned strictly on environment resets."""
observation: CodeObservation = Field(..., description="Initial environment observation upon reset")
class TaskInfo(BaseModel):
"""Metadata regarding an available task scenario."""
id: str = Field(..., description="Task UUID or unique string identifier")
language: str = Field(..., description="Source code language for the flaw context")
bug_class: str = Field(..., description="The classification parameter of the embedded bug")
difficulty: str = Field(..., description="The difficulty tier indicator (e.g. easy, medium)")
Action = CodeReviewAction
Observation = CodeObservation
Reward = float
|