3️⃣ Building a more complex agent: WikiGame

# Retrieve a Wikipedia page from its title
page = wikipedia.page("Large language model")

# Access basic page information
print("Title:", page.title)
print("\nURL", page.url)
print(f"\nSummary (word count {len(page.summary.split())}):", page.summary)
print(
    f"\nContent (word count {len(page.content.split())}):",
    page.content[:1000],
    "......",
)
print(f"""\nLinks (link count {len(page.links)}): [{", ".join(page.links[:7])}, ......]""")

try:
    page = wikipedia.page("Python")
except DisambiguationError as e:
    print(type(e), "\n\n", e)

try:
    page = wikipedia.page("Animalss", auto_suggest=False)
except Exception as e:
    print(type(e), "\n\n", e)

# Fixes PageError by allowing redirects
page = wikipedia.page("Animalss", redirect=True)
print(page.title)

# Fixes DisambiguationError by selecting the first option

try:
    page = wikipedia.page("Python")
except DisambiguationError as e:
    page = wikipedia.page(e.options[0])
print(page.title)

def get_page(title: str) -> WikipediaPage:
    """
    Get a Wikipedia page object given a title. If the title is ambiguous, choose the first option.
    If the title is not found, try to find a similar title.

    Args:
        title (str): The title of the Wikipedia page

    Returns:
        WikipediaPage: The Wikipedia page
    """
    try:
        return wikipedia.page(title, auto_suggest=False, redirect=True)
    except DisambiguationError as e:
        return wikipedia.page(e.options[0], auto_suggest=False, redirect=True)
    except PageError:
        return wikipedia.page(title, auto_suggest=True, redirect=True)

def get_permitted_links(current_page: WikipediaPage) -> list[str]:
    """
    Get "permitted" links (i.e. links that are in the content of the page) from a Wikipedia page.

    Args:
        current_page (WikipediaPage): The current Wikipedia page

    Returns:
        list[str]: A list of permitted links from current_page

    """
    raise NotImplementedError("You need to implement the get_permitted_links function")


tests.test_get_permitted_links(get_permitted_links)

class WikiGame:
    def __init__(
        self,
        starting_page: str,
        goal_page: str,
    ):
        """
        This task simulates the Wikipedia game, where the agent starts on one Wikipedia page and
        attempts to navigate to a goal page using only links found in the main content of Wikipedia
        pages.

        Args:
            starting_page (str): The page the agent starts on.
            goal_page (str): The page the agent is trying to reach.

        Attributes:
            page_history (list[str]): The history of pages visited by the agent.
            starting_page (WikipediaPage): The starting page of the game.
            goal_page (WikipediaPage): The goal page of the game.
            current_page (WikipediaPage): The current page the agent is on.

        """
        self.page_history: list[str] = [starting_page]
        self.starting_page: WikipediaPage = self.get_page(starting_page)
        self.goal_page: WikipediaPage = self.get_page(goal_page)
        self.current_page: WikipediaPage = self.starting_page

    # ========================= Helper Functions (given) =========================

    # Get page and page summary
    @staticmethod
    def get_page(title: str) -> WikipediaPage:
        """
        Get a Wikipedia page object given a title. If the title is ambiguous, choose the first
        option. If the title is not found, try to find a similar title.

        Args:
            title (str): The title of the Wikipedia page

        Returns:
            WikipediaPage: The Wikipedia page
        """
        try:
            return wikipedia.page(title, auto_suggest=False, redirect=True)
        except DisambiguationError as e:
            return wikipedia.page(e.options[0], auto_suggest=False, redirect=True)
        except PageError:
            return wikipedia.page(title, auto_suggest=True, redirect=True)

    def get_page_summary(self, page: WikipediaPage | None = None) -> str:
        """
        Get summary of a wikipedia page, to the last full stop within the first 500 characters.
        This can be used to give a brief overview of a page to the agent.

        Args:
            page (WikipediaPage): The Wikipedia page object.

        Returns:
            str: The summary of the Wikipedia page.
        """
        page = page if page else self.goal_page
        summary = page.content[:500]
        last_period_index = summary.rfind(".")
        return summary[: last_period_index + 1] if last_period_index != -1 else summary

    # Get and check permitted links
    def get_permitted_links(self) -> list[str]:
        """
        Returns a list of permitted links (i.e. links in the main page content) for the current page.

        Returns:
            list[str]: The permitted links.
        """
        all_links = self.current_page.links
        content_lower = self.current_page.content.lower()
        permitted_links = [link for link in all_links if link.lower() in content_lower]
        if self.current_page.title in permitted_links:
            permitted_links.remove(self.current_page.title)
        return permitted_links

    def is_permitted_link(self, link: str) -> bool:
        """
        Returns True if the link is in the permitted links for the current page, False otherwise.

        Args:
            link (str): The link to check.

        Returns:
            bool: True if the link is permitted, False otherwise
        """
        return link.lower() in (x.lower() for x in self.get_permitted_links())

    # ========================= Task State Management (given) =========================

    def check_win(self) -> bool:
        return self.current_page == self.goal_page

