""" Example: MCP ReAct Agent A complete ReAct agent that uses MCP tools to play text adventure games. This is a working example students can learn from. """ import json import os import re from collections import defaultdict from dataclasses import dataclass, field from typing import Optional from dotenv import load_dotenv from huggingface_hub import InferenceClient # Load environment variables load_dotenv() # Set USE_LOCAL_MODEL=1 in your .env to use a locally downloaded model USE_LOCAL_MODEL = os.getenv("USE_LOCAL_MODEL", "0").strip() in ("1", "true", "yes") LOCAL_MODEL_ID = os.getenv("LOCAL_MODEL_ID", "Qwen/Qwen2.5-3B-Instruct") # ============================================================================= # LLM Configuration - DO NOT MODIFY # ============================================================================= # Model to use (fixed for fair evaluation) LLM_MODEL = "Qwen/Qwen2.5-72B-Instruct" # Initialize the LLM client based on mode _local_pipeline = None if USE_LOCAL_MODEL: import torch from transformers import pipeline as _hf_pipeline _local_pipeline = _hf_pipeline( "text-generation", model=LOCAL_MODEL_ID, torch_dtype=torch.bfloat16, device_map="auto", ) LLM_CLIENT = None else: _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 = 300) -> str: """ Call the LLM with the given prompt. Use this function in your agent. Args: prompt: The user prompt (current game state, history, etc.) system_prompt: The system prompt (instructions for the agent) seed: Random seed for reproducibility max_tokens: Maximum tokens in response (default: 300) Returns: The LLM's response text Example: response = call_llm( prompt="You are in a forest. What do you do?", system_prompt=SYSTEM_PROMPT, seed=42, ) """ messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}, ] if USE_LOCAL_MODEL and _local_pipeline is not None: outputs = _local_pipeline( messages, max_new_tokens=max_tokens, temperature=0.0001, # Near-deterministic (0.0 unsupported by some backends) do_sample=True, ) return outputs[0]["generated_text"][-1]["content"] response = LLM_CLIENT.chat.completions.create( model=LLM_MODEL, messages=messages, temperature=0.0, # Deterministic for reproducibility max_tokens=max_tokens, seed=seed, ) return response.choices[0].message.content @dataclass class RunResult: """Result of running the agent. Do not modify this class.""" 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 Prompt - Customize this for your agent # ============================================================================= SYSTEM_PROMPT = """You are an expert text adventure game player. Your goal is to explore, collect treasures, and maximize your score. AVAILABLE TOOLS: 1. get_valid_actions - Get valid actions for the current location. 2. play_action - Execute game commands (north, take lamp, open mailbox, etc.). 3. memory - Get current game state and score. 4. get_map - See explored locations and connections. 5. inventory - Check what you're carrying. VALID GAME COMMANDS for play_action: - Movement: north, south, east, west, up, down, enter, exit, northeast, northwest, southeast, southwest - Objects: take , drop , open , close , examine - Light: turn on lamp, turn off lamp - Combat: attack with - Other: inventory, look, read , wait CRITICAL INSTRUCTION: You MUST respond ONLY with a valid JSON object. Do not include markdown blocks, pleasantries, or extra text. Use this exact JSON schema: { "reflection": "Briefly analyze the last outcome, current room, and state your strategy.", "global_notes": "(String) Update your ongoing notes. Track unsolved puzzles, locked doors, or cross-room clues. Keep the most critical information that will help you later. This is your persistent memory across rooms. You should rewrite this every turn to keep it up to date. If you have nothing new to add, just repeat the most important notes.", "tool_name": "", "tool_args": { "action": "" } } STRATEGY: 0. DONT WASTE TO MANY STEPS IN A ROOM TRYING TO FIND EVERY SINGLE INTERACTION. IF YOU HAVE TRIED THE MOST OBVIOUS INTERACTIONS AND ACTIONS IN A ROOM AND HAVEN'T FOUND ANY NEW INFORMATION, IT'S OFTEN BEST TO MOVE ON AND EXPLORE NEW AREAS. YOU CAN ALWAYS COME BACK LATER IF YOU FIND CLUES THAT POINT BACK TO THAT ROOM. LIKE GOOD RULE OF THUMB IS 10 STEPS MAX AND 1-4 STEPS IN AVERAGE. 0. Bias heavily towards discovering new locations in the early game to gather information in your memory bank (DONT FORGET TO UPDATE YOUR KNOWLEDGE OF THE GAME) and find useful items. Avoid getting stuck trying to 100% explore a single room before moving on. 0. Use your global_notes to remember things you need to come back to later. 1. Pay attention to the "Known Exits" in your HUD to navigate efficiently. 2. Carefully read the observation for unmapped doors, passages, or directions. 3. If it is dark, prioritize finding or using a light source. 4. If an action fails (e.g., "Need fire for burning"), DO NOT try it again until you have changed your state (e.g., found a fire). THIS IS THE MOST IMPORTANT RULE TO AVOID GETTING STUCK IN LOOPS. If an action doesn't yield new information, try something else to change your state. Also dont try similar actions that are likely to fail for the same reason (e.g. if "take lamp" fails because it's out of reach, "take all" will likely fail too). 5. Do not repeat actions that yielded no results. 6. Avoid immediate backtracking unless you hit a dead end. 7. DISAMBIGUATION: If the game asks a clarification question like "Which do you mean, the X or the Y?" or "What that mean, X or Y?", your next action MUST simply be the exact name of the object you want (e.g., "X" or "Y"). Do not write a full command. 8. TRUST THE HINTS: If a cardinal direction (north, northeast, up, down, etc.) appears in the "Valid verbs/hints" list, YOU CAN GO THAT WAY. Do not assume an exit is blocked just because it isn't explicitly described in the room text. """ # ============================================================================= # Student Agent Implementation # ============================================================================= @dataclass class LocationLog: """Tracks everything known about a specific room/location (Graph Node).""" valid_actions: list[str] = field(default_factory=list) promising_actions: set[str] = field(default_factory=set) action_history: list[dict] = field(default_factory=list) exits: dict[str, str] = field(default_factory=dict) # Graph Edges class StudentAgent: INVALID_VERB_MAP = { "check": "examine", "inspect": "examine", "search": "look", "grab": "take", "pick": "take", "use": "examine", "investigate": "examine", } MOVEMENT_ACTIONS = { "north", "south", "east", "west", "up", "down", "enter", "exit", "northeast", "northwest", "southeast", "southwest", } SCORE_PATTERNS = [r"Score:\s*(\d+)", r"score[:\s]+(\d+)", r"\[Score:\s*(\d+)"] GAME_OVER_PHRASES = [ "game over", "you have died", "you are dead", "*** you have died ***", ] def __init__(self): self.location_logs: defaultdict[str, LocationLog] = defaultdict(LocationLog) self.global_history: list[dict] = [] self.location_path: list[tuple[str, str]] = [] self.current_location: str = "Unknown" self.steps_in_location: int = 0 self.score: int = 0 self.agent_notes: str = "No notes yet. Start exploring!" self.game_name: str = "" async def run( self, client, game: str, max_steps: int, seed: int, verbose: bool = False ) -> RunResult: self.game_name = game locations_visited = set() result_history = [] moves = 0 tools = await client.list_tools() tool_names = [t.name for t in tools] # Initial Observation result = await client.call_tool("play_action", {"action": "verbose"}) observation = self._extract_result(result) initial_memory = await self._get_memory_text(client, tool_names) self.current_location = self._extract_location(observation, initial_memory) locations_visited.add(self.current_location) valid_actions = await self._get_valid_actions(client, tool_names) self.location_logs[self.current_location].valid_actions = valid_actions promising = self._extract_promising_actions(observation, valid_actions, seed) self.location_logs[self.current_location].promising_actions.update(promising) if verbose: print(f"\n{observation}") # Main ReAct Loop for step in range(1, max_steps + 1): prompt = self._build_prompt(observation) response = call_llm(prompt, SYSTEM_PROMPT, seed + step, max_tokens=800) reflection, tool_name, tool_args = self._parse_json_response( response, tool_names ) if verbose: print(f"\n--- Step {step} ---") print(f"[REFLECTION] {reflection}") print(f"[TOOL] {tool_name}({tool_args})") tool_name, tool_args = self._validate_tool_call( tool_name, tool_args, tool_names ) # Loop Detection Check if tool_name == "play_action": action = tool_args.get("action", "look") recent_plays = [ entry.get("args", {}).get("action") for entry in self.global_history[-2:] if entry.get("tool") == "play_action" ] if len(recent_plays) == 2 and all(a == action for a in recent_plays): if verbose: print("[WARNING] Loop detected - forcing 'look'") action = "look" tool_args = {"action": "look"} moves += 1 previous_location = self.current_location attempted_action = ( tool_args.get("action", "") if tool_name == "play_action" else "" ) # Execute Tool try: result = await client.call_tool(tool_name, tool_args) observation = self._extract_result(result) if ( len(observation.strip()) < 30 and tool_name == "play_action" and attempted_action in self.MOVEMENT_ACTIONS ): if verbose: print("[INFO] Brief description detected. Forcing 'look'...") look_result = await client.call_tool( "play_action", {"action": "look"} ) observation = self._extract_result(look_result) if verbose: print(f"[RESULT] {observation[:200]}...") except Exception as e: observation = f"Error: {e}" if verbose: print(f"[ERROR] {e}") # Graph & State Updates memory_text = await self._get_memory_text(client, tool_names) is_new_location, location = self._did_enter_new_location( observation, memory_text ) if is_new_location: if attempted_action in self.MOVEMENT_ACTIONS: self.location_logs[previous_location].exits[attempted_action] = ( location ) self.location_path.append((previous_location, attempted_action)) locations_visited.add(location) valid_actions = await self._get_valid_actions(client, tool_names) self.location_logs[location].valid_actions = valid_actions promising = self._extract_promising_actions( observation, valid_actions, seed + step ) self.location_logs[location].promising_actions.update(promising) # Record action cleanly without an extra LLM call if tool_name == "play_action": # Save just the first 120 chars of the outcome to keep context tight clean_outcome = observation.replace("\n", " ").strip() short_outcome = clean_outcome[:120] + ( "..." if len(clean_outcome) > 120 else "" ) self.location_logs[location].action_history.append( {"action": action, "summary": short_outcome} ) self.global_history.append( { "step": step, "tool": tool_name, "args": tool_args, "result": self._normalize_prompt_text(observation), } ) self._update_score(observation) result_history.append( (reflection, f"{tool_name}({tool_args})", observation[:100]) ) if self._is_game_over(observation): if verbose: print("\n*** GAME OVER ***") break return RunResult( final_score=self.score, max_score=350, moves=moves, locations_visited=locations_visited, game_completed=self._is_game_over(observation), history=result_history, ) def _build_prompt(self, observation: str) -> str: """Builds the HUD (Heads-Up Display) layout for the LLM context.""" parts = [] parts.append(f"You are playing: {self.game_name}") # --- MEMORY BANK --- parts.append("=== YOUR MEMORY BANK (Global Notes) ===") parts.append(self.agent_notes) parts.append("") # --- GAME STATE --- parts.append("=== GAME STATE ===") parts.append(f"Location: {self.current_location}") parts.append(f"Score: {self.score}") parts.append(f"Turns spent in this room: {self.steps_in_location}") if len(self.location_path) > 0: path_strings = [ f"[{room}] --({direction})-->" for room, direction in self.location_path[-4:] ] parts.append( f"Recent Path: {' '.join(path_strings)} [{self.current_location}]" ) if ( len(self.location_path) >= 2 and self.location_path[-1][0] == self.current_location ): parts.append( "\n[WARNING: You are looping between rooms. Explore a new direction!]" ) # --- ROOM KNOWLEDGE --- parts.append("\n=== ROOM KNOWLEDGE ===") current_log = self.location_logs[self.current_location] if current_log.exits: parts.append("Mapped Exits (Already Explored):") for direction, destination in current_log.exits.items(): parts.append(f"- {direction} leads to {destination}") parts.append( "WARNING: Read the room description carefully for unexplored exits!" ) else: parts.append("Known Exits: None confirmed yet.") if current_log.valid_actions: parts.append( "\nValid verbs/hints (from game engine): " + ", ".join(current_log.valid_actions[:12]) ) if current_log.promising_actions: # Create a set of everything we've already tried in this room tried_actions = { attempt["action"].lower().strip() for attempt in current_log.action_history } # Filter the promising actions untested_promising = [ action for action in current_log.promising_actions if action.lower().strip() not in tried_actions ] if untested_promising: parts.append( "Promising targets to interact with: " + ", ".join(untested_promising[:8]) ) else: parts.append( "Promising targets: (You have already tried all obvious interactions here. Look for new exits!)" ) if current_log.action_history: parts.append("\nThings you've already tried here:") # Deduplicate history to prevent "A-B-A" spam from wiping context unique_attempts = {} for attempt in current_log.action_history: unique_attempts[attempt["action"]] = attempt["summary"] # Show the last 6 UNIQUE actions tried in this room for action, summary in list(unique_attempts.items())[-6:]: parts.append(f"- '{action}' -> {summary}") # --- CURRENT OBSERVATION --- parts.append("\n=== CURRENT OBSERVATION ===") current_situation = self._normalize_prompt_text(observation) parts.append(current_situation) parts.append("\n=== ACTION REQUIRED ===") parts.append("What is your next move? Output strictly valid JSON.") print("\n--- PROMPT TO LLM ---") print("\n".join(parts)) print("--- END OF PROMPT ---\n") return "\n".join(parts) def _parse_json_response( self, response: str, valid_tools: list[str] ) -> tuple[str, str, dict]: """Strict JSON parsing to extract reflection, tool, and args.""" reflection = "No reflection provided." tool_name = "play_action" tool_args = {"action": "look"} # Strip out markdown block if the LLM wrapped the JSON clean_response = response.strip() if clean_response.startswith("```json"): clean_response = clean_response[7:] if clean_response.startswith("```"): clean_response = clean_response[3:] if clean_response.endswith("```"): clean_response = clean_response[:-3] try: data = json.loads(clean_response.strip()) reflection = data.get("reflection", reflection) tool_name = data.get("tool_name", tool_name) tool_args = data.get("tool_args", tool_args) if "global_notes" in data: raw_notes = data["global_notes"] if isinstance(raw_notes, list): # If the LLM returned an array, join it into a bulleted string self.agent_notes = "\n".join(f"- {str(note)}" for note in raw_notes) elif isinstance(raw_notes, str) and raw_notes.strip(): # If it's a normal string, just strip it self.agent_notes = raw_notes.strip() except json.JSONDecodeError: # Fallback if the model fails to output valid JSON reflection = f"Failed to parse JSON. Raw output: {response[:100]}..." print(f"[JSON ERROR] Could not parse: {clean_response}") return reflection, tool_name, tool_args def _validate_tool_call( self, tool_name: str, tool_args: dict, valid_tools: list[str] ) -> tuple[str, dict]: if tool_name not in valid_tools: if tool_name in ["action", "do", "command"]: tool_name = "play_action" elif tool_name in ["map", "location"]: tool_name = "get_map" elif tool_name in ["mem", "state", "status"]: tool_name = "memory" elif tool_name in ["inv", "items"]: tool_name = "inventory" else: tool_name = "play_action" if tool_name == "play_action": tool_args["action"] = self._normalize_action( tool_args.get("action", "look") ) return tool_name, tool_args def _normalize_action(self, action: str) -> str: words = action.lower().split() if words and words[0] in self.INVALID_VERB_MAP: words[0] = self.INVALID_VERB_MAP[words[0]] action = " ".join(words) action = ( action.lower().strip().replace("**", "").replace("*", "").replace("`", "") ) return " ".join(action.split()) def _normalize_prompt_text(self, text: str) -> str: if not text: return "" normalized = text.replace("\r\n", "\n") kept_lines = [ line for line in normalized.split("\n") if not re.match(r"^\s*\[Score:\s*\d+\s*\|\s*Moves:\s*\d+\]\s*$", line) ] normalized = "\n".join(kept_lines).strip() return re.sub(r"\n{3,}", "\n\n", normalized) def _extract_result(self, result) -> str: 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 _extract_location(self, observation: str, memory_text: str = "") -> str: # 1. First, try to get the exact location from the memory text if memory_text: match = re.search(r"-\s*Location:\s*(.+)", memory_text) if match and match.group(1).strip(): raw_loc = match.group(1).strip() # USE YOUR NEW METHOD HERE: return self._normalize_location(raw_loc) # 2. Fallback to guessing from the observation text if not observation: return "Unknown" lines = [line.strip() for line in observation.split("\n") if line.strip()] if not lines: return "Unknown" if lines[0].startswith("[") and self.current_location: return self.current_location return lines[0] def _did_enter_new_location( self, observation: str, memory_text: str = "" ) -> tuple[bool, str]: location = self._extract_location(observation, memory_text) if location != self.current_location: self.current_location = location self.steps_in_location = 0 return True, location self.steps_in_location += 1 return False, location async def _get_valid_actions(self, client, tool_names: list[str]) -> list[str]: if "get_valid_actions" not in tool_names: return [] try: raw = await client.call_tool("get_valid_actions", {}) text = self._extract_result(raw) except Exception: return [] actions = [] for line in text.split("\n"): if line.strip().startswith("-"): action = line.strip()[1:].strip().lower() if action and action not in actions: actions.append(action) return actions async def _get_memory_text(self, client, tool_names: list[str]) -> str: if "memory" not in tool_names: return "" try: raw = await client.call_tool("memory", {}) return self._extract_result(raw) except Exception: return "" def _extract_promising_actions( self, observation: str, valid_actions: list[str], seed: int ) -> list[str]: valid_hint = ", ".join(valid_actions[:30]) if valid_actions else "" extraction_prompt = f""" Look for affordances in the observation below. Output ONLY a JSON array of promising text-adventure commands. CRITICAL RULES: 1. If the text mentions doorways, passages, or tunnels in specific directions (north, south, east, west, up, down, northeast, etc.), YOU MUST INCLUDE THOSE DIRECTIONS as commands (e.g., ["north", "southwest"]). 2. Ignore nonsensical actions like "eat pants" or "put torch across wall". Observation: {observation[:800]} """ actions: list[str] = [] try: response = call_llm( prompt=extraction_prompt, system_prompt="You extract promising text-adventure commands and return strict JSON arrays. Output nothing else.", seed=seed, max_tokens=100, ) # Strip markdown if present clean_resp = response.strip() if clean_resp.startswith("```json"): clean_resp = clean_resp[7:] if clean_resp.startswith("```"): clean_resp = clean_resp[3:] if clean_resp.endswith("```"): clean_resp = clean_resp[:-3] parsed = json.loads(clean_resp.strip()) if isinstance(parsed, list): for item in parsed: if isinstance(item, str) and item.strip(): actions.append(item.strip().lower()) except Exception: pass if not actions and valid_actions: priorities = [ "open mailbox", "take lamp", "examine", "read", "north", "south", "east", "west", "up", "down", ] for candidate in priorities: for valid_action in valid_actions: if candidate in valid_action and valid_action not in actions: actions.append(valid_action) break if len(actions) >= 5: break normalized_actions = [] for action in actions: if ( action.strip().lower() and action.strip().lower() not in normalized_actions ): normalized_actions.append(action.strip().lower()) return normalized_actions[:5] def _update_score(self, text: str) -> None: for pattern in self.SCORE_PATTERNS: match = re.search(pattern, text, re.IGNORECASE) if match: self.score = max(self.score, int(match.group(1))) def _is_game_over(self, text: str) -> bool: text_lower = text.lower() return any(phrase in text_lower for phrase in self.GAME_OVER_PHRASES) def _normalize_location(self, raw_location) -> str: """Normalize Jericho location object/string to a clean room name.""" if raw_location is None: return "Unknown" name_attr = getattr(raw_location, "name", None) if isinstance(name_attr, str) and name_attr.strip(): return name_attr.strip() text = str(raw_location).strip() if not text: return "Unknown" # Jericho object repr format example: # Obj93: Outside Parent0 Sibling0 Child87 Attributes ... match = re.search(r":\s*(.+?)\s+Parent\d+", text) if match: candidate = match.group(1).strip() if candidate: return candidate if ":" in text: candidate = text.split(":", 1)[1].strip() if candidate: return candidate.split()[0] return text # ============================================================================= # Local Testing # ============================================================================= async def test_agent(): from fastmcp import Client agent = StudentAgent() async with Client("mcp_server.py") as client: result = await agent.run( client=client, game="zork1", max_steps=20, seed=42, verbose=True, ) print(f"\n{'=' * 50}") print(f"Final 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())