docs/agent_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cda38fbe-8df8-4893-aa0d-4fe67acbb646",
   "metadata": {},
   "source": [
    "# Tutorial: Creating Language Agents to interact with Aviary environments\n",
    "\n",
    "---\n",
    "\n",
    "This tutorial will guide you in creating language agents that interact with [Aviary](https://github.com/Future-House/aviary) environments. We'll start by introducing a few basic concepts from the Aviary and [LDP](https://github.com/Future-House/ldp) frameworks, then show how to define agents that can engage with existing environments. First, we’ll build a simple agent, then extend it to create one that follows system prompts for guided interactions.\n",
    "\n",
    "---\n",
    "\n",
    "## **Pre-requisites**\n",
    "\n",
    "To get started, you’ll need to install `Aviary`, `LDP`, and `pydantic`, all available on PyPI. Note that `Aviary` is listed on PyPI as `fhaviary`. Run the command below to install these packages:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f219d662-93a9-4026-80a2-efdecd017e85",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"fhaviary[gsm8k,hotpotqa]\" ldp pydantic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2314d74-c13b-407b-b9b9-0297ea9cb78b",
   "metadata": {},
   "source": [
    "You will need to set the `OPENAI_API_KEY` or the corresponding key for any other API you wish to access."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ccd741e-da9a-492e-8f81-7806093c08ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"your_openai_api_key\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3be864c0-0c1e-4ec8-889c-d777582d7f57",
   "metadata": {},
   "source": [
    "# 1. Background\n",
    "\n",
    "[Aviary](https://github.com/Future-House/aviary) is our framework supporting diverse language environments, where actions are tools available to agents. [LDP](https://github.com/Future-House/ldp) is our framework for creating and training language agents.\n",
    "\n",
    "Below, we briefly define some key classes and concepts from these libraries for context:\n",
    "\n",
    "**From Aviary**  \n",
    "- `Message`: Used by language agents and environments for communication. Messages include attributes like `content` ot `role` (`system`, `user`, `assistant`, `tool` ), matching OpenAI's conventions.\n",
    "- `Environment`: An environment is a stateful system or \"world\" where an agent operates by taking actions. In Aviary, these actions are called tools. The environment presents states that the agent observes (totally or partially), prompting it to use tools to affect outcomes. Each action taken yields a reward and leads to a new state.\n",
    "- `Tool`: Defines an environmental tool that an agent can use to accomplish its task. Different environments offer different tools. Most tools take arguments and multiple tools can be called in parallel.\n",
    "- `ToolRequestMessage`, `ToolResponseMessage`: These are specialized subclasses of `Message` used for tool requests and responses. Typically, a language agent sends a `ToolRequestMessage` to the environment to request the execution of a specific tool, and the environment responds with a `ToolResponseMessage` containing the result. The role of `ToolRequestMessage` is always `assistant` and the role of the `ToolResponseMessage` is always `tool`.\n",
    "\n",
    "**From LDP**   \n",
    "- `Agent`: An entity that interacts with the environment, mapping observations to tool request actions.\n",
    "- `Compute Graph`: A network of operations representing data flow through an agent, storing information and parameters useful for learning. Similar to compute graphs in libraries like TensorFlow or PyTorch, LDP handles this internally, so in-depth knowledge isn’t necessary for this tutorial. However, some methods will have a `@compute_graph` decorator due to this structure.\n",
    "- `Op`: Represents an operation within the agent. LDP includes various operations (Ops), such as API LLM calls, API embedding calls, or PyTorch module handling. These operations form the compute graph.\n",
    "- `OpResult`: the output of an `Op`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57a52b51-b9b9-4430-b649-2bf778b231f4",
   "metadata": {},
   "source": [
    "## Section 2: Defining a Simple Agent\n",
    "\n",
    "We will start by defining the simplest agent possible.\n",
    "\n",
    "### 2.1 Defining an Agent’s State\n",
    "\n",
    "In typical reinforcement learning (RL), an agent receives an observation and generates an action while maintaining an internal state. However, as a design choice here, we make the agents independent of the state and provide them with the current state at each step. This means that the agent will receive an “agent state” containing all relevant information needed to make the next decision at any given step, instead of keeping it internally. This approach enables us, for example, to use a single agent instance to process multiple rollouts in parallel.\n",
    "\n",
    "The state includes all relevant information about a rollout in progress and required by the agent to make the next decision. In the simplest case, this comprises the current list of messages from the start of the episode and the tools available to the agent. We define such a state as follows:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dc2d741-1ff5-4a71-8f49-8f29cd52c4e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Self\n",
    "\n",
    "from aviary.core import Message, Tool, ToolRequestMessage, ToolResponseMessage\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "from ldp.agent import Agent\n",
    "from ldp.graph import OpResult\n",
    "\n",
    "\n",
    "class MySimpleAgentState(BaseModel):\n",
    "    \"\"\"Simple bucket to store available tools and previous messages.\"\"\"\n",
    "\n",
    "    tools: list[Tool] = Field(default_factory=list)\n",
    "    messages: list[Message] = Field(default_factory=list)\n",
    "\n",
    "    def get_next_state(\n",
    "        self,\n",
    "        obs: list[Message] | None = None,\n",
    "    ) -> Self:\n",
    "        \"\"\"\n",
    "        Return the next agent state based on current state and optional messages.\n",
    "\n",
    "        Args:\n",
    "            obs: Optional observation messages to use in creating the next state.\n",
    "\n",
    "        Returns:\n",
    "            The next agent state.\n",
    "        \"\"\"\n",
    "        return type(self)(\n",
    "            tools=self.tools,\n",
    "            messages=self.messages + (obs or []),\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29d5303b-90aa-4297-8f5a-c4c2b522592a",
   "metadata": {},
   "source": [
    "### 2.2 Defining the Agent \n",
    "\n",
    "The `MySimpleAgent` class is a basic agent that utilizes a language model (LLM) to decide which tools to invoke based on observations from its environment. It manages a `MySimpleAgentState` to store relevant information. The agent's core functionality involves querying an LLM to determine the next tool request and updating its state accordingly.\n",
    "\n",
    "All agents in our framework must implement two asynchronous methods:\n",
    "\n",
    "1. `init_state`: Initializes the agent’s internal state by accepting a list of tools the agent can use (provided by the environment).\n",
    "2. `get_asv`: Processes observations to update the agent's state and decide which tool to use next. It returns the tool request (action), the updated agent state, and a value (hence the name `get_asv`). The state value is the the cumulative reward an agent to achieve from the given state . This value indicates the quality of the current state, which may benefit certain learning algorithms. In our case, we set it to 0.0, as we do not have a way to estimate the value.\n",
    "\n",
    "This setup enables `MySimpleAgent` to iteratively assess and respond to new information, invoking the appropriate tool with each decision step.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6556a4ce-0b00-43c9-b7da-daa9049bd5a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any\n",
    "\n",
    "from ldp.graph import LLMCallOp, compute_graph\n",
    "\n",
    "\n",
    "class MySimpleAgent(BaseModel, Agent[MySimpleAgentState]):\n",
    "    \"\"\"Simple agent that can invoke tools with a language model.\"\"\"\n",
    "\n",
    "    llm_model: dict[str, Any] = Field(\n",
    "        default={\n",
    "            \"model\": \"gpt-4o-2024-08-06\",\n",
    "            \"temperature\": 0.1,\n",
    "        },\n",
    "        description=\"Configuration for the LLM object.\",\n",
    "    )\n",
    "\n",
    "    def __init__(self, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "\n",
    "        # Create a Op that the agent will use to call the LLM API\n",
    "        self._llm_call_op = LLMCallOp()\n",
    "\n",
    "    async def init_state(self, tools: list[Tool]) -> MySimpleAgentState:\n",
    "        return MySimpleAgentState(tools=tools)\n",
    "\n",
    "    @compute_graph()\n",
    "    async def get_asv(\n",
    "        self, agent_state: MySimpleAgentState, obs: list[Message]\n",
    "    ) -> tuple[OpResult[ToolRequestMessage], MySimpleAgentState, float]:\n",
    "        # Obtain the next agent state, given the environment observation\n",
    "        next_state = agent_state.get_next_state(obs)\n",
    "\n",
    "        # Call agent language model to ask for the next tool to use\n",
    "        result = await self._llm_call_op(\n",
    "            self.llm_model, msgs=next_state.messages, tools=next_state.tools\n",
    "        )\n",
    "\n",
    "        # Extend the the agent state with the new ToolRequestMessage\n",
    "        next_state.messages = [*next_state.messages, result.value]\n",
    "\n",
    "        # Agent returns an OpResult, the next agent state and the value, which we set to 0.0\n",
    "        return result, next_state, 0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d4b6f87-e9f2-40d5-a92b-84fffe0bce66",
   "metadata": {},
   "source": [
    "### 2.3 Using the Agent\n",
    "\n",
    "Let's define an environment from aviary and ensure our agent can interact with it. The GSM8K environment presents simple mathematical problems to the agent and includes both a calculator tool and a tool for submitting the final answer. The agent receives a reward only after submitting a correct answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9676b2e1-2d05-4c7f-8c9e-f6beb98d25bf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: Weng earns $12 an hour for babysitting. Yesterday, she just did 50 minutes of babysitting. How much did she earn?\n",
      "Step 1 - agent action: calculator({'expr': '12 * (50/60)'}), environment answer: 10 , environment reward 0.0\n",
      "Step 2 - agent action: submit_answer({'answer': '10'}), environment answer: True , environment reward 1.0\n",
      "Finished! \n",
      "\n",
      "Question: Betty is saving money for a new wallet which costs $100. Betty has only half of the money she needs. Her parents decided to give her $15 for that purpose, and her grandparents twice as much as her parents. How much more money does Betty need to buy the wallet?\n",
      "Step 1 - agent action: calculator({'expr': '100 / 2'}), environment answer: 50 , environment reward 0.0\n",
      "Step 2 - agent action: calculator({'expr': '50 + 15'}), environment answer: 65 agent action: calculator({'expr': '15 * 2'}), environment answer: 30 , environment reward 0.0\n",
      "Step 3 - agent action: calculator({'expr': '65 + 30'}), environment answer: 95 , environment reward 0.0\n",
      "Step 4 - agent action: calculator({'expr': '100 - 95'}), environment answer: 5 , environment reward 0.0\n",
      "Step 5 - agent action: submit_answer({'answer': '5'}), environment answer: True , environment reward 1.0\n",
      "Finished! \n",
      "\n",
      "Question: Julie is reading a 120-page book. Yesterday, she was able to read 12 pages and today, she read twice as many pages as yesterday. If she wants to read half of the remaining pages tomorrow, how many pages should she read?\n",
      "Step 1 - agent action: calculator({'expr': '12 * 2'}), environment answer: 24 , environment reward 0.0\n",
      "Step 2 - agent action: calculator({'expr': '120 - (12 + 24)'}), environment answer: 84 , environment reward 0.0\n",
      "Step 3 - agent action: calculator({'expr': '84 / 2'}), environment answer: 42 , environment reward 0.0\n",
      "Step 4 - agent action: submit_answer({'answer': '42'}), environment answer: True , environment reward 1.0\n",
      "Finished! \n",
      "\n"
     ]
    }
   ],
   "source": [
    "from aviary.envs.gsm8k import GSM8kDataset\n",
    "\n",
    "\n",
    "async def main(idx: int = 0):\n",
    "    env = GSM8kDataset(split=\"train\").get_new_env_by_idx(idx)\n",
    "    agent = MySimpleAgent()\n",
    "\n",
    "    # Get initial question, available tools from the environment\n",
    "    obs, tools = await env.reset()\n",
    "    print(f\"Question: {obs[0].content}\")\n",
    "\n",
    "    # Get initial agent state\n",
    "    agent_state = await agent.init_state(tools=tools)\n",
    "\n",
    "    step = 1\n",
    "    done = False\n",
    "    while not done:\n",
    "        action, agent_state, _ = await agent.get_asv(agent_state, obs)\n",
    "        obs, reward, done, _ = await env.step(action.value)\n",
    "        print(\n",
    "            f\"Step {step} - {print_action_obs(action, obs)}, environment reward {reward}\"\n",
    "        )\n",
    "        step += 1\n",
    "\n",
    "    print(\"Finished! \\n\")\n",
    "\n",
    "\n",
    "def print_action_obs(action: ToolRequestMessage, obs: list[ToolResponseMessage]):\n",
    "    tool_calls = action.value.tool_calls\n",
    "    msg = \"\"\n",
    "    for tool_call, tool_answer in zip(tool_calls, obs, strict=True):\n",
    "        tool_name = tool_call.function.name\n",
    "        tool_args = tool_call.function.arguments\n",
    "        msg += f\"agent action: {tool_name}({tool_args}), environment answer: {tool_answer.content} \"\n",
    "    return msg\n",
    "\n",
    "\n",
    "for i in range(3):\n",
    "    await main(i)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4bb99a0-3477-41e0-9af1-6de7a52f6dc6",
   "metadata": {},
   "source": [
    "## Section 3: Specific Prompt Agent\n",
    "\n",
    "Let's extend our agent a bit further. It is common practice to provide system-level guidelines to an agent to help guide its behavior and improve its responses. These textual guidelines provide context or rules that help the agent interpret its environment, make better decisions, and align with specific objectives. We’ll incorporate these guidelines to enhance our agent’s performance and ensure it consistently follows our intended framework.\n",
    "\n",
    "### 3.1 Defining the Agent \n",
    "\n",
    "The changes are minimal with respect to our `MySimpleAgent` class, and we can reuse the `MySimpleAgentState`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "440f6480-27f3-4bf6-9fb7-563962257a93",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any\n",
    "\n",
    "from ldp.graph import compute_graph\n",
    "\n",
    "\n",
    "class MyGuidedAgent(BaseModel, Agent[MySimpleAgentState]):\n",
    "    \"\"\"Simple agent that can invoke tools with a language model.\"\"\"\n",
    "\n",
    "    llm_model: dict[str, Any] = Field(\n",
    "        default={\n",
    "            \"model\": \"gpt-4o-2024-08-06\",\n",
    "            \"temperature\": 0.1,\n",
    "        },\n",
    "        description=\"Configuration for the LLM object.\",\n",
    "    )\n",
    "\n",
    "    guidelines_msg: Message = Field(\n",
    "        default=Message(role=\"system\", content=\"\"),\n",
    "        description=\"Initial guidelines to be shown to the LLM.\",\n",
    "    )\n",
    "\n",
    "    def __init__(self, guidelines: str = \"\", **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "\n",
    "        self.guidelines_msg = Message(role=\"system\", content=\"Guidelines:\" + guidelines)\n",
    "\n",
    "        # Create a Op that calls an LLM\n",
    "        self._llm_call_op = LLMCallOp()\n",
    "\n",
    "    async def init_state(self, tools: list[Tool]) -> MySimpleAgentState:\n",
    "        return MySimpleAgentState(tools=tools)\n",
    "\n",
    "    @compute_graph()\n",
    "    async def get_asv(\n",
    "        self, agent_state: MySimpleAgentState, obs: list[Message]\n",
    "    ) -> tuple[OpResult[ToolRequestMessage], MySimpleAgentState, float]:\n",
    "        # Obtain the next agent state, given the environment observation\n",
    "        next_state = agent_state.get_next_state(obs)\n",
    "\n",
    "        # Call agent language model to ask for the next tool to use\n",
    "        result = await self._llm_call_op(\n",
    "            self.llm_model,\n",
    "            msgs=[\n",
    "                # We prepend the system guidelines here!\n",
    "                self.guidelines_msg,\n",
    "                *next_state.messages,\n",
    "            ],\n",
    "            tools=next_state.tools,\n",
    "        )\n",
    "\n",
    "        # Extend the the agent state with the new ToolRequestMessage\n",
    "        next_state.messages = [*next_state.messages, result.value]\n",
    "\n",
    "        # Agent returns an OpResult, the next agent state and the value, which we set to 0.0\n",
    "        return result, next_state, 0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c32cdcf2-1ec7-459c-98a2-a09638c56028",
   "metadata": {},
   "source": [
    "### 3.2 Using the Agent\n",
    "\n",
    "An simple example where our improved agent can be useful is the HotPotQA environment. The HotPotQA is a question-answering environment. However, the environment only considers the answer correct if it exactly matches the ground truth. Most agents tend to respond with a full sentence that, although containing the correct answer, is marked as incorrect. By providing specific guidelines, we can guide our agent to answer with only the precise information requested. Let's test this by running the environment first without guidelines, and then with them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d02bd30-abef-462e-9ed3-82fd2a2580c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: Which magazine was started first Arthur's Magazine or First for Women?\n",
      "Step 1 - agent action: search({'entity': \"Arthur's Magazine\"}), environment answer: Arthur's Magazine (1844–1846) was an American literary periodical published in Philadelphia in the 19th century. Edited by Timothy Shay Arthur, it featured work by Edgar A. Poe, J.H. Ingraham, Sarah Josepha Hale, Thomas G. Spear, and others.[1][2] In May 1846 it was merged into Godey's Lady's Book.[3]. agent action: search({'entity': 'First for Women'}), environment answer: First for Women is a woman's magazine published by A360media in the US.[1] The magazine was started in 1989 by Bauer Media Group.[2] In 2011 the circulation of the magazine was 1,310,696 copies.[3] In April 2024, the magazine became a weekly.. This women's magazine–related article is a stub. You can help Wikipedia by expanding it.. See tips for writing articles about magazines. Further suggestions might be found on the article's talk page.. , environment reward 0.0\n",
      "Step 2 - agent action: submit_answer({'answer': \"Arthur's Magazine was started first in 1844, while First for Women was started in 1989.\"}), environment answer: Finished. , environment reward 1.0\n",
      "Finished! \n",
      "\n",
      "Question: Which magazine was started first Arthur's Magazine or First for Women?\n",
      "Step 1 - agent action: search({'entity': \"Arthur's Magazine\"}), environment answer: Arthur's Magazine (1844–1846) was an American literary periodical published in Philadelphia in the 19th century. Edited by Timothy Shay Arthur, it featured work by Edgar A. Poe, J.H. Ingraham, Sarah Josepha Hale, Thomas G. Spear, and others.[1][2] In May 1846 it was merged into Godey's Lady's Book.[3]. agent action: search({'entity': 'First for Women'}), environment answer: First for Women is a woman's magazine published by A360media in the US.[1] The magazine was started in 1989 by Bauer Media Group.[2] In 2011 the circulation of the magazine was 1,310,696 copies.[3] In April 2024, the magazine became a weekly.. This women's magazine–related article is a stub. You can help Wikipedia by expanding it.. See tips for writing articles about magazines. Further suggestions might be found on the article's talk page.. , environment reward 0.0\n",
      "Step 2 - agent action: submit_answer({'answer': \"Arthur's Magazine\"}), environment answer: Finished. , environment reward 1.0\n",
      "Finished! \n",
      "\n"
     ]
    }
   ],
   "source": [
    "from aviary.envs.hotpotqa import HotPotQADataset\n",
    "\n",
    "\n",
    "async def main(idx: int = 0, guidelines: str = \"\"):\n",
    "    env = HotPotQADataset(split=\"train\").get_new_env_by_idx(idx)\n",
    "    agent = MyGuidedAgent(guidelines=guidelines)\n",
    "\n",
    "    # Get initial question, available tools from the environment\n",
    "    obs, tools = await env.reset()\n",
    "    print(f\"{obs[0].content}\")\n",
    "\n",
    "    # Get initial agent state\n",
    "    agent_state = await agent.init_state(tools=tools)\n",
    "\n",
    "    step = 1\n",
    "    done = False\n",
    "    while not done:\n",
    "        action, agent_state, _ = await agent.get_asv(agent_state, obs)\n",
    "        obs, reward, done, _ = await env.step(action.value)\n",
    "        print(\n",
    "            f\"Step {step} - {print_action_obs(action, obs)}, environment reward {reward}\"\n",
    "        )\n",
    "        step += 1\n",
    "\n",
    "    print(\"Finished! \\n\")\n",
    "\n",
    "\n",
    "for i in range(1):\n",
    "    await main(i)\n",
    "\n",
    "guidelines = (\n",
    "    \"Answer the questions. \"\n",
    "    \"Return only the specific information asked. \"\n",
    "    \"If asked for a name, only the name; if asked for a year, only the year, etc.\"\n",
    ")\n",
    "\n",
    "for i in range(1):\n",
    "    await main(i, guidelines)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e75015a-8847-4631-a866-d4e4ec208dde",
   "metadata": {},
   "source": [
    "## Section 3: Advanced Agent Definition\n",
    "\n",
    "Many more options are possible. New agents can be created and extended by modifying their `get_asv` method, and if additional information is needed, their state can also be adjusted. Ideas include agents that reflect before taking action, receive feedback from an oracle, or simulate multiple scenarios before making a decision."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}