Spaces:
Sleeping
Sleeping
| """ | |
| Memory-driven agent with two-phase LLM approach: | |
| 1. Action Selection: Choose promising actions (max 10) and pick best one | |
| 2. Outcome Summarization: Summarize action result for memory storage | |
| Strategy: | |
| - Location memory tracks: valid_actions, tried_actions, promising_actions, results | |
| - Results store: {observation, summary, success, key_info} for each action | |
| - Agent sees concise summaries in context, not overwhelming full observations | |
| - Forces agent to "listen" by summarizing outcomes | |
| """ | |
| import json | |
| import os | |
| import re | |
| from dataclasses import dataclass, field | |
| from typing import Optional, Dict, Set, Any | |
| from dotenv import load_dotenv | |
| from huggingface_hub import InferenceClient | |
| load_dotenv() | |
| # ============================================================================= | |
| # LLM Configuration - DO NOT MODIFY | |
| # ============================================================================= | |
| LLM_MODEL = "Qwen/Qwen2.5-72B-Instruct" | |
| _hf_token = os.getenv("HF_TOKEN") | |
| if not _hf_token: | |
| raise ValueError("HF_TOKEN not found. Set it in your .env file.") | |
| LLM_CLIENT = InferenceClient(token=_hf_token) | |
| def call_llm(prompt: str, system_prompt: str, seed: int, max_tokens: int = 512) -> str: | |
| messages = [ | |
| {"role": "system", "content": system_prompt}, | |
| {"role": "user", "content": prompt}, | |
| ] | |
| response = LLM_CLIENT.chat.completions.create( | |
| model=LLM_MODEL, | |
| messages=messages, | |
| temperature=0.0, | |
| max_tokens=max_tokens, | |
| seed=seed, | |
| ) | |
| return response.choices[0].message.content | |
| class RunResult: | |
| final_score: int | |
| max_score: int | |
| moves: int | |
| locations_visited: set[str] | |
| game_completed: bool | |
| error: Optional[str] = None | |
| history: list[tuple[str, str, str]] = field(default_factory=list) | |
| # ============================================================================= | |
| # System Prompts - Two Phase Approach | |
| # ============================================================================= | |
| ACTION_SELECTION_SYSTEM_PROMPT = """You are playing a text adventure game to maximize score. | |
| GOAL: Explore systematically, solve puzzles, collect items, and maximize score. | |
| YOU WILL RECEIVE: | |
| - Current observation | |
| - Location memory with: | |
| * VALID ACTIONS (from game engine - verified to work at this location) | |
| * TRIED ACTIONS with summaries of outcomes (concise, LLM-generated) | |
| * Previous promising actions (if you've been here before) | |
| YOUR TASK - ACTION SELECTION: | |
| 1. Analyze the valid actions available | |
| 2. Consider which actions you've already tried and their outcomes | |
| 3. Identify up to 10 PROMISING ACTIONS from available options (PROMISING = likely to yield progress based on memory and observation) | |
| 4. Choose one action to try next from the promising action list | |
| STRATEGY GUIDELINES: | |
| - Prioritize untried actions from the valid list | |
| - When uncertain about next steps, prioritize exploration of the current location (move around, examine and interact with objects) | |
| - Pick up valuable items: lamp, lantern, torch, sword, keys, treasures, tools | |
| - If you wander if you have an item for a task, check inventory | |
| - Make sure you have a light sources before entering dark areas; do not leave it unattended once acquired | |
| - PROMISING ACTIONS lead to 1) discovering new locations and 2) interacting with objects meaningfully (examine, take, open, read, push, pull, turn) | |
| - If stagnating (no progress), try different location or different action type | |
| - Learn from previous outcomes: avoid repeating failures unless context changed | |
| - Look for negative game outcomes ("not see that there": object moved - explore elsewhere) | |
| - Follow object/animal/character movements mentioned as important in observations | |
| RESPOND IN THIS EXACT FORMAT (no markdown): | |
| THOUGHT: <your strategic reasoning about the current situation> | |
| PROMISING_ACTIONS: <JSON array of up to 10 promising actions to consider> | |
| REASONING: <evaluate the promising actions and explain which one is best and why> | |
| CHOSEN_ACTION: <the single best action to execute based on your reasoning> | |
| Example: | |
| THOUGHT: I'm in a dark area and need light. Valid actions show "turn on lamp" which I haven't tried yet. | |
| PROMISING_ACTIONS: ["turn on lamp", "examine lamp", "go back", "look", "inventory"] | |
| REASONING: Among these options, "turn on lamp" is best because lamps are critical for exploring dark areas and should be activated before moving forward. "examine lamp" might give info but won't solve the darkness. Going back or just looking won't help without light. | |
| CHOSEN_ACTION: turn on lamp | |
| """ | |
| OUTCOME_SUMMARY_SYSTEM_PROMPT = """You are analyzing the outcome of an action in a text adventure game. | |
| YOUR TASK - OUTCOME SUMMARIZATION: | |
| Given an action and its observation result, create a concise summary that captures: | |
| 1. What happened (1-2 sentences max) | |
| 2. Whether it succeeded, partially succeeded, or failed | |
| 3. Key information to remember for future decisions | |
| Be CONCISE but capture critical details like: | |
| - Items acquired/lost | |
| - New areas discovered | |
| - Interactions with objects/characters | |
| - Obstacles encountered | |
| - Score changes | |
| - Object movements | |
| - State changes (doors opened, lights turned on, etc.) | |
| RESPOND IN THIS EXACT FORMAT (no markdown): | |
| OUTCOME_SUMMARY: <1-2 sentence summary> | |
| SUCCESS: <yes/no/partial> | |
| KEY_INFO: <key detail to remember> | |
| Example action: "take lamp" | |
| Example observation: "Taken. The brass lamp is now in your inventory. [Score: 5 | Moves: 3]" | |
| Example response: | |
| OUTCOME_SUMMARY: Successfully picked up the brass lamp and added it to inventory. | |
| SUCCESS: yes | |
| KEY_INFO: Lamp acquired - can use for dark areas | |
| """ | |
| # ============================================================================= | |
| # Student Agent with Two-Phase LLM Approach | |
| # ============================================================================= | |
| class StudentAgent: | |
| def __init__(self): | |
| self.score = 0 | |
| self.moves = 0 | |
| # Enhanced location memory | |
| self.location_memory: Dict[str, Dict[str, Any]] = {} | |
| # Structure: { | |
| # "Location Name": { | |
| # "valid_actions": [...], | |
| # "tried_actions": set(), | |
| # "promising_actions": [...], | |
| # "visited": count, | |
| # "results": { | |
| # "action": { | |
| # "observation": "full text", | |
| # "summary": "concise summary", | |
| # "success": "yes/no/partial", | |
| # "key_info": "important detail" | |
| # } | |
| # } | |
| # } | |
| # } | |
| # Global tracking | |
| self.locations_visited: set[str] = set() | |
| self.inventory_items: set[str] = set() | |
| # Stagnation detection | |
| self.last_score_change_move: int = 0 | |
| # MCP client handle | |
| self.env_handle = None | |
| def _extract_result(self, result) -> str: | |
| """Extract text from MCP tool result.""" | |
| if hasattr(result, "content") and result.content: | |
| return result.content[0].text | |
| if isinstance(result, list) and result: | |
| return result[0].text if hasattr(result[0], "text") else str(result[0]) | |
| return str(result) | |
| def _parse_action_selection(self, response: str) -> tuple[str, list[str], str, str]: | |
| """ | |
| Parse action selection response. | |
| Returns: (thought, promising_actions, reasoning, chosen_action) | |
| """ | |
| thought = "Proceed with exploration" | |
| promising_actions = [] | |
| reasoning = "" | |
| chosen_action = "look" | |
| for line in (response or "").splitlines(): | |
| line_stripped = line.strip() | |
| line_upper = line_stripped.upper() | |
| if line_upper.startswith("THOUGHT:"): | |
| thought = line_stripped.split(":", 1)[1].strip() | |
| elif line_upper.startswith("PROMISING_ACTIONS:"): | |
| actions_str = line_stripped.split(":", 1)[1].strip() | |
| try: | |
| # Try to parse as JSON array | |
| promising_actions = json.loads(actions_str) | |
| if not isinstance(promising_actions, list): | |
| promising_actions = [] | |
| except: | |
| # Fallback: comma-separated | |
| promising_actions = [a.strip().strip('"\'') for a in actions_str.split(",")] | |
| promising_actions = [a for a in promising_actions if a][:10] | |
| elif line_upper.startswith("REASONING:"): | |
| reasoning = line_stripped.split(":", 1)[1].strip() | |
| elif line_upper.startswith("CHOSEN_ACTION:"): | |
| chosen_action = line_stripped.split(":", 1)[1].strip() | |
| chosen_action = chosen_action.strip('"\'').strip() | |
| # Ensure we have at least something | |
| if not chosen_action or chosen_action == "": | |
| chosen_action = "look" | |
| return thought, promising_actions, reasoning, chosen_action | |
| def _parse_outcome_summary(self, response: str) -> tuple[str, str, str]: | |
| """ | |
| Parse outcome summary response. | |
| Returns: (summary, success, key_info) | |
| """ | |
| summary = "Action executed" | |
| success = "unknown" | |
| key_info = "" | |
| for line in (response or "").splitlines(): | |
| line_stripped = line.strip() | |
| line_upper = line_stripped.upper() | |
| if line_upper.startswith("OUTCOME_SUMMARY:"): | |
| summary = line_stripped.split(":", 1)[1].strip() | |
| elif line_upper.startswith("SUCCESS:"): | |
| success = line_stripped.split(":", 1)[1].strip().lower() | |
| elif line_upper.startswith("KEY_INFO:"): | |
| key_info = line_stripped.split(":", 1)[1].strip() | |
| return summary, success, key_info | |
| def _update_score_moves(self, obs: str) -> None: | |
| """Extract score and moves from observation.""" | |
| m = re.search(r"\[Score:\s*(\d+)\s*\|\s*Moves:\s*(\d+)\]", obs) | |
| if m: | |
| new_score = int(m.group(1)) | |
| if new_score > self.score: | |
| self.last_score_change_move = int(m.group(2)) | |
| self.score = max(self.score, new_score) | |
| self.moves = max(self.moves, int(m.group(2))) | |
| def _is_game_over(self, obs: str) -> bool: | |
| t = (obs or "").lower() | |
| return any(p in t for p in ["game over", "you have died", "you are dead", "*** you have died ***"]) | |
| async def _get_current_location_name(self, client) -> str: | |
| """Get current location name from Jericho's get_player_location().""" | |
| try: | |
| res = await client.call_tool("get_player_location", {}) | |
| loc_info = self._extract_result(res) | |
| return loc_info.strip() | |
| except Exception as e: | |
| return "Unknown" | |
| async def _get_valid_actions_for_location(self, client) -> list[str]: | |
| """Get valid actions from Jericho API.""" | |
| try: | |
| res = await client.call_tool("get_valid_actions", {}) | |
| actions_str = self._extract_result(res) | |
| if actions_str.startswith("{"): | |
| data = json.loads(actions_str) | |
| return data.get("actions", []) | |
| return [a.strip() for a in actions_str.split(",") if a.strip()] | |
| except Exception: | |
| return [] | |
| def _format_location_memory_for_action_selection(self, loc_name: str) -> str: | |
| """Format location memory for action selection prompt.""" | |
| if loc_name not in self.location_memory: | |
| return "=== LOCATION MEMORY: First visit - no memory yet ===\n" | |
| mem = self.location_memory[loc_name] | |
| visit_count = mem["visited"] | |
| valid_actions = mem["valid_actions"] | |
| tried_actions = mem["tried_actions"] | |
| results = mem["results"] | |
| promising_actions = mem.get("promising_actions", []) | |
| parts = [f"=== LOCATION MEMORY: {loc_name} (visited {visit_count} times) ==="] | |
| # Valid actions from game engine | |
| parts.append(f"\nVALID ACTIONS ({len(valid_actions)} available):") | |
| parts.append(f"{', '.join(sorted(valid_actions))}") | |
| parts.append("NOTE: These are location-specific actions verified by game engine.") | |
| parts.append("Universal commands (look, inventory, wait) also work but aren't listed here.") | |
| # Previous promising actions (if any) | |
| if promising_actions: | |
| parts.append(f"\nPREVIOUS PROMISING ACTIONS:") | |
| parts.append(f"{', '.join(promising_actions)}") | |
| # Tried actions with SUMMARIES (not full observations) | |
| if results: | |
| parts.append(f"\nTRIED ACTIONS ({len(tried_actions)} total):") | |
| # Show most recent 8 with summaries | |
| for action, action_result in list(results.items())[-8:]: | |
| summary = action_result.get("summary", "No summary") | |
| success = action_result.get("success", "unknown") | |
| key_info = action_result.get("key_info", "") | |
| status_icon = "+" if success == "yes" else "-" if success == "no" else "~" | |
| parts.append(f" [{status_icon}] {action}") | |
| parts.append(f" → {summary}") | |
| if key_info: | |
| parts.append(f" [*] {key_info}") | |
| else: | |
| parts.append("\nTRIED ACTIONS: (none yet at this location)") | |
| return "\n".join(parts) | |
| def _build_action_selection_prompt(self, observation: str, current_location: str) -> str: | |
| """Build prompt for Phase 1: Action Selection.""" | |
| parts = [] | |
| # Game state | |
| parts.append(f"=== GAME STATE ===") | |
| parts.append(f"Score: {self.score} | Moves: {self.moves}") | |
| parts.append(f"Locations visited: {len(self.locations_visited)}") | |
| if self.inventory_items: | |
| parts.append(f"Inventory: {', '.join(sorted(self.inventory_items))}") | |
| # Location memory (key context) | |
| parts.append("\n" + self._format_location_memory_for_action_selection(current_location)) | |
| # Current observation | |
| parts.append(f"\n=== CURRENT OBSERVATION ===") | |
| parts.append(observation) | |
| # Add strategic hints based on game state | |
| hints = [] | |
| # Stagnation warning | |
| moves_since_progress = self.moves - self.last_score_change_move | |
| if moves_since_progress > 10: | |
| hints.append(f"[!] No score progress in {moves_since_progress} moves!") | |
| hints.append(" Consider: exploring new locations or trying different action types") | |
| # Check for "not see that there" patterns | |
| if current_location in self.location_memory: | |
| mem = self.location_memory[current_location] | |
| recent_failures = sum(1 for action, result in list(mem["results"].items())[-3:] | |
| if "not see that there" in result.get("observation", "").lower()) | |
| if recent_failures >= 2: | |
| hints.append("[!] Multiple 'not see that there' errors - object likely moved elsewhere!") | |
| # Check for object movement in observation | |
| obs_lower = observation.lower() | |
| if any(phrase in obs_lower for phrase in ["ran to", "run to", "went to", "moved to"]): | |
| hints.append("[+] Object movement detected in observation - consider following it!") | |
| if hints: | |
| parts.append("\n=== STRATEGIC HINTS ===") | |
| parts.extend(hints) | |
| parts.append("\n=== YOUR TASK ===") | |
| parts.append("Select up to 10 promising actions and choose the best one to execute.") | |
| return "\n".join(parts) | |
| def _build_outcome_summary_prompt(self, action: str, observation: str) -> str: | |
| """Build prompt for Phase 2: Outcome Summarization.""" | |
| return f"""Action executed: "{action}" | |
| Observation received: | |
| {observation} | |
| Analyze this outcome and provide a concise summary.""" | |
| async def run(self, client, game: str, max_steps: int, seed: int, verbose: bool = False) -> RunResult: | |
| result_history: list[tuple[str, str, str]] = [] | |
| moves_used = 0 | |
| # Discover available tools | |
| tools = await client.list_tools() | |
| tool_names = {t.name for t in tools} | |
| # Initial look | |
| try: | |
| res = await client.call_tool("play_action", {"action": "look"}) | |
| obs = self._extract_result(res) | |
| moves_used += 1 | |
| self._update_score_moves(obs) | |
| except Exception as e: | |
| return RunResult(self.score, 350, moves_used, self.locations_visited, False, error=str(e), history=result_history) | |
| if verbose: | |
| print(f"\n{obs}") | |
| # Initial inventory check | |
| if "inventory" in tool_names: | |
| try: | |
| inv_res = await client.call_tool("inventory", {}) | |
| inv_text = self._extract_result(inv_res).lower() | |
| moves_used += 1 | |
| for item in ["torch", "lamp", "lantern", "sword", "key"]: | |
| if item in inv_text: | |
| self.inventory_items.add(item) | |
| if verbose and self.inventory_items: | |
| print(f"[Starting inventory: {', '.join(self.inventory_items)}]") | |
| except: | |
| pass | |
| # Main game loop | |
| for step in range(1, max_steps - moves_used + 1): | |
| # Get current location | |
| current_location = await self._get_current_location_name(client) | |
| self.locations_visited.add(current_location) | |
| # Initialize or update location memory | |
| if current_location not in self.location_memory: | |
| valid_actions = await self._get_valid_actions_for_location(client) | |
| self.location_memory[current_location] = { | |
| "tried_actions": set(), | |
| "valid_actions": valid_actions, | |
| "promising_actions": [], | |
| "visited": 1, | |
| "results": {} | |
| } | |
| if verbose: | |
| print(f"\n[New location: {current_location}]") | |
| print(f"[Valid actions: {len(valid_actions)}]") | |
| else: | |
| self.location_memory[current_location]["visited"] += 1 | |
| # ======================================================== | |
| # PHASE 1: ACTION SELECTION (LLM Call #1) | |
| # ======================================================== | |
| prompt1 = self._build_action_selection_prompt(obs, current_location) | |
| llm_response1 = call_llm(prompt1, ACTION_SELECTION_SYSTEM_PROMPT, seed + step, max_tokens=768) | |
| thought, promising_actions, reasoning, chosen_action = self._parse_action_selection(llm_response1) | |
| # Store promising actions in memory | |
| self.location_memory[current_location]["promising_actions"] = promising_actions | |
| if verbose: | |
| print(f"\n--- Step {step} ---") | |
| print(f"THOUGHT: {thought}") | |
| print(f"PROMISING: {promising_actions}") | |
| print(f"REASONING: {reasoning}") | |
| print(f"CHOSEN: {chosen_action}") | |
| # ======================================================== | |
| # EXECUTE ACTION | |
| # ======================================================== | |
| try: | |
| res = await client.call_tool("play_action", {"action": chosen_action}) | |
| obs = self._extract_result(res) | |
| moves_used += 1 | |
| # Track score changes | |
| old_score = self.score | |
| self._update_score_moves(obs) | |
| if self.score > old_score: | |
| self.last_score_change_move = moves_used | |
| if verbose: | |
| print(f"Observation: {obs}") | |
| # ======================================================== | |
| # PHASE 2: OUTCOME SUMMARIZATION (LLM Call #2) | |
| # ======================================================== | |
| prompt2 = self._build_outcome_summary_prompt(chosen_action, obs) | |
| llm_response2 = call_llm(prompt2, OUTCOME_SUMMARY_SYSTEM_PROMPT, seed + step + 10000, max_tokens=256) | |
| summary, success, key_info = self._parse_outcome_summary(llm_response2) | |
| if verbose: | |
| print(f"SUMMARY: {summary}") | |
| print(f"SUCCESS: {success}") | |
| if key_info: | |
| print(f"KEY_INFO: {key_info}") | |
| # ======================================================== | |
| # UPDATE MEMORY with summarized outcome | |
| # ======================================================== | |
| mem = self.location_memory[current_location] | |
| mem["tried_actions"].add(chosen_action) | |
| mem["results"][chosen_action] = { | |
| "observation": obs, # Full text preserved | |
| "summary": summary, # LLM-generated summary | |
| "success": success, # yes/no/partial | |
| "key_info": key_info # Important detail | |
| } | |
| # Update inventory tracking | |
| if "take" in chosen_action.lower() and success == "yes": | |
| words = chosen_action.split() | |
| if len(words) >= 2: | |
| item = words[-1] | |
| self.inventory_items.add(item) | |
| # Record in history | |
| result_history.append((thought, f"play_action({chosen_action})", obs)) | |
| if self._is_game_over(obs): | |
| break | |
| except Exception as e: | |
| result_history.append((thought, f"play_action({chosen_action})", f"Error: {e}")) | |
| return RunResult(self.score, 350, moves_used, self.locations_visited, False, error=str(e), history=result_history) | |
| if moves_used >= max_steps: | |
| break | |
| return RunResult( | |
| final_score=self.score, | |
| max_score=350, | |
| moves=moves_used, | |
| locations_visited=self.locations_visited, | |
| game_completed=self._is_game_over(obs), | |
| history=result_history, | |
| ) | |
| # ============================================================================= | |
| # For local testing | |
| # ============================================================================= | |
| async def test_agent(): | |
| """Test the agent locally.""" | |
| from fastmcp import Client | |
| agent = StudentAgent() | |
| async with Client("mcp_server.py") as client: | |
| result = await agent.run( | |
| client=client, | |
| game="lostpig", | |
| max_steps=50, | |
| seed=42, | |
| verbose=True, | |
| ) | |
| print(f"\nFinal Score: {result.final_score}") | |
| print(f"Moves: {result.moves}") | |
| print(f"Locations: {len(result.locations_visited)}") | |
| if __name__ == "__main__": | |
| import asyncio | |
| asyncio.run(test_agent()) | |