| from __future__ import annotations |
| from enum import Enum |
| from typing import Any, Dict, List, Optional |
| from pydantic import BaseModel, Field |
|
|
|
|
| class Severity(str, Enum): |
| low = "low" |
| medium = "medium" |
| high = "high" |
| critical = "critical" |
|
|
|
|
| class ActionType(str, Enum): |
| add_comment = "add_comment" |
| flag_bug = "flag_bug" |
| flag_vulnerability = "flag_vulnerability" |
| suggest_fix = "suggest_fix" |
| request_changes = "request_changes" |
| approve = "approve" |
| submit_report = "submit_report" |
|
|
|
|
| class TaskId(str, Enum): |
| bug_detection = "bug_detection" |
| security_audit = "security_audit" |
| architecture_review = "architecture_review" |
|
|
|
|
| class CodeFile(BaseModel): |
| filename: str = Field(..., description="Relative path of the file") |
| language: str = Field(..., description="Programming language") |
| content: str = Field(..., description="Full source code content") |
| line_count: int = Field(..., description="Number of lines") |
|
|
|
|
| class ReviewComment(BaseModel): |
| step: int |
| action: ActionType |
| line: Optional[int] = None |
| message: str |
| severity: Optional[Severity] = None |
|
|
|
|
| class Observation(BaseModel): |
| task_id: TaskId |
| task_description: str |
| files: List[CodeFile] |
| history: List[ReviewComment] = Field(default_factory=list) |
| step: int = Field(0) |
| max_steps: int = Field(20) |
| done: bool = Field(False) |
| hint: Optional[str] = Field(None) |
|
|
|
|
| class Action(BaseModel): |
| action_type: ActionType |
| line: Optional[int] = None |
| filename: Optional[str] = None |
| message: Optional[str] = None |
| severity: Optional[Severity] = None |
| cve_category: Optional[str] = None |
| original_code: Optional[str] = None |
| suggested_code: Optional[str] = None |
| report: Optional[Dict[str, Any]] = None |
| reason: Optional[str] = None |
|
|
|
|
| class RewardBreakdown(BaseModel): |
| coverage: float = 0.0 |
| precision: float = 0.0 |
| severity_match: float = 0.0 |
| false_positive_penalty: float = 0.0 |
| efficiency: float = 0.0 |
| report_quality: float = 0.0 |
|
|
|
|
| class Reward(BaseModel): |
| value: float = Field(..., ge=0.0, le=1.0) |
| breakdown: RewardBreakdown |
| delta: float = 0.0 |
| done: bool = False |
| info: Dict[str, Any] = Field(default_factory=dict) |
|
|
|
|
| class EnvironmentState(BaseModel): |
| task_id: TaskId |
| step: int |
| max_steps: int |
| done: bool |
| ground_truth_count: int |
| found_count: int |
| false_positive_count: int |
| cumulative_reward: float |
| current_score: float |
|
|
|
|
| class StepResponse(BaseModel): |
| observation: Observation |
| reward: Reward |
| done: bool |
| info: Dict[str, Any] = Field(default_factory=dict) |
|
|
|
|
| class ResetResponse(BaseModel): |
| observation: Observation |
| info: Dict[str, Any] = Field(default_factory=dict) |
|
|
|
|
| class ValidateResponse(BaseModel): |
| valid: bool |
| version: str |
| tasks: List[str] |
| message: str |
|
|
