agent.py

"""
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


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


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