class GetContentTool:
    """
    The GetContentTool retrieves the full content of the current Wikipedia page, marking all
    accessible links within the main content by wrapping them in <link></link> tags.

    This is an example of a tool that provides the agent with detailed page content to enable
    reasoning about possible next steps in the Wikipedia game.
    """

    name = "get_content"

    @staticmethod
    def execute(task: WikiGame) -> str:
        """
        Get all the content for the wikipedia page you are currently on. Anything which corresponds
        to a link is wrapped in <link></link> tags.

        Args:
            task (WikiGame): The current task object.

        Returns:
            str: The content of the page with links wrapped
        """
        content = task.current_page.content
        permitted_links = get_permitted_links(task.current_page)
        for word in sorted(permitted_links, key=len, reverse=True):
            content = re.sub(
                r"""(\s|[,.)!?;:'"])(""" + re.escape(word) + r""")(\s|[,.)!?;:'"s])""",
                r"\1<link>\2</link>\3",
                content,
                count=1,
                flags=re.IGNORECASE,
            )
        return content

    @property
    def description(self):
        """
        Provides the description of the getContent tool

        Returns:
            dict: The description of the tool for the API
        """
        raise NotImplementedError("You need to implement the description property for the GetContentTool")


class MovePageTool:
    """
    The MovePageTool allows the agent to navigate to a different Wikipedia page using a valid link
    found in the current page content.

    This is an example of a tool that modifies the task state dynamically based on inputs from the
    agent.
    """

    name = "move_page"

    @staticmethod
    def execute(new_page: str, task: WikiGame) -> str:
        """
        Changes your current page to a specified new page which is accessible via a link from the
        current page. You can only call this function once at a time, as it will take you to a
        different page.

        Args:
            task (WikiGame): The current task object.
            new_page (str): The title of the new page to move to.

        Returns:
            str: A message indicating the result of the move
        """
        raise NotImplementedError("You need to implement the execute method for the MovePageTool")

    @property
    def description(self):
        """
        Provides the description of the move_page tool

        Returns:
            dict: The description of the move_page tool for the API
        """
        raise NotImplementedError("You need to implement the description property for the MovePageTool")


tests.test_get_content_tool(GetContentTool)
tests.test_move_page_tool(MovePageTool)
GetContentTool_inst = GetContentTool()
MovePageTool_inst = MovePageTool()
wiki_game_tools = [GetContentTool_inst, MovePageTool_inst]

class WikiAgent(SimpleAgent):
    """
    WikiAgent is an LLM-based agent designed to navigate the Wikipedia game by integrating decision-
    making & tool use. It inherits from SimpleAgent, but should be modified to work effectively in
    the Wikipedia game context.

    Attributes:
        model (str): The model used for generating responses (inherited)
        tools (list[Any]): List of tools (inherited)
        client (OpenAI): OpenAI client for API calls (inherited)
        task (WikiGame): The current task being executed
        chat_history (list[dict]): History of interactions (inherited)

    Inherited Methods:
        get_response(use_tool: bool = True) -> ChatCompletionMessage:
            Get response from the model (inherited)

        execute_tool_calls(message: ChatCompletionMessage) -> list[str]:
            Execute tool calls from the model's response (inherited)


    Methods:
        handle_refusal(response: ChatCompletionMessage):
            Handle refusals from the model response (implemented below)

        reset_history():
            Empty self.chat_history of the agent (implemented below)

        system_instruction() -> dict:
            Generate instructions for the game. Formatted as a system prompt (to be implemented)

        on_page_instruction() -> dict:
            Tell the agent what page they are on, and what page they are trying to get to. Formatted
            as a user prompt (to be implemented)

        next_step_instruction() -> dict:
            Ask the agent "What's your next step?" after making a tool call. Formatted as a user
            prompt (to be implemented)

        handle_tool_calls(response: ChatCompletionMessage):
            Handle tool calls from the model response (to be implemented)

        start():
            Put the starting instructions in agent.chat_history when the agent starts a new page or
            starts the game (to be implemented)

        run(with_tool: bool = True) -> bool:
            Run one loop of the Wikipedia agent (to be implemented)

    """

    def __init__(
        self,
        task: WikiGame,
        tools: list[Any],
        model="gpt-4o-mini",
        chat_history: list[dict] = None,
        verbose: bool = True,
        temperature=1,
    ):
        super().__init__(model=model, tools=tools, task=task, temperature=temperature)

        self.chat_history = chat_history if chat_history else []
        self.full_chat_history = chat_history if chat_history else []
        # ^ All messages that have been sent in the chat history.
        # We have to erase each time a new page is reached for context window reasons.
        self.verbose = verbose
        self.start()

    def handle_refusal(self, response: ChatCompletionMessage):
        """
        Handles refusals in the wikipedia game context:

        Args:
            response (ChatCompletionMessage): The response from the model
        """
        self.update_history({"role": "assistant", "content": response.refusal})
        if self.verbose:
            print(f"\nMODEL REFUSAL: {response.refusal}")

    def reset_history(self):
        """
        Empty self.chat_history of the agent.
        """
        self.chat_history = []

    # ========================= Prompting (to implement) =========================
    @property
    def system_instruction(self) -> dict:
        """
        Generate the starting instructions for the game, formatted as a system prompt.

        Returns:
            dict: The starting instructions.
        """
        raise NotImplementedError("You need to implement the system_instruction property")

    @property
    def on_page_instruction(self) -> dict:
        """
        Tell the agent what page they are on and give a summary of the page, formatted as a user prompt.

        Returns:
            dict: The instructions for the current page.
        """
        raise NotImplementedError("You need to implement the on_page_instruction property")

    @property
    def next_step_instruction(self) -> dict:
        """
        Ask the agent "What's your next step?" after making a tool call, formatted as a user prompt.

        Returns:
            dict: The instructions for the next step.
        """
        raise NotImplementedError("You need to implement the next_step_instruction property")

    # ========================= Chat History Management (to implement) =========================

    def update_history(
        self,
        message: dict[str, str]
        | ChatCompletionMessage
        | list[dict[str, str] | ChatCompletionMessage],
    ):
        """
        Update self.chat_history and self.full_chat_history with a message or list of messages.

        Args:
            message: The message(s) to add to the chat history
        """
        raise NotImplementedError("You need to implement the update_history method")

    def start(self):
        """
        A function to put the starting instructions in agent.chat_history when the agent starts a
        new page or starts the game.
        """
        raise NotImplementedError("You need to implement the start method")

    # ========================= Task Execution (to implement) =========================

    def handle_tool_calls(self, response: ChatCompletionMessage):
        """
        Handles tool_calls in the wikipedia game context:
            - Adds the model response to the chat_history
            - Executes the tool calls using execute_tool_calls
            - Appends the tool responses to the chat_history
            - If the agent has moved to a new page:
                - Reset the chat_history, and call start()
            - Otherwise
                - Get the next_step_message instruction from the task and append it to chat_history

        Args:
            response (ChatCompletionMessage): The response from the model
        """
        raise NotImplementedError("You need to implement the handle_tool_calls method")

    def run(self):
        """
        This function runs the agent in the wikipedia game context. It:
            - Gets the current task instruction
            - Gets the response from the model
            - Handles the response in the cases:
                - tool calls (using handle_tool_calls)
                - refusals (using handle_refusal)
                - no tool calls (using update_history)
        """
        raise NotImplementedError("You need to implement the run method")


tests.test_wiki_agent(WikiAgent)

def agent_loop(agent, num_loops=10):
    """
    Run the agent loop for a given number of loops

    Args:
        agent (WikiAgent): The agent to run
        game (WikiGame): The game to play
        num_loops (int): The number of loops to run
    """
    raise NotImplementedError("You need to implement the agent_loop function")

game_1 = WikiGame("Elizabeth I", "United States")
agent = WikiAgent(task=game_1, tools=wiki_game_tools, model="gpt-4o-mini", temperature=0)
agent_loop(agent, 30)

game_2 = WikiGame("County Seat", "Saint Pierre and Miquelon")
agent = WikiAgent(task=game_2, tools=wiki_game_tools, model="gpt-4o-mini", temperature=0)
agent_loop(agent, 30)

for message in agent.full_chat_history:
    try:
        if message["role"] == "tool" and message["name"] == "get_content":
            print(f"{message['role']}:\n {message['content'][:200]} ...")
        else:
            print(f"{message['role']}:\n {message['content']}")
    except:
        print(f"{message.role}: {message.content}")