Migrate dataclass to pydantic BaseModel

Author: john-b-yang

codeclash/agents/utils.py

@@ -1,9 +1,9 @@
 import os
-from dataclasses import asdict, dataclass
 from pathlib import Path
 
 from dotenv import load_dotenv
 from jinja2 import Template
+from pydantic import BaseModel
 
 load_dotenv()
 
@@ -16,8 +16,7 @@ def resolve_api_key(model: str) -> str:
     return ""
 
 
-@dataclass
-class GameContext:
+class GameContext(BaseModel):
     """
     A class that gives agent access to a partial view of the game state.
 
@@ -38,11 +37,11 @@ class GameContext:
     working_dir: str
 
     def _render_prompt_templates(self) -> dict:
-        context = asdict(self)
+        context = self.dict()
         return {key: Template(template_str).render(**context) for key, template_str in self.prompts.items()}
 
     def to_template_vars(self) -> dict[str, str]:
         """Convert the GameContext to a dictionary for rendering prompts in the agent"""
-        out = asdict(self) | self._render_prompt_templates()
+        out = self.dict() | self._render_prompt_templates()
         out.pop("prompts")
         return out

codeclash/games/abstract.py

@@ -2,36 +2,33 @@
 import subprocess
 import time
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
 from minisweagent.environments.docker import DockerEnvironment
+from pydantic import BaseModel
 
 from codeclash.agents.abstract import Player
 from codeclash.constants import DIR_LOGS, DIR_WORK, GH_ORG
 from codeclash.utils.environment import assert_zero_exit_code, copy_between_containers
 from codeclash.utils.log import get_logger
 
 
-@dataclass
-class RoundStats:
+class RoundStats(BaseModel):
     winner: str
     scores: dict[str, float]  # Map of player to game metric (e.g. # of wins, assets accumulated)
-    details: dict[str, Any] = None  # Optional, for game-specific info
+    details: dict[str, Any] | None = None  # Optional, for game-specific info
 
     def __str__(self) -> str:
         return "\n".join([f"- Winner: {self.winner}", f"- Scores: {self.scores}"])
 
 
-@dataclass
-class RoundData:
+class RoundData(BaseModel):
     logs: list[str]
     results: list[str]
 
 
-@dataclass
-class RoundRecord:
+class RoundRecord(BaseModel):
     data: RoundData
     stats: RoundStats
 

codeclash/games/battlesnake/main.py

@@ -75,7 +75,7 @@ def execute_round(self, agents: list[Player]) -> RoundData:
                 self.environment.execute(f"rm -f game/{output_file}")
                 time.sleep(0.05)
 
-            return RoundData(log_outputs, result_outputs)
+            return RoundData(logs=log_outputs, results=result_outputs)
         finally:
             # Kill all python servers when done
             self.environment.execute("pkill -f 'python main.py' || true")

codeclash/games/robocode/main.py

@@ -108,4 +108,4 @@ def execute_round(self, agents: list[Player]) -> RoundData:
         # Clean up the results file
         self.environment.execute(f"rm -f {results_file}")
 
-        return RoundData([response["output"]], [result_output])
+        return RoundData(logs=[response["output"]], results=[result_output])