Ref: Move most of main() logic to Game class

Author: klieret

codeclash/agents/abstract.py

@@ -2,7 +2,7 @@
 from abc import ABC, abstractmethod
 
 from dotenv import load_dotenv
-from minisweagent import Environment
+from minisweagent.environments.docker import DockerEnvironment
 
 from codeclash.agents.utils import GameContext
 from codeclash.constants import GH_ORG
@@ -16,7 +16,7 @@ class Player(ABC):
     def __init__(
         self,
         config: dict,
-        environment: Environment,
+        environment: DockerEnvironment,
         game_context: GameContext,
     ):
         self.config = config

codeclash/agents/minisweagent.py

@@ -8,8 +8,9 @@
 
 import yaml
 from jinja2 import Template
-from minisweagent import Environment, Model
+from minisweagent import Model
 from minisweagent.agents.default import AgentConfig, DefaultAgent
+from minisweagent.environments.docker import DockerEnvironment
 from minisweagent.models.litellm_model import LitellmModel
 from minisweagent.run.utils.save import save_traj
 from rich.console import Console
@@ -28,7 +29,7 @@ class ClashAgent(DefaultAgent):
     def __init__(
         self,
         model: Model,
-        env: Environment,
+        env: DockerEnvironment,
         name: str,
         game_context: GameContext,
         *,
@@ -69,7 +70,7 @@ class MiniSWEAgent(Player):
     """Player with agentic code editing capabilities"""
 
     def __init__(
-        self, config: dict, environment: Environment, game_context: GameContext
+        self, config: dict, environment: DockerEnvironment, game_context: GameContext
     ):
         super().__init__(config, environment=environment, game_context=game_context)
 

codeclash/games/__init__.py

@@ -7,16 +7,17 @@
 
 
 # might consider postponing imports to avoid loading things we don't need
-def get_game(config: dict) -> CodeGame:
+def get_game(config: dict, **kwargs) -> CodeGame:
     game = {
-        x.name: x for x in [
+        x.name: x
+        for x in [
             BattleCodeGame,
             BattleSnakeGame,
             CoreWarGame,
             RoboCodeGame,
-            RobotRumbleGame
+            RobotRumbleGame,
         ]
     }.get(config["game"]["name"])
     if game is None:
         raise ValueError(f"Unknown game: {config['game']['name']}")
-    return game(config)
+    return game(config, **kwargs)

codeclash/games/abstract.py

@@ -24,14 +24,15 @@
 class CodeGame(ABC):
     name: str
 
-    def __init__(self, config: dict):
+    def __init__(self, config: dict, *, push_agent: bool = False):
         self.url_gh: str = f"git@github.com:{GH_ORG}/{self.name}.git"
         self.artifacts: list[Path] = []
         """Artifact objects that we might want to clean up after the game."""
         self.scoreboard: list[tuple[int, str]] = []
         """List of (round number, winner (player id))"""
         self.game_config: dict = config["game"]
         self.config: dict = config
+        self.push_agent: bool = push_agent
         self.rounds: int = self.game_config.get("rounds", 1)
         self.round: int = 0
         self.game_id: str = f"{self.name}{time.strftime('%y%m%d%H%M%S')}"
@@ -43,6 +44,13 @@ def __init__(self, config: dict):
         self.environment: DockerEnvironment = self.get_environment()
         assert len(config["players"]) >= 2, "At least two players are required"
 
+        # Initialize agents
+        from codeclash.agents import get_agent
+
+        self.agents: list[Player] = []
+        for agent_conf in config["players"]:
+            self.agents.append(get_agent(agent_conf, config["prompts"], self))
+
     @property
     def image_name(self) -> str:
         return f"codeclash/{self.name.lower()}"
@@ -192,6 +200,21 @@ def _post_round_setup(self, agents: list[Player]):
                 )
         self.logger.info(f"Round {self.round} completed.")
 
+    def run(self, cleanup: bool = False):
+        """
+        Run the full game with all configured agents.
+        """
+        try:
+            for _ in range(self.rounds):
+                self.run_round(self.agents)
+                for agent in self.agents:
+                    agent.run()
+        finally:
+            self.end(cleanup)
+            if self.push_agent:
+                for agent in self.agents:
+                    agent.push()
+
     def run_round(self, agents: list[Player]):
         """
         Run a single round of the game with the given agents.

codeclash/games/battlecode/main.py

@@ -8,8 +8,8 @@
 class BattleCodeGame(CodeGame):
     name: str = "BattleCode"
 
-    def __init__(self, config):
-        super().__init__(config)
+    def __init__(self, config, **kwargs):
+        super().__init__(config, **kwargs)
         assert len(config["players"]) == 2, "BattleCode is a two-player game"
         self.run_cmd_round: str = "python run.py run"
         for arg, val in self.game_config.get("args", {}).items():

codeclash/games/battlesnake/main.py

@@ -9,8 +9,8 @@
 class BattleSnakeGame(CodeGame):
     name: str = "BattleSnake"
 
-    def __init__(self, config):
-        super().__init__(config)
+    def __init__(self, config, **kwargs):
+        super().__init__(config, **kwargs)
         self.run_cmd_round: str = "./battlesnake play"
         for arg, val in self.game_config.get("args", {}).items():
             if isinstance(val, bool):

codeclash/games/corewar/main.py

@@ -7,8 +7,8 @@
 class CoreWarGame(CodeGame):
     name: str = "CoreWar"
 
-    def __init__(self, config):
-        super().__init__(config)
+    def __init__(self, config, **kwargs):
+        super().__init__(config, **kwargs)
         self.run_cmd_round: str = "./src/pmars"
         for arg, val in self.game_config.get("args", {}).items():
             if isinstance(val, bool):

codeclash/games/robocode/main.py

@@ -8,8 +8,8 @@
 class RoboCodeGame(CodeGame):
     name: str = "RoboCode"
 
-    def __init__(self, config):
-        super().__init__(config)
+    def __init__(self, config, **kwargs):
+        super().__init__(config, **kwargs)
         self.run_cmd_round: str = "./robocode.sh"
         for arg, val in self.game_config.get("args", {}).items():
             if isinstance(val, bool):

codeclash/games/robotrumble/main.py

@@ -6,8 +6,8 @@
 class RobotRumbleGame(CodeGame):
     name: str = "RobotRumble"
 
-    def __init__(self, config):
-        super().__init__(config)
+    def __init__(self, config, **kwargs):
+        super().__init__(config, **kwargs)
         assert len(config["players"]) == 2, "RobotRumble is a two-player game"
         self.run_cmd_round: str = "./rumblebot run term"
 

main.py

@@ -2,30 +2,15 @@
 
 import yaml
 
-from codeclash.agents import get_agent
-from codeclash.agents.abstract import Player
 from codeclash.games import get_game
 from codeclash.games.abstract import CodeGame
 
 
 def main(config_path: str, cleanup: bool = False, push_agent: bool = False):
     with open(config_path, "r") as f:
         config = yaml.safe_load(f)
-    game: CodeGame = get_game(config)
-    agents: list[Player] = []
-    for agent_conf in config["players"]:
-        agents.append(get_agent(agent_conf, config["prompts"], game))
-
-    try:
-        for _ in range(game.rounds):
-            game.run_round(agents)
-            for agent in agents:
-                agent.run()
-    finally:
-        game.end(cleanup)
-        if push_agent:
-            for agent in agents:
-                agent.push()
+    game: CodeGame = get_game(config, push_agent=push_agent)
+    game.run(cleanup=cleanup)
 
 
 if __name__ == "__main__":