import json import os import re from dataclasses import dataclass, field from typing import Optional, List, Dict from dotenv import load_dotenv from huggingface_hub import InferenceClient from groq import Groq load_dotenv() # ============================================================================= # LLM Configuration - DO NOT MODIFY # ============================================================================= # Model to use (fixed for fair evaluation) LLM_MODEL = "Qwen/Qwen2.5-72B-Instruct" # Initialize the LLM client (uses HF_TOKEN from environment) _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}, ] 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 @dataclass 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) # ============================================================================= # The "State-Injecting" System Prompt # ============================================================================= SYSTEM_PROMPT = """You are an expert text adventure player. OBJECTIVE: Explore, collect treasures, and maximize score. TOOLS 1. play_action: Execute commands (north, take sword, etc.) 2. inventory: Check what you are carrying VALID COMMANDS for play_action (you must use one of these): - Move: n, s, e, w, ne, nw, se, sw, up, down, enter, exit - Perception : look, examine , look into , look under , listen - Action : take, drop, open, close, examine, read, break, climb, unlock, push, pull, burn - Complex: turn on/off , attack with , get with . INTERACTION RULES: 1. EXAMINE + LOOK INTO: you MUST 'look into' AND 'examine' EVERY item (stairs, chest, statue...). 2. TAKE ITEMS: If you see an item, 'take' it immediately. 3. LISTEN: Noise or sound -> 'listen' 4. ANTI-LOOP: if did not work, try , and then move on. EXPLORATION RULES: 1. EXHAUSTIVE SEARCH: Try EVERY direction that is not blocked and not known yet. 2. AWARENESS: If there is a single OBVIOUS direction hinted at, go for it. 3. DIAGONAL: Also try 'ne', 'se', 'nw', 'sw'. 4. VERTICAL: Try 'u', 'd'. RESPONSE FORMAT (Strict JSON-like): THOUGHT: TOOL: ARGS: EXAMPLE: THOUGHT: I see a fountain. I will look into the fountain. TOOL: play_action ARGS: {"action": "look into fountain"} """ ACTIONS_EXTRACTION_PROMPT = """ You are helping an AI agent play a text-adventure game. You receive an observation describing a new location. Your task is to extract all relevant items in this location. INPUT: RULES: 1. ONLY include objects that can be interacted with by the player. 2. NO: field, moon, house... 3. YES: statue, fountain, chest, stairs, doors... 5. Keep the names of the items exactly as they appear in the text. 4. The list of items can be empty. RESPONSE FORMAT (only the list of items, possibly empty): , ..., """ ACTIONS_SUMMARIZATION_PROMPT = """ You are helping an AI agent play a text-adventure game. You receive an observation describing the outcome of an action made by the agent. Summarize this observation. INPUT: -> TASK: 1. Summarize the observation in one short sentence, by insisting on important properties. 2. Remove all the flavor and go straight to the point. 3. If the action has failed or there is nothing worth noting, your summary must be only 'FAILED'. 4. Do not invent objects that are not in the description ! OUTPUT: """ REVERSE_ACTIONS = { "north" : "south", "south" : "north", "east" : "west", "west" : "east", "up" : "down", "down" : "up", "enter" : "exit", "exit" : "enter", "n" : "s", "s" : "n", "e" : "w", "w" : "e", "u" : "d", "d" : "u", "northeast" : "southwest", "northwest" : "southeast", "southeast" : "northwest", "southwest" : "northeast", "ne" : "sw", "nw" : "se", "se" : "nw", "sw" : "ne" } # ============================================================================= # Optimized Agent # ============================================================================= class StudentAgent: def __init__(self): self.history: List[Dict] = [] # Full logs self.score: int = 0 self.visited_locations: set[str] = set() self.last_observation: str = "" self.results : dict[str, dict[str, str]] = dict() self.items : dict[str, str] = dict() self.last_action : str = None self.last_thought : str = None self.descriptions : str = None self.current_location : str = "" self.descriptions : dict[str, str] = dict() self.explored_locations : dict[str, set] = dict() self.tried_directions : dict[str, set] = dict() self.time_in_current_location = 0 async def run(self, client, game: str, max_steps: int, seed: int, verbose: bool = False) -> RunResult: moves = 0 # Initial Look result = await client.call_tool("play_action", {"action": "look"}) observation = self._clean_observation(self._extract_result(result)) self.last_observation = observation self.last_action = "look" new_location = self._extract_location(observation) self.current_location = new_location self.explored_locations[new_location] = set() self.results[new_location] = dict() self.items[new_location] = call_llm(observation, ACTIONS_EXTRACTION_PROMPT, seed) self.descriptions[new_location] = observation if verbose: print(f"START: {observation[:100]}...") for step in range(1, max_steps + 1): self.time_in_current_location += 1 # Update map if self.last_action in [ "north", "south", "east", "west", "up", "down", "enter", "exit", "n", "s", "e", "w", "u", "d", "northeast", "northwest", "southeast", "southwest", "ne", "nw", "se", "sw" ]: new_location = self._extract_location(observation) if new_location != self.current_location: self.items[new_location] = call_llm(observation, ACTIONS_EXTRACTION_PROMPT, 0) self.explored_locations[self.current_location].add(f"{self.last_action} -> {new_location}") if new_location != "Blocked": self.time_in_current_location = 1 if new_location not in self.explored_locations: self.explored_locations[new_location] = set() self.results[new_location] = dict() self.explored_locations[new_location].add(f"{REVERSE_ACTIONS[self.last_action]} -> {self.current_location}") self.current_location = new_location self.descriptions[new_location] = observation else: last_action_and_result = f"{self.last_action} -> {observation}" summary = call_llm(last_action_and_result, ACTIONS_SUMMARIZATION_PROMPT, seed + step) self.results[self.current_location][self.last_action] = summary inventory_result = await client.call_tool("inventory", {}) inventory = self._extract_result(inventory_result) prompt = self._build_prompt(observation, inventory) print("---------------------------------------------------") print(prompt) print("---------------------------------------------------") response = call_llm(prompt, SYSTEM_PROMPT, seed + step) thought, tool_name, tool_args = self._parse_response(response) self.last_thought = thought if verbose: print(f"\n--- Step {step} ---") print(f"Thought: {thought}") print(f"Action: {tool_args}") try: result = await client.call_tool(tool_name, tool_args) action = tool_args.get("action", tool_name) self.last_action = action raw_result = self._extract_result(result) observation = self._clean_observation(raw_result) self._update_score(raw_result) except Exception as e: observation = f"System Error: {e}" self.history.append({ "step": step, "tool": tool_name, "args": tool_args, "result": observation, "thought": thought }) if len(self.history) > 10: self.history.pop(0) moves += 1 if self._is_game_over(observation): break return RunResult( final_score=self.score, max_score=0, # Unknown in generic agent moves=moves, locations_visited=self.visited_locations, game_completed=self._is_game_over(observation) ) def _extract_location(self, observation: str, max_length: int = 25) -> str: lines = observation.strip().split('\n') for line in lines: cleaned_line = line.strip() if not cleaned_line: continue if len(cleaned_line) >= max_length: continue if re.match(r'^[a-zA-Z0-9 ]+$', cleaned_line): return cleaned_line return "Blocked" def _smart_truncate(self, text: str, max_length: int = 80) -> str: """ Truncates text to the nearest sentence ending (., !, ?) before max_length. If no punctuation is found, it falls back to a hard cut. """ # 1. Clean up newlines first clean_text = text.replace("\n", " ").strip() # 2. If it's already short enough, return it if len(clean_text) <= max_length: return clean_text # 3. Take the substring of max_length truncated = clean_text[:max_length] # 4. Find the last sentence ending punctuation # We search for the LAST occurrence of ., !, or ? import re match = re.search(r'[.!?](?!.*[.!?])', truncated) if match: # Cut at the punctuation + 1 (to include the punctuation) return truncated[:match.end()] + "..." # Fallback: If no punctuation found, cut at the last space to avoid splitting a word last_space = truncated.rfind(' ') if last_space != -1: return truncated[:last_space] + "..." return truncated + "..." def get_mini_map(self) -> str: parts = [f"KNOWN CONNECTIONS FROM {self.current_location}:"] for exit in self.explored_locations[self.current_location]: parts.append(f" > {exit}") return "\n".join(parts) def _build_prompt(self, current_obs: str, inventory : str) -> str: """ Constructs a prompt that includes the 'Short Term Memory' so the LLM knows what it just tried. """ parts = [] if self.history: parts.append("\nRECENT HISTORY (Read this to avoid loops!):") for h in self.history[-5:]: action = h['args'].get('action', 'check') parts.append(f" > {action}") ''' res_summary = self._smart_truncate(h['result'], 80) parts.append(f"- Action: {action} -> Result: {res_summary}") ''' if len(self.history) >= 2: last_action = self.history[-1]['args'].get('action') second_last = self.history[-2]['args'].get('action') if last_action == second_last: parts.append("\nWARNING: You just repeated an action. TRY SOMETHING DIFFERENT.\n") if self.time_in_current_location >= 10: parts.append("\nWARNING: You spent the last 10 turns in this location. TRY GOING ELSEWHERE.") parts.append(f"\nIMPORTANT ITEMS IN {self.current_location}:") parts.append(self.items[self.current_location]) parts.append(f"\nYOU ALREADY TRIED:") parts.append(self._format_actions_results()) parts.append(f"\n{self.get_mini_map()}") parts.append(f"\n{inventory}") parts.append(f"\nCURRENT OBSERVATION :") parts.append(current_obs) parts.append("\nBased on the history and observation, what is your next move?") return "\n".join(parts) def _format_actions_results(self) -> str: parts = [] failed = [] for action, result in self.results[self.current_location].items(): if result.startswith("FAILED"): failed.append(action) else: parts.append(f" > {action} -> {result}") if failed: failed_str = "THIS FAILED: " + ", ".join(failed[:5]) + "..." parts.append(failed_str) return "\n".join(parts) def _clean_observation(self, text: str) -> str: """Removes 'Score' lines to prevent LLM confusion.""" text = re.sub(r'\[?Score:.*\]?', '', text, flags=re.IGNORECASE) return text.strip() def _parse_response(self, text: str): """Robust parsing that handles messy LLM output.""" thought = "Deciding next move..." tool_name = "play_action" tool_args = {"action": "look"} # Extract THOUGHT if "THOUGHT:" in text: thought = text.split("THOUGHT:")[1].split("TOOL:")[0].strip() # Extract TOOL if "TOOL:" in text: tool_part = text.split("TOOL:")[1].split("ARGS:")[0].strip() tool_name = tool_part.lower() # Extract ARGS if "ARGS:" in text: args_part = text.split("ARGS:")[1].strip() try: # Try JSON parse tool_args = json.loads(args_part) except: # Fallback regex for simple actions import re match = re.search(r'action["\']?\s*:\s*["\']([^"\']+)["\']', args_part) if match: tool_args = {"action": match.group(1)} return thought, tool_name, tool_args def _extract_result(self, result) -> str: """Helper to get text from MCP result object.""" if hasattr(result, 'content') and result.content: return result.content[0].text return str(result) def _update_score(self, text: str): match = re.search(r'Score:\s*(\d+)', text, re.IGNORECASE) if match: self.score = max(self.score, int(match.group(1))) def _is_game_over(self, text: str) -> bool: return "*** you have died ***" in text.lower() or "game over" in text.lower() 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="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())