environment/models.py

from typing import Optional, List, Dict, Any, Literal
from pydantic import BaseModel, Field
from enum import Enum


class ReviewActionType(str, Enum):
    ADD_COMMENT = "add_comment"
    SUGGEST_FIX = "suggest_fix"
    APPROVE = "approve"
    REQUEST_CHANGES = "request_changes"
    MARK_AS_RESOLVED = "mark_as_resolved"


class Comment(BaseModel):
    line_number: int = Field(..., ge=1, description="Line number in the diff")
    content: str = Field(..., min_length=1, description="Comment content")
    is_issue: bool = Field(True, description="Whether this is pointing out an issue")
    severity: Optional[Literal["low", "medium", "high", "critical"]] = Field(
        None, description="Issue severity"
    )
    resolved: bool = Field(False, description="Whether the comment has been addressed")


class Suggestion(BaseModel):
    original_line: int = Field(..., ge=1, description="Line number with issue")
    suggested_code: str = Field(..., description="Suggested replacement code")
    explanation: str = Field(..., description="Explanation of the suggestion")
    accepted: bool = Field(False, description="Whether the suggestion was accepted")


class ReviewAction(BaseModel):
    action_type: ReviewActionType = Field(..., description="Type of action to take")
    comments: List[Comment] = Field(default_factory=list, description="Comments to add")
    suggestions: List[Suggestion] = Field(default_factory=list, description="Suggested fixes")
    final_decision: Optional[Literal["approved", "changes_requested"]] = Field(
        None, description="Final review decision"
    )


class CodeContext(BaseModel):
    file_path: str = Field(..., description="Path to the file being reviewed")
    file_extension: str = Field(..., description="File extension")
    code_diff: str = Field(..., description="Git diff of changes")
    surrounding_code: str = Field(..., description="Surrounding code context")
    language: str = Field(..., description="Programming language")
    line_count: int = Field(..., ge=1, description="Number of lines changed")


class TaskMetadata(BaseModel):
    task_id: str = Field(..., description="Unique task identifier")
    task_name: str = Field(..., description="Human-readable task name")
    difficulty: Literal["easy", "medium", "hard"] = Field(..., description="Task difficulty")
    description: str = Field(..., description="Task description")
    expected_issues: List[Dict[str, Any]] = Field(default_factory=list, description="Expected issues to find")


class ReviewState(BaseModel):
    code_context: CodeContext = Field(..., description="Code being reviewed")
    task_metadata: TaskMetadata = Field(..., description="Current task information")
    comments_made: List[Comment] = Field(default_factory=list, description="Comments already made")
    suggestions_made: List[Suggestion] = Field(default_factory=list, description="Suggestions already made")
    current_step: int = Field(0, description="Current step number")
    is_complete: bool = Field(False, description="Whether review is complete")
    final_decision: Optional[str] = Field(None, description="Final decision if made")
    last_action_valid: bool = Field(True, description="Whether last action was valid")
    last_error: Optional[str] = Field(None, description="Last error message if any")


class Observation(BaseModel):
    code_diff: str = Field(..., description="Current code diff")
    file_context: str = Field(..., description="Surrounding file context")
    file_path: str = Field(..., description="File being reviewed")
    language: str = Field(..., description="Programming language")
    task_description: str = Field(..., description="Description of the review task")
    task_difficulty: str = Field(..., description="Task difficulty level")
    current_step: int = Field(0, description="Current step number")
    max_steps: int = Field(50, description="Maximum steps allowed")
    previous_comments: List[Comment] = Field(default_factory=list, description="Comments made so far")
    previous_suggestions: List[Suggestion] = Field(default_factory=list, description="Suggestions made so far")
    review_complete: bool = Field(False, description="Whether review is complete")
    final_decision_made: Optional[str] = Field(None, description="Final decision if made")