agent.py

"""
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 <item>, drop <item>, open <thing>, close <thing>, examine <thing>
- Light: turn on lamp, turn off lamp
- Combat: attack <enemy> with <weapon>
- Other: inventory, look, read <thing>, 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_name>",
  "tool_args": {
    "action": "<game_command>"
  }
}


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())