Update models.py

models.py

@@ -2,21 +2,25 @@ from pydantic import BaseModel
 from typing import List, Optional
 
 class Observation(BaseModel):
-    pr_code: str
-    comments: List[str]
+    pr_title: str                      # title of the pull request
+    pr_description: str                # description of the change
+    code_snippet: str                  # the code under review
+    comments: List[str]                # existing comments
     agent_comment: Optional[str] = None
     step: int = 0
     done: bool = False
 
 class Action(BaseModel):
-    action_type: str  # "write_comment", "skip", "done"
+    action_type: str   # "write_comment", "skip", "done"
     comment_text: Optional[str] = None
 
 class Reward(BaseModel):
     value: float
 
 class State(BaseModel):
-    pr_code: str
+    pr_title: str
+    pr_description: str
+    code_snippet: str
     comments: List[str]
     agent_comment: Optional[str]
     step: int