feat: add --no-screenshot flag to balatrollm

Author: S1M0N38

src/balatrollm/bot.py

@@ -7,6 +7,7 @@
 from pathlib import Path
 from typing import Any
 
+from balatrobot.enums import State
 from openai import (
     APIConnectionError,
     APIStatusError,
@@ -18,7 +19,6 @@
 from openai.types.chat import ChatCompletion
 
 from balatrobot import BalatroClient, BalatroError
-from balatrobot.enums import State
 
 from .config import Config, load_model_config
 from .data_collection import ChatCompletionError, ChatCompletionResponse, StatsCollector
@@ -227,9 +227,10 @@ async def _make_llm_request_with_retries(
                 request_id = str(time.time_ns() // 1_000_000)
                 response = await self.llm_client.chat.completions.create(**request_data)
                 self.responses.append(response)
-                self.balatro_client.screenshot(
-                    self.data_collector.screenshot_dir / f"{custom_id}.png"
-                )
+                if self.config.take_screenshots:
+                    self.balatro_client.screenshot(
+                        self.data_collector.screenshot_dir / f"{custom_id}.png"
+                    )
                 self.data_collector.write_response(
                     id=str(time.time_ns() // 1_000_000),
                     custom_id=custom_id,
@@ -399,7 +400,7 @@ async def _init_game(self, base_dir: Path = Path.cwd()) -> dict[str, Any]:
             "stake": self.config.stake,
             "seed": self.config.seed,
             "challenge": self.config.challenge,
-            "log_path": self.data_collector.run_dir / "gamestates.jsonl",
+            # "log_path": self.data_collector.run_dir / "gamestates.jsonl",
         }
 
         return self.balatro_client.send_message("start_run", start_run_args)

src/balatrollm/cli.py

@@ -55,6 +55,7 @@ async def cmd_balatrollm(args) -> None:
         config = Config(
             model=args.model,
             strategy=args.strategy,
+            take_screenshots=args.take_screenshots,
         )
 
     # Parse seeds
@@ -149,6 +150,7 @@ async def _port_worker(
             model=config.model,
             strategy=config.strategy,
             seed=seed,
+            take_screenshots=config.take_screenshots,
         )
 
         # Create bot for this run
@@ -284,6 +286,12 @@ def _create_argument_parser() -> argparse.ArgumentParser:
         action="append",
         help="Port for BalatroBot client connection (can specify multiple, default: 12346)",
     )
+    parser.add_argument(
+        "--no-screenshot",
+        action="store_false",
+        dest="take_screenshots",
+        help="Disable taking screenshots during gameplay",
+    )
 
     # Benchmark subcommand
     benchmark_parser = subparsers.add_parser(

src/balatrollm/config.py

@@ -76,6 +76,7 @@ class Config:
         stake: Difficulty stake level (default: 1).
         seed: Game seed for reproducible runs (default: 'OOOO155').
         challenge: Optional challenge mode identifier.
+        take_screenshots: Whether to take screenshots during gameplay (default: True).
         version: Software version string.
         name: Human-readable configuration name.
         description: Configuration description.
@@ -89,6 +90,7 @@ class Config:
     stake: int = 1
     seed: str = "OOOO155"
     challenge: str | None = None
+    take_screenshots: bool = True
     version: str = __version__
     name: str = "Unknown Name"
     description: str = "Unknown Description"
@@ -109,6 +111,7 @@ def from_defaults(cls) -> "Config":
             stake=1,
             seed="OOOO155",
             challenge=None,
+            take_screenshots=True,
             version="",
             name="Unknown Name",
             description="Unknown Description",