run-llama--llama_index/docs/examples/agent/agent_builder.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "91c0b9fd-213a-4da8-b84b-c766b424716c",
   "metadata": {},
   "source": [
    "# GPT Builder Demo\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/agent/agent_builder.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "Inspired by GPTs interface, presented at OpenAI Dev Day 2023. Construct an agent with natural language.\n",
    "\n",
    "Here you can build your own agent...with another agent!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e112b8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-embeddings-openai\n",
    "%pip install llama-index-llms-openai\n",
    "%pip install llama-index-readers-file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93ff34e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b6d7505-d582-465e-b86c-eaf2cf8c28f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core import Settings\n",
    "\n",
    "llm = OpenAI(model=\"gpt-4o\")\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = OpenAIEmbedding(model=\"text-embedding-3-small\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "839cb488-912e-4a34-88a4-98c751798fcc",
   "metadata": {},
   "source": [
    "## Define Candidate Tools\n",
    "\n",
    "We also define a tool retriever to retrieve candidate tools.\n",
    "\n",
    "In this setting we define tools as different Wikipedia pages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc91fd57-a681-4c18-991c-6f011d180dea",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50797099-bff6-40f8-b245-62a80b07e7db",
   "metadata": {},
   "outputs": [],
   "source": [
    "wiki_titles = [\"Toronto\", \"Seattle\", \"Chicago\", \"Boston\", \"Houston\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5aba2fc-9bde-44e7-8a69-8a25ffa8de73",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import requests\n",
    "\n",
    "for title in wiki_titles:\n",
    "    response = requests.get(\n",
    "        \"https://en.wikipedia.org/w/api.php\",\n",
    "        params={\n",
    "            \"action\": \"query\",\n",
    "            \"format\": \"json\",\n",
    "            \"titles\": title,\n",
    "            \"prop\": \"extracts\",\n",
    "            # 'exintro': True,\n",
    "            \"explaintext\": True,\n",
    "        },\n",
    "    ).json()\n",
    "    page = next(iter(response[\"query\"][\"pages\"].values()))\n",
    "    wiki_text = page[\"extract\"]\n",
    "\n",
    "    data_path = Path(\"data\")\n",
    "    if not data_path.exists():\n",
    "        Path.mkdir(data_path)\n",
    "\n",
    "    with open(data_path / f\"{title}.txt\", \"w\") as fp:\n",
    "        fp.write(wiki_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a034ebf2-4a31-488b-bfbf-777dbc768426",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load all wiki documents\n",
    "city_docs = {}\n",
    "for wiki_title in wiki_titles:\n",
    "    city_docs[wiki_title] = SimpleDirectoryReader(\n",
    "        input_files=[f\"data/{wiki_title}.txt\"]\n",
    "    ).load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96ccd68b-45fe-43aa-a209-b2fd5d2aa75d",
   "metadata": {},
   "source": [
    "### Build Query Tool for Each Document"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc9e8634-18e3-4762-8e9e-792a5ce8e934",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.core.tools import QueryEngineTool\n",
    "from llama_index.core import VectorStoreIndex\n",
    "\n",
    "# Build tool dictionary\n",
    "tool_dict = {}\n",
    "\n",
    "for wiki_title in wiki_titles:\n",
    "    # build vector index\n",
    "    vector_index = VectorStoreIndex.from_documents(\n",
    "        city_docs[wiki_title],\n",
    "    )\n",
    "    # define query engines\n",
    "    vector_query_engine = vector_index.as_query_engine(llm=llm)\n",
    "\n",
    "    # define tools\n",
    "    vector_tool = QueryEngineTool.from_defaults(\n",
    "        query_engine=vector_query_engine,\n",
    "        name=wiki_title,\n",
    "        description=(\"Useful for questions related to\" f\" {wiki_title}\"),\n",
    "    )\n",
    "    tool_dict[wiki_title] = vector_tool"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d2a3aeb-11ee-4dd2-aabb-d148213e234a",
   "metadata": {},
   "source": [
    "### Define Tool Retriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70d41c03-4110-4990-990b-7a3d706c0c84",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define an \"object\" index and retriever over these tools\n",
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.core.objects import ObjectIndex\n",
    "\n",
    "tool_index = ObjectIndex.from_objects(\n",
    "    list(tool_dict.values()),\n",
    "    index_cls=VectorStoreIndex,\n",
    ")\n",
    "tool_retriever = tool_index.as_retriever(similarity_top_k=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1642f27f-457a-4cd7-b543-9f81a04a42da",
   "metadata": {},
   "source": [
    "### Load Data\n",
    "\n",
    "Here we load wikipedia pages from different cities."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "063cc9a7-d74c-4e08-9110-a63e11841d7f",
   "metadata": {},
   "source": [
    "## Define Meta-Tools for GPT Builder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d76c29b-4b24-4d47-bd1e-027af9427f6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import FunctionAgent\n",
    "from llama_index.core.llms import ChatMessage\n",
    "from llama_index.core import ChatPromptTemplate\n",
    "from typing import List\n",
    "\n",
    "GEN_SYS_PROMPT_STR = \"\"\"\\\n",
    "Task information is given below. \n",
    "\n",
    "Given the task, please generate a system prompt for an OpenAI-powered bot to solve this task: \n",
    "{task} \\\n",
    "\"\"\"\n",
    "\n",
    "gen_sys_prompt_messages = [\n",
    "    ChatMessage(\n",
    "        role=\"system\",\n",
    "        content=\"You are helping to build a system prompt for another bot.\",\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=GEN_SYS_PROMPT_STR),\n",
    "]\n",
    "\n",
    "GEN_SYS_PROMPT_TMPL = ChatPromptTemplate(gen_sys_prompt_messages)\n",
    "\n",
    "\n",
    "agent_cache = {}\n",
    "\n",
    "\n",
    "async def create_system_prompt(task: str):\n",
    "    \"\"\"Create system prompt for another agent given an input task.\"\"\"\n",
    "    llm = OpenAI(llm=\"gpt-4\")\n",
    "    fmt_messages = GEN_SYS_PROMPT_TMPL.format_messages(task=task)\n",
    "    response = await llm.achat(fmt_messages)\n",
    "    return response.message.content\n",
    "\n",
    "\n",
    "async def get_tools(task: str):\n",
    "    \"\"\"Get the set of relevant tools to use given an input task.\"\"\"\n",
    "    subset_tools = await tool_retriever.aretrieve(task)\n",
    "    return [t.metadata.name for t in subset_tools]\n",
    "\n",
    "\n",
    "def create_agent(system_prompt: str, tool_names: List[str]):\n",
    "    \"\"\"Create an agent given a system prompt and an input set of tools.\"\"\"\n",
    "    llm = OpenAI(model=\"gpt-4o\")\n",
    "    try:\n",
    "        # get the list of tools\n",
    "        input_tools = [tool_dict[tn] for tn in tool_names]\n",
    "\n",
    "        agent = FunctionAgent(\n",
    "            tools=input_tools, llm=llm, system_prompt=system_prompt\n",
    "        )\n",
    "        agent_cache[\"agent\"] = agent\n",
    "        return_msg = \"Agent created successfully.\"\n",
    "    except Exception as e:\n",
    "        return_msg = f\"An error occurred when building an agent. Here is the error: {repr(e)}\"\n",
    "    return return_msg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02ebe043-73f1-4f09-88e1-011ceb4ed05d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.tools import FunctionTool\n",
    "\n",
    "system_prompt_tool = FunctionTool.from_defaults(fn=create_system_prompt)\n",
    "get_tools_tool = FunctionTool.from_defaults(fn=get_tools)\n",
    "create_agent_tool = FunctionTool.from_defaults(fn=create_agent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c42842b3-438e-484e-9598-ad7dc7e2de09",
   "metadata": {},
   "outputs": [],
   "source": [
    "GPT_BUILDER_SYS_STR = \"\"\"\\\n",
    "You are helping to construct an agent given a user-specified task. You should generally use the tools in this order to build the agent.\n",
    "\n",
    "1) Create system prompt tool: to create the system prompt for the agent.\n",
    "2) Get tools tool: to fetch the candidate set of tools to use.\n",
    "3) Create agent tool: to create the final agent.\n",
    "\"\"\"\n",
    "\n",
    "prefix_msgs = [ChatMessage(role=\"system\", content=GPT_BUILDER_SYS_STR)]\n",
    "\n",
    "\n",
    "builder_agent = FunctionAgent(\n",
    "    tools=[system_prompt_tool, get_tools_tool, create_agent_tool],\n",
    "    prefix_messages=prefix_msgs,\n",
    "    llm=OpenAI(model=\"gpt-4o\"),\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e135c1a-fcd5-40bc-b92a-5c1ad6ad9a50",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Called tool create_system_prompt with input {'task': 'Tell me about Toronto'}\n",
      "Got output: \"Generate a brief summary about Toronto, including its history, culture, landmarks, and notable features.\"\n",
      "Called tool get_tools with input {'task': 'Tell me about Toronto'}\n",
      "Got output: ['Toronto']\n",
      "Called tool create_agent with input {'system_prompt': 'Generate a brief summary about Toronto, including its history, culture, landmarks, and notable features.', 'tool_names': ['Toronto']}\n",
      "Got output: Agent created successfully.\n",
      "Result: I have created an agent that can provide information about Toronto, including its history, culture, landmarks, and notable features. You can now ask the agent any questions you have about Toronto!\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.agent.workflow import ToolCallResult\n",
    "\n",
    "handler = builder_agent.run(\"Build an agent that can tell me about Toronto.\")\n",
    "async for event in handler.stream_events():\n",
    "    if isinstance(event, ToolCallResult):\n",
    "        print(\n",
    "            f\"Called tool {event.tool_name} with input {event.tool_kwargs}\\nGot output: {event.tool_output}\"\n",
    "        )\n",
    "\n",
    "result = await handler\n",
    "print(f\"Result: {result}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee65b244-a6f0-447f-a88d-b7cbdfe8a74a",
   "metadata": {},
   "outputs": [],
   "source": [
    "city_agent = agent_cache[\"agent\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c66eea2-21a9-4e3d-a3a2-f4219476903e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Toronto is home to a diverse array of parks and public spaces, offering both urban and natural environments. Key downtown parks include Allan Gardens, Christie Pits, and Trinity Bellwoods Park. For waterfront views, Tommy Thompson Park and the Toronto Islands are popular destinations. In the city's outer areas, large parks like High Park, Humber Bay Park, and Morningside Park provide expansive green spaces. Additionally, parts of Rouge National Urban Park, the largest urban park in North America, are located within Toronto. The city also features notable squares such as Nathan Phillips Square, Yonge–Dundas Square, and Harbourfront Square. Approximately 12.5% of Toronto's land is dedicated to parkland, offering facilities for various activities, including winter sports like ice skating and skiing.\n"
     ]
    }
   ],
   "source": [
    "response = await city_agent.run(\"Tell me about the parks in Toronto\")\n",
    "print(str(response))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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
}