diff --git a/agent_strategy.py b/agent_strategy.py index baaf3e0..d1ad4a2 100644 --- a/agent_strategy.py +++ b/agent_strategy.py @@ -4,6 +4,8 @@ from dataclasses import dataclass, field from enum import Enum +from memory_reader import ROUTE_1 + class GamePhase(Enum): """Different phases of the game.""" @@ -249,7 +251,7 @@ def suggest_action_for_goal( # early is a dead end -- runs reached the table row and # stalled there. Leave and go north until Route 1 has # been seen; only then approach the table. - if 0x0B not in self.visited_maps: + if ROUTE_1 not in self.visited_maps: return "DOWN" # Exit the lab (door is at the south) if y > 4: return "UP" @@ -272,7 +274,7 @@ def suggest_action_for_goal( return "RIGHT" # Try different direction return "UP" return "UP" - elif map_id == 0x0B: # Route 1 + elif map_id == ROUTE_1: # Route 1 # Continue north toward Viridian if pos and len(pos) >= 2: if self.is_at_boundary(pos, "UP", map_id): @@ -506,7 +508,7 @@ def check_goal_completion( if "explore_route_1" not in self.completed_goals: map_info = memory_data.get("current_map", {}) map_id = map_info.get("map_id", 0) - if map_id == 0x0B: # Route 1 + if map_id == ROUTE_1: # Route 1 self.mark_goal_complete("explore_route_1") def get_progress_summary(self) -> dict: diff --git a/game_state.py b/game_state.py index becb485..1f6fc1d 100644 --- a/game_state.py +++ b/game_state.py @@ -706,8 +706,10 @@ def execute_action(self, action: str) -> bool: try: # WAIT N is in ticks of 15 frames (~0.25s of game time each); # raw single-frame waits were far too short to let loading - # screens and transitions play out + # screens and transitions play out. Cap N so a model reply + # of WAIT 999999 cannot stall the emulator. units = int(action.split()[1]) if len(action.split()) > 1 else 10 + units = max(0, min(units, 30)) for _ in range(units * 15): self.pyboy.tick() return True diff --git a/memory_reader.py b/memory_reader.py index 7a32b05..c83ae99 100644 --- a/memory_reader.py +++ b/memory_reader.py @@ -30,15 +30,17 @@ class MemoryAddresses: PARTY_POKEMON_START = 0xD16B # First Pokemon data structure POKEMON_DATA_SIZE = 44 # Size of each Pokemon data structure - # Pokemon data structure offsets (relative to party Pokemon start) + # Pokemon data structure offsets (relative to party Pokemon start). + # pret pokered party_struct: species, current HP, box level, status, + # types, ... level at 33, max HP at 34. Offset 3 is box level, not max HP. POKEMON_SPECIES = 0 # Pokemon species ID POKEMON_HP_CURRENT = 1 # Current HP (2 bytes, little-endian) - POKEMON_HP_MAX = 3 # Max HP (2 bytes, little-endian) - POKEMON_STATUS = 5 # Status condition (0x00 = normal) - POKEMON_TYPE1 = 6 - POKEMON_TYPE2 = 7 + POKEMON_STATUS = 4 # Status condition (0x00 = normal) + POKEMON_TYPE1 = 5 + POKEMON_TYPE2 = 6 + POKEMON_EXP = 14 # Experience (3 bytes) POKEMON_LEVEL = 33 # Level - POKEMON_EXP = 34 # Experience (3 bytes) + POKEMON_HP_MAX = 34 # Max HP (2 bytes, little-endian) # Inventory INVENTORY_COUNT = 0xD31D # Number of items in inventory @@ -386,33 +388,38 @@ def read_full_game_state(self) -> dict[str, any]: 0x08: "Cinnabar Island", 0x09: "Indigo Plateau", 0x0A: "Saffron City", - 0x0B: "Route 1", - 0x0C: "Route 2", - 0x0D: "Route 3", - 0x0E: "Route 4", - 0x0F: "Route 5", - 0x10: "Route 6", - 0x11: "Route 7", - 0x12: "Route 8", - 0x13: "Route 9", - 0x14: "Route 10", - 0x15: "Route 11", - 0x16: "Route 12", - 0x17: "Route 13", - 0x18: "Route 14", - 0x19: "Route 15", - 0x1A: "Route 16", - 0x1B: "Route 17", - 0x1C: "Route 18", - 0x1D: "Route 19", - 0x1E: "Route 20", - 0x1F: "Route 21", - 0x20: "Route 22", - 0x21: "Route 23", - 0x22: "Route 24", - 0x23: "Route 25", + # 0x0B is UNUSED_MAP_0B in pret pokered map_constants.asm. + # Route 1 is 0x0C; Route 25 is 0x24. + 0x0C: "Route 1", + 0x0D: "Route 2", + 0x0E: "Route 3", + 0x0F: "Route 4", + 0x10: "Route 5", + 0x11: "Route 6", + 0x12: "Route 7", + 0x13: "Route 8", + 0x14: "Route 9", + 0x15: "Route 10", + 0x16: "Route 11", + 0x17: "Route 12", + 0x18: "Route 13", + 0x19: "Route 14", + 0x1A: "Route 15", + 0x1B: "Route 16", + 0x1C: "Route 17", + 0x1D: "Route 18", + 0x1E: "Route 19", + 0x1F: "Route 20", + 0x20: "Route 21", + 0x21: "Route 22", + 0x22: "Route 23", + 0x23: "Route 24", + 0x24: "Route 25", } +# pret pokered: Route 1 is map 0x0C (0x0B is unused). +ROUTE_1 = 0x0C + def get_map_name(map_id: int) -> str: """Get map name from map ID. diff --git a/ocr_enhancer.py b/ocr_enhancer.py index 58e7341..fadd35c 100644 --- a/ocr_enhancer.py +++ b/ocr_enhancer.py @@ -256,24 +256,23 @@ def extract_text_enhanced( reverse=True, ) - # Extract text from each region - texts = [] + # Keep each string paired with the region that produced it. Indexing + # the original region list with the filtered text list mis-labels + # dialog after a short OCR result is dropped. + extracted: list[tuple[dict, str]] = [] for region in regions: text = self.extract_text_from_region(image, region) if text and len(text.strip()) > 2: - texts.append(text) - - # Combine texts, prioritizing dialog - if texts: - # If we have dialog text, use it primarily - dialog_texts = [ - t - for i, t in enumerate(texts) - if regions[i]["type"] == TextRegion.DIALOG_BOX - ] - if dialog_texts: - return " ".join(dialog_texts) - else: - return " ".join(texts) - - return "" + extracted.append((region, text)) + + if not extracted: + return "" + + dialog_texts = [ + text + for region, text in extracted + if region["type"] == TextRegion.DIALOG_BOX + ] + if dialog_texts: + return " ".join(dialog_texts) + return " ".join(text for _, text in extracted) diff --git a/pokemon_agent.py b/pokemon_agent.py index 9f0454c..f0a7551 100644 --- a/pokemon_agent.py +++ b/pokemon_agent.py @@ -3,9 +3,9 @@ Version: 0.0.7 """ -import contextlib import logging import random +import re import time from collections import Counter, deque from dataclasses import dataclass @@ -20,6 +20,13 @@ logger = logging.getLogger(__name__) +# Whole-word button tokens. Exact token match keeps START from collapsing +# into A ("A" is a substring of "START") and ignores prose like "GO EAST". +_BUTTON_TOKENS = frozenset({"START", "SELECT", "RIGHT", "LEFT", "DOWN", "UP", "A", "B"}) +# WAIT N is applied as N * 15 frames. Cap N so a model reply cannot stall the emulator. +_MAX_WAIT_UNITS = 30 +_BUTTON_WORD = re.compile(r"[A-Z]+") + @dataclass class Observation: @@ -864,7 +871,6 @@ def _dialog_shortcut_policy(self, obs: "Observation") -> str | None: def _llm_policy(self, obs: "Observation") -> str: """Ask the LLM for an action; falls back to heuristics on failure.""" - llm_start_time = None try: prompt = self.get_prompt() @@ -877,16 +883,11 @@ def _llm_policy(self, obs: "Observation") -> str: ) llm_duration = time.time() - llm_start_time - # Record LLM metrics + # Call counts live on the provider (it has token and error detail). + # Recording again here doubled total_calls and distorted latency. if self.metrics: self.metrics.performance.record_llm_time(llm_duration) - self.metrics.llm.record_call(llm_duration) except Exception as e: - llm_duration = ( - time.time() - llm_start_time if llm_start_time is not None else 0.0 - ) - if self.metrics: - self.metrics.llm.record_call(llm_duration, error=True) logger.warning(f"LLM call failed: {e}") logger.warning("Using fallback action based on game state") return self._llm_fallback(obs) @@ -913,50 +914,56 @@ def _llm_fallback(self, obs: "Observation") -> str: # during character creation) return "A" + def _match_button_token(self, text: str) -> str | None: + """First whole-word button in text. + + Substring matching treated START as A because ``"A" in "START"``. + Tokens are matched exactly, so START and SELECT stay intact and + prose like "GO EAST" does not become A. + """ + for token in _BUTTON_WORD.findall(text): + if token in _BUTTON_TOKENS: + return token + return None + + def _match_wait_units(self, text: str) -> int | None: + """WAIT count from a reply, clamped to ``_MAX_WAIT_UNITS``.""" + match = re.search(r"\bWAIT(?:\s+(\d+))?\b", text) + if match is None: + return None + if match.group(1) is None: + return 10 + return min(int(match.group(1)), _MAX_WAIT_UNITS) + + def _block_creation_b(self, obs: "Observation", button: str) -> str: + """B cancels a new game; substitute A while naming is in progress.""" + if button == "B" and self._in_creation_window(): + logger.info( + f"[CHARACTER_CREATION] Blocked B button, using A instead (step {obs.step_count})" + ) + return "A" + return button + def _parse_llm_response(self, obs: "Observation", response: str) -> str: """Extract a valid button action from the LLM's raw response.""" response_clean = response.strip().upper() - # CRITICAL: Final check - prevent B during character creation - if self._in_creation_window() and ( - response_clean == "B" or response_clean.startswith("B") - ): - logger.info( - f"[CHARACTER_CREATION] Blocked B press from LLM, using A instead (step {obs.step_count})" - ) - response_clean = "A" - - valid_buttons = ["UP", "DOWN", "LEFT", "RIGHT", "A", "B", "SELECT", "START"] - - # Check for direct match first (most common case) - for button in valid_buttons: - if button in response_clean: - # CRITICAL: Block B during character creation even if found in response - if button == "B" and self._in_creation_window(): - logger.info( - f"[CHARACTER_CREATION] Blocked B button, using A instead (step {obs.step_count})" - ) - return self._cache_action(obs, "A") - # Extract the button - if "," in response_clean: - # Handle comma-separated - parts = [p.strip() for p in response_clean.split(",")] - valid_parts = [p for p in parts if p in valid_buttons] - action = ( - ", ".join(valid_parts[:2]) if valid_parts else button - ) # Max 2 actions - else: - action = button - return self._cache_action(obs, action) - - # Check for WAIT command - if "WAIT" in response_clean: - parts = response_clean.split() - action = "WAIT 10" - if len(parts) >= 2: - with contextlib.suppress(ValueError): - action = f"WAIT {int(parts[1])}" - return self._cache_action(obs, action) + if "," in response_clean: + buttons = [] + for part in response_clean.split(","): + button = self._match_button_token(part) + if button is not None: + buttons.append(self._block_creation_b(obs, button)) + if buttons: + return self._cache_action(obs, ", ".join(buttons[:2])) + else: + button = self._match_button_token(response_clean) + if button is not None: + return self._cache_action(obs, self._block_creation_b(obs, button)) + + wait_units = self._match_wait_units(response_clean) + if wait_units is not None: + return self._cache_action(obs, f"WAIT {wait_units}") # Fallback: START while still orienting, A afterwards action = "START" if len(self.action_history) < 5 else "A" @@ -1008,29 +1015,12 @@ def step(self) -> dict: completed_goals=completed, ) - # Check if strategy suggests an action + # Do not apply a strategy suggestion here. suggest_action_for_goal + # returns raw A/DOWN/START for menus and dialog, and taking it + # skipped the policy chain (naming scripts, menu escape, dialog + # loops) on most steps. Overworld routing still runs inside + # get_action() via _route_policy. action = None - if self.strategy: - current_goal = self.strategy.get_current_goal() - if current_goal: - memory_data = { - "player_position": pre_state.get("player_position"), - "current_map": pre_state.get("current_map", {}), - "party": pre_state.get("party", []), - "health": pre_state.get("health", {}), - } - suggested_action = self.strategy.suggest_action_for_goal( - current_goal, pre_game_state, memory_data - ) - # Use strategy suggestion if available and not in pure exploration mode - # Calculate action diversity for adaptive exploration - action_diversity = ( - self.calculate_action_entropy(window=15) / 3.0 - ) # Normalize to 0-1 range - if suggested_action and not self.strategy.should_explore( - self.stuck_count, action_diversity - ): - action = suggested_action # Force exploration when stuck (before getting action). # Overworld only: random arrows inside a menu, dialog, or transition diff --git a/tests/conftest.py b/tests/conftest.py index cdc0bd8..d270835 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -74,11 +74,13 @@ def sample_memory_data(): 0xD163: 1, # 1 Pokemon in party # First Pokemon data (simplified) 0xD16B: 25, # Species (Pikachu) - 0xD16C: 0x20, # HP current low byte + 0xD16C: 0x20, # HP current low byte (offset 1) 0xD16D: 0x03, # HP current high byte (0x0320 = 800) - 0xD16E: 0x20, # HP max low byte - 0xD16F: 0x03, # HP max high byte (0x0320 = 800) - 0xD19C: 50, # Level (0xD16B + 33 = 0xD19C) + 0xD16E: 0x00, # Box level (offset 3), not max HP + 0xD16F: 0x00, # Status (offset 4) + 0xD18C: 50, # Level (0xD16B + 33) + 0xD18D: 0x20, # Max HP low byte (offset 34) + 0xD18E: 0x03, # Max HP high byte (0x0320 = 800) # Inventory 0xD31D: 2, # 2 items 0xD31E: 1, # Item 1 ID diff --git a/tests/test_agent_strategy_route.py b/tests/test_agent_strategy_route.py index e5b7828..cca4108 100644 --- a/tests/test_agent_strategy_route.py +++ b/tests/test_agent_strategy_route.py @@ -1,6 +1,7 @@ """Route-to-Oak direction tests for the get_starter goal.""" from agent_strategy import AgentStrategy +from memory_reader import ROUTE_1 def suggest(map_id, pos, visited=()): @@ -32,9 +33,9 @@ def test_pallet_town_heads_north(): def test_oaks_lab_heads_to_ball_table_then_interacts(): - # Only after Route 1 (0x0B) has been seen — that is when Oak's cutscene + # Only after Route 1 (0x0C) has been seen — that is when Oak's cutscene # has fired and the balls become takeable. - r1 = (0x0B,) + r1 = (ROUTE_1,) assert suggest(0x28, (4, 11), r1) == "UP" # deep in the lab: go up assert suggest(0x28, (4, 4), r1) == "RIGHT" # at table row: go right assert suggest(0x28, (7, 4), r1) == "A" # at the table: interact @@ -50,7 +51,7 @@ def test_lab_is_a_dead_end_before_the_oak_cutscene(): def test_lab_approaches_table_after_route_1_seen(): s = AgentStrategy() - s.visited_maps.add(0x0B) # Route 1 visited -> Oak cutscene has fired + s.visited_maps.add(ROUTE_1) # Route 1 visited -> Oak cutscene has fired goal = next(g for g in s.goals if g.name == "get_starter") md = {"current_map": {"map_id": 0x28}, "player_position": (4, 11), "party": []} assert s.suggest_action_for_goal(goal, "overworld", md) == "UP" @@ -165,12 +166,12 @@ def test_step_updates_map_tracking(mock_llm_provider, mock_pyboy): "game_state": "overworld", "party": [], "player_position": (5, 6), - "current_map": {"map_id": 0x0B, "map_name": "Route 1"}, + "current_map": {"map_id": ROUTE_1, "map_name": "Route 1"}, } ) gs.execute_action = Mock(return_value=True) agent.step() - assert 0x0B in agent.strategy.visited_maps + assert ROUTE_1 in agent.strategy.visited_maps class TestEdgeScan: diff --git a/tests/test_game_state.py b/tests/test_game_state.py index 62da811..da8b24b 100644 --- a/tests/test_game_state.py +++ b/tests/test_game_state.py @@ -145,6 +145,15 @@ def test_execute_action_wait(self, mock_pyboy): assert result is True assert mock_pyboy.tick.call_count == 5 * 15 # WAIT units are 15 frames each + def test_execute_action_wait_is_capped(self, mock_pyboy): + """Huge WAIT counts must not stall the emulator.""" + game_state = GameState(mock_pyboy, ocr_enabled=False) + + result = game_state.execute_action("WAIT 999999") + + assert result is True + assert mock_pyboy.tick.call_count == 30 * 15 + def test_execute_action_invalid(self, mock_pyboy): """Test executing invalid action.""" game_state = GameState(mock_pyboy, ocr_enabled=False) diff --git a/tests/test_memory_reader.py b/tests/test_memory_reader.py index cb623bc..05476a4 100644 --- a/tests/test_memory_reader.py +++ b/tests/test_memory_reader.py @@ -149,12 +149,6 @@ def getitem(address): mock_pyboy.memory.__getitem__ = Mock(side_effect=getitem) - # Level is at PARTY_POKEMON_START + POKEMON_LEVEL offset - # PARTY_POKEMON_START = 0xD16B, POKEMON_LEVEL = 33 - # So level address = 0xD16B + 33 = 0xD18C (not 0xD19C in sample data) - # Update sample data to have level at correct address - sample_memory_data[0xD18C] = 50 # Level at correct offset (0xD16B + 33) - reader = MemoryReader(mock_pyboy) party = reader.read_pokemon_party() @@ -162,9 +156,13 @@ def getitem(address): assert party[0]["species"] == 25 assert party[0]["level"] == 50 assert party[0]["hp_current"] == 800 + # Max HP is party_struct offset 34, not the box-level byte at offset 3. assert party[0]["hp_max"] == 800 + assert party[0]["status"] == 0 assert party[0]["hp_percent"] == 100.0 assert not party[0]["fainted"] + assert MemoryAddresses.POKEMON_HP_MAX == 34 + assert MemoryAddresses.POKEMON_STATUS == 4 def test_read_pokemon_party_empty(self, mock_pyboy): """Test reading empty party.""" @@ -308,7 +306,10 @@ def test_get_map_name(self): """Test getting map name from ID.""" assert get_map_name(0x00) == "Pallet Town" assert get_map_name(0x01) == "Viridian City" - assert get_map_name(0x0B) == "Route 1" + # pret map_constants: 0x0B is unused, Route 1 is 0x0C, Route 25 is 0x24 + assert get_map_name(0x0C) == "Route 1" + assert get_map_name(0x24) == "Route 25" + assert get_map_name(0x0B) == "Map 0B" def test_get_map_name_unknown(self): """Test getting name for unknown map ID.""" diff --git a/tests/test_metrics_integration.py b/tests/test_metrics_integration.py index 2a27c34..c13dd26 100644 --- a/tests/test_metrics_integration.py +++ b/tests/test_metrics_integration.py @@ -142,6 +142,40 @@ def mock_get_info(): assert llm_stats["total_calls"] >= 0 assert llm_stats["latency"]["avg"] >= 0 + def test_llm_call_recorded_once(self, mock_pyboy): + """One generate() records one LLM call, not one in the provider and one in the agent.""" + metrics = MetricsCollector() + with patch.object( + OllamaProvider, "_list_available_models", return_value=["llama3.2"] + ): + llm_provider = OllamaProvider(model="llama3.2", metrics=metrics) + llm_provider.client = Mock() + llm_provider.client.chat = Mock( + return_value={"message": {"content": "UP"}, "eval_count": 5} + ) + + game_state = GameState(mock_pyboy, ocr_enabled=False, metrics=metrics) + agent = PokemonAgent(llm_provider, game_state, metrics=metrics) + game_state.get_game_info = Mock( + return_value={ + "screen_text": "open grass", + "frame_count": 40, + "game_state": "overworld", + "party": [{"species": 1, "level": 6}], + "player_position": (3, 3), + "current_map": {}, + } + ) + game_state.execute_action = Mock(return_value=True) + agent.step_count = 10 + agent.new_game_started = True + agent._creation_over = True + + agent.get_action() + + assert metrics.get_all_stats()["llm"]["total_calls"] == 1 + assert llm_provider.client.chat.call_count == 1 + def test_metrics_summary_generation(self): """Test that metrics summary can be generated.""" metrics = MetricsCollector() diff --git a/tests/test_ocr_enhancer.py b/tests/test_ocr_enhancer.py new file mode 100644 index 0000000..22b12a9 --- /dev/null +++ b/tests/test_ocr_enhancer.py @@ -0,0 +1,26 @@ +"""Tests for OCR region association.""" + +from unittest.mock import Mock + +import numpy as np + +from ocr_enhancer import OCREnhancer, TextRegion + + +def test_dialog_text_stays_paired_after_short_reads_are_dropped(): + """A dropped short read must not shift later strings onto earlier regions.""" + enhancer = OCREnhancer() + image = np.zeros((144, 160, 3), dtype=np.uint8) + regions = [ + {"type": TextRegion.MENU, "confidence": 0.2, "bbox": (0, 0, 10, 10)}, + {"type": TextRegion.MENU, "confidence": 0.4, "bbox": (0, 20, 10, 10)}, + {"type": TextRegion.DIALOG_BOX, "confidence": 0.9, "bbox": (0, 104, 40, 20)}, + ] + enhancer.detect_text_regions = Mock(return_value=regions) + enhancer.extract_text_from_region = Mock( + side_effect=["HI", "OPEN MENU", "OAK SPEAKS"] + ) + + text = enhancer.extract_text_enhanced(image, prioritize_dialog=False) + + assert text == "OAK SPEAKS" diff --git a/tests/test_pokemon_agent.py b/tests/test_pokemon_agent.py index 7f9b25b..61bf3c7 100644 --- a/tests/test_pokemon_agent.py +++ b/tests/test_pokemon_agent.py @@ -3,7 +3,7 @@ from unittest.mock import Mock from game_state import GameState -from pokemon_agent import PokemonAgent +from pokemon_agent import Observation, PokemonAgent class TestPokemonAgent: @@ -406,3 +406,85 @@ def test_counter_resets_when_position_changes(self, mock_llm_provider, mock_pybo ) agent.get_action() assert agent._dialog_loop_steps == 1 + + +class TestPolicyChainNotSkipped: + """step() must run early-game and menu policies, not a raw strategy button.""" + + def test_step_uses_naming_script(self, mock_llm_provider, mock_pyboy): + game_state = GameState(mock_pyboy, ocr_enabled=False) + agent = PokemonAgent(mock_llm_provider, game_state) + game_state.get_game_info = Mock( + return_value={ + "screen_text": "NEW NAME RED ASH JACK", + "frame_count": 100, + "game_state": "menu", + "party": [], + "player_position": (0, 0), + } + ) + game_state.execute_action = Mock(return_value=True) + + result = agent.step() + + # Strategy's start_game menu suggestion is A. The naming script is DOWN. + assert result["action"] == "DOWN" + + def test_step_closes_lingering_menu(self, mock_llm_provider, mock_pyboy): + game_state = GameState(mock_pyboy, ocr_enabled=False) + agent = PokemonAgent(mock_llm_provider, game_state) + game_state.get_game_info = Mock( + return_value={ + "screen_text": "", + "frame_count": 100, + "game_state": "menu", + "party": [{"species": 1, "level": 6}], + "player_position": (5, 5), + } + ) + game_state.execute_action = Mock(return_value=True) + agent.new_game_started = True + agent.character_creation_steps = 60 + agent.step_count = 200 + + actions = [agent.step()["action"] for _ in range(3)] + + # Strategy would keep returning DOWN. Menu escape presses B. + assert actions[-1] == "B" + + +class TestParseLlmResponse: + """Button tokens are whole words, and WAIT is capped.""" + + def _agent(self, mock_llm_provider, mock_pyboy): + game_state = GameState(mock_pyboy, ocr_enabled=False) + return PokemonAgent(mock_llm_provider, game_state) + + def _obs(self) -> Observation: + return Observation( + game_info={"frame_count": 1}, + screen_text="", + game_state="title_screen", + step_count=10, + ) + + def test_start_is_not_a(self, mock_llm_provider, mock_pyboy): + agent = self._agent(mock_llm_provider, mock_pyboy) + obs = self._obs() + assert agent._parse_llm_response(obs, "START") == "START" + assert agent._parse_llm_response(obs, "PRESS START") == "START" + assert agent._parse_llm_response(obs, "GO EAST") == "START" + + def test_prose_without_a_button_falls_back(self, mock_llm_provider, mock_pyboy): + agent = self._agent(mock_llm_provider, mock_pyboy) + agent.action_history = ["A", "A", "A", "A", "A"] + assert agent._parse_llm_response(self._obs(), "GO EAST") == "A" + + def test_wait_is_capped(self, mock_llm_provider, mock_pyboy): + agent = self._agent(mock_llm_provider, mock_pyboy) + assert agent._parse_llm_response(self._obs(), "WAIT 999999") == "WAIT 30" + assert agent._parse_llm_response(self._obs(), "WAIT") == "WAIT 10" + + def test_comma_separated_keeps_two_buttons(self, mock_llm_provider, mock_pyboy): + agent = self._agent(mock_llm_provider, mock_pyboy) + assert agent._parse_llm_response(self._obs(), "UP, DOWN, LEFT") == "UP, DOWN"