llama_hub/llama_packs/agents/llm_compiler/llm_compiler.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f2cf3ee9-d4dd-44d1-a529-6e92d9f8e023",
   "metadata": {},
   "source": [
    "# LLM Compiler Agent Cookbook\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/run-llama/llama-hub/blob/main/llama_hub/llama_packs/agents/llm_compiler/llm_compiler.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "**NOTE**: Full credits to the [source repo for LLMCompiler](https://github.com/SqueezeAILab/LLMCompiler). A lot of our implementation was lifted from this repo (and adapted with LlamaIndex modules).\n",
    "\n",
    "In this cookbook, we show how to use our LLMCompiler agent implementation for various settings. This includes using some simple function tools to do math, but also to answer multi-part queries for more advanced RAG use cases over multiple documents.\n",
    "\n",
    "We see that the LLMCompilerAgent is capable of parallel function calling, giving results much more quickly than sequential execution through ReAct.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6f6b9616-fbd3-4686-9ce3-2893bdfad2e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🌍 To view the Phoenix app in your browser, visit http://127.0.0.1:6006/\n",
      "📺 To view the Phoenix app in a notebook, run `px.active_session().view()`\n",
      "📖 For more information on how to use Phoenix, check out https://docs.arize.com/phoenix\n"
     ]
    }
   ],
   "source": [
    "# Phoenix can display in real time the traces automatically\n",
    "# collected from your LlamaIndex application.\n",
    "import phoenix as px\n",
    "\n",
    "# Look for a URL in the output to open the App in a browser.\n",
    "px.launch_app()\n",
    "# The App is initially empty, but as you proceed with the steps below,\n",
    "# traces will appear automatically as your LlamaIndex application runs.\n",
    "\n",
    "import llama_index\n",
    "\n",
    "llama_index.set_global_handler(\"arize_phoenix\")\n",
    "\n",
    "# Run all of your LlamaIndex applications as usual and traces\n",
    "# will be collected and displayed in Phoenix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "338345d0-0b03-44f0-b812-d98b530d047b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93038a05-9e15-46a8-9b17-56fd1bf5f4c6",
   "metadata": {},
   "source": [
    "### Download Llama Pack\n",
    "\n",
    "Here we download the Llama Pack.\n",
    "\n",
    "**NOTE**: This only works if `skip_load=True`, because we are loading an entire directory of files instead of just a single file.\n",
    "\n",
    "Instead of directly using the pack though, we will show how to directly import some of the underlying modules to build your custom agents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "653fd259-499a-4705-aff6-16890e3f5373",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Option: if developing with the llama_hub package\n",
    "# from llama_hub.llama_packs.agents.llm_compiler.step import LLMCompilerAgentWorker\n",
    "\n",
    "# Option: download_llama_pack\n",
    "from llama_index.llama_pack import download_llama_pack\n",
    "\n",
    "download_llama_pack(\n",
    "    \"LLMCompilerAgentPack\",\n",
    "    \"./agent_pack\",\n",
    "    skip_load=True,\n",
    "    # leave the below line commented out if using the notebook on main\n",
    "    # llama_hub_url=\"https://raw.githubusercontent.com/run-llama/llama-hub/jerry/add_llm_compiler_pack/llama_hub\"\n",
    ")\n",
    "from agent_pack.step import LLMCompilerAgentWorker"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b82925e9-0c6b-4e1a-82f7-5d238cfd4bc6",
   "metadata": {},
   "source": [
    "## Test LLMCompiler Agent with Simple Functions\n",
    "\n",
    "Here we test the LLMCompilerAgent with simple math functions (add, multiply) to illustrate how it works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e6fdf3ad-0ded-4165-bdc6-870f91dc723d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from typing import Sequence, List\n",
    "\n",
    "from llama_index.llms import OpenAI, ChatMessage\n",
    "from llama_index.tools import BaseTool, FunctionTool\n",
    "\n",
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc38b021-97c8-4d00-8781-91ce081547c0",
   "metadata": {},
   "source": [
    "### Define Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c13da398-bd9e-4f2e-82d2-2dbb62a3ac52",
   "metadata": {},
   "outputs": [],
   "source": [
    "def multiply(a: int, b: int) -> int:\n",
    "    \"\"\"Multiple two integers and returns the result integer\"\"\"\n",
    "    return a * b\n",
    "\n",
    "\n",
    "multiply_tool = FunctionTool.from_defaults(fn=multiply)\n",
    "\n",
    "\n",
    "def add(a: int, b: int) -> int:\n",
    "    \"\"\"Add two integers and returns the result integer\"\"\"\n",
    "    return a + b\n",
    "\n",
    "\n",
    "add_tool = FunctionTool.from_defaults(fn=add)\n",
    "\n",
    "tools = [multiply_tool, add_tool]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9a77611c-5df1-488e-8856-6bc5cdb4249e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"{'title': 'multiply', 'type': 'object', 'properties': {'a': {'title': 'A', 'type': 'integer'}, 'b': {'title': 'B', 'type': 'integer'}}, 'required': ['a', 'b']}\""
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "multiply_tool.metadata.fn_schema_str"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9ce60c7-9625-42a2-87d6-a8fe7854a4ca",
   "metadata": {},
   "source": [
    "### Setup LLMCompiler Agent\n",
    "\n",
    "We import the `LLMCompilerAgentWorker` and combine it with the `AgentRunner`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "faf87b3e-a027-4da2-a402-a704cd03a728",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.agent import AgentRunner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "35b3fa8a-647f-4268-aac2-4816fcfbdecc",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(model=\"gpt-4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "413a6043-4824-4a12-9de7-fbfe15c4e541",
   "metadata": {},
   "outputs": [],
   "source": [
    "callback_manager = llm.callback_manager"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1c14e9bb-b401-414e-916e-4c575b26b5e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_worker = LLMCompilerAgentWorker.from_tools(\n",
    "    tools, llm=llm, verbose=True, callback_manager=callback_manager\n",
    ")\n",
    "agent = AgentRunner(agent_worker, callback_manager=callback_manager)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7f2da7b-e611-4441-ab0e-22b8748ec658",
   "metadata": {},
   "source": [
    "### Test out some Queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a0980a07-d2f5-434a-b9e9-80deabc0d1a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Running step 7fbd5304-8f67-46f6-882d-9c15ace75d80 for task dbbb4991-5347-4302-a358-6fe5d0705bdc.\n",
      "> Step count: 0\n",
      "\u001b[1;3;38;5;200m> Plan: 1. multiply(121, 3)\n",
      "2. add($1, 42)\n",
      "3. join()<END_OF_PLAN>\n",
      "\u001b[0m\u001b[1;3;34mRan task: multiply. Observation: 363\n",
      "\u001b[0m\u001b[1;3;34mRan task: add. Observation: 405\n",
      "\u001b[0m\u001b[1;3;34mRan task: join. Observation: None\n",
      "\u001b[0m\u001b[1;3;38;5;200m> Thought: The result of the operation is 405.\n",
      "\u001b[0m\u001b[1;3;38;5;200m> Answer: 405\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "response = agent.chat(\"What is (121 * 3) + 42?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "6ebe5e1f-4929-4f47-9af2-1da8faebe638",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AgentChatResponse(response='405', sources=[], source_nodes=[])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a3dfd5c4-633c-4729-a337-a4fcf13d8a76",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[ChatMessage(role=<MessageRole.USER: 'user'>, content='What is (121 * 3) + 42?', additional_kwargs={}),\n",
       " ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>, content='405', additional_kwargs={})]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent.memory.get_all()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1e836cb-f481-4033-afca-211a52ef0665",
   "metadata": {},
   "source": [
    "## Try out LLMCompiler for RAG\n",
    "\n",
    "Now let's try out the LLMCompiler for RAG use cases. Specifically, we load a dataset of Wikipedia articles about various cities and ask questions over them."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "975c5d36-6d80-4004-a565-15351b7ac702",
   "metadata": {},
   "source": [
    "### Setup Data\n",
    "\n",
    "We use our `WikipediaReader` to load data for various cities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a50dccdc-491e-4ff2-8657-5dfea6fab284",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.readers import WikipediaReader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2647c3be-1637-4b88-ae76-bb513f7e9682",
   "metadata": {},
   "outputs": [],
   "source": [
    "wiki_titles = [\"Toronto\", \"Seattle\", \"Chicago\", \"Boston\", \"Miami\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8eeb7add-aa6c-4f47-8d8c-6a6d222adc55",
   "metadata": {},
   "outputs": [],
   "source": [
    "city_docs = {}\n",
    "reader = WikipediaReader()\n",
    "for wiki_title in wiki_titles:\n",
    "    docs = reader.load_data(pages=[wiki_title])\n",
    "    city_docs[wiki_title] = docs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4af02a14-0d58-4570-93c5-78e1c87c8c72",
   "metadata": {},
   "source": [
    "### Setup LLM + Service Context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6425ad0e-2f8e-4101-8730-387f4aab8bfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import ServiceContext\n",
    "from llama_index.llms import OpenAI\n",
    "from llama_index.callbacks import CallbackManager\n",
    "\n",
    "llm = OpenAI(temperature=0, model=\"gpt-4\")\n",
    "service_context = ServiceContext.from_defaults(llm=llm)\n",
    "callback_manager = CallbackManager([])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3a52155-1790-44e6-b6be-14964b0b25f1",
   "metadata": {},
   "source": [
    "### Define Toolset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "91324b15-9f98-4f52-a832-2eadc76b73c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import load_index_from_storage, StorageContext\n",
    "from llama_index.node_parser import SentenceSplitter\n",
    "from llama_index.tools import QueryEngineTool, ToolMetadata\n",
    "from llama_index import VectorStoreIndex\n",
    "import os\n",
    "\n",
    "node_parser = SentenceSplitter()\n",
    "\n",
    "# Build agents dictionary\n",
    "query_engine_tools = []\n",
    "\n",
    "for idx, wiki_title in enumerate(wiki_titles):\n",
    "    nodes = node_parser.get_nodes_from_documents(city_docs[wiki_title])\n",
    "\n",
    "    if not os.path.exists(f\"./data/{wiki_title}\"):\n",
    "        # build vector index\n",
    "        vector_index = VectorStoreIndex(\n",
    "            nodes, service_context=service_context, callback_manager=callback_manager\n",
    "        )\n",
    "        vector_index.storage_context.persist(persist_dir=f\"./data/{wiki_title}\")\n",
    "    else:\n",
    "        vector_index = load_index_from_storage(\n",
    "            StorageContext.from_defaults(persist_dir=f\"./data/{wiki_title}\"),\n",
    "            service_context=service_context,\n",
    "            callback_manager=callback_manager,\n",
    "        )\n",
    "    # define query engines\n",
    "    vector_query_engine = vector_index.as_query_engine()\n",
    "\n",
    "    # define tools\n",
    "    query_engine_tools.append(\n",
    "        QueryEngineTool(\n",
    "            query_engine=vector_query_engine,\n",
    "            metadata=ToolMetadata(\n",
    "                name=f\"vector_tool_{wiki_title}\",\n",
    "                description=(\n",
    "                    \"Useful for questions related to specific aspects of\"\n",
    "                    f\" {wiki_title} (e.g. the history, arts and culture,\"\n",
    "                    \" sports, demographics, or more).\"\n",
    "                ),\n",
    "            ),\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29750eb8-b117-4777-b228-32e8becf662c",
   "metadata": {},
   "source": [
    "### Setup LLMCompilerAgent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "dde0c37b-5660-43a0-a22f-78d4521204a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.agent import AgentRunner\n",
    "from llama_index.llms import OpenAI\n",
    "\n",
    "llm = OpenAI(model=\"gpt-4\")\n",
    "agent_worker = LLMCompilerAgentWorker.from_tools(\n",
    "    query_engine_tools,\n",
    "    llm=llm,\n",
    "    verbose=True,\n",
    "    callback_manager=callback_manager,\n",
    ")\n",
    "agent = AgentRunner(agent_worker, callback_manager=callback_manager)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf2e41c1-3c68-4cc0-85f1-94eb21746bac",
   "metadata": {},
   "source": [
    "### Test out Queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "54484552-e53f-4548-8719-8c188179b596",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Running step 3aef03b1-9f79-4581-8859-680e0c589daa for task f76c6082-2f96-4cb9-8f28-b3d377f40874.\n",
      "> Step count: 0\n",
      "\u001b[1;3;38;5;200m> Plan: 1. vector_tool_Miami('demographics')\n",
      "2. vector_tool_Chicago('demographics')\n",
      "Thought: I can answer the question now.\n",
      "3. join()<END_OF_PLAN>\n",
      "\u001b[0m\u001b[1;3;34mRan task: vector_tool_Miami. Observation: Miami is the largest city in South Florida and the second-largest in Florida, with a population of 442,241 as of the 2020 census. The city is part of the Miami metropolitan area, which has over 6 million residents. Miami's population is diverse, with a significant Hispanic majority. In 2010, 34.4% of city residents were of Cuban origin, and 15.8% had a Central American background. The Non-Hispanic White population of Miami has been growing, making up 14.0% of the population in 2020. The Non-Hispanic Black population has been declining, making up 11.7% of the population in 2020. Other ethnic groups include those of Asian ancestry, who accounted for 1.0% of Miami's population in 2010. Christianity is the most-practiced religion in Miami, followed by Judaism, Islam, Buddhism, Hinduism, and other religions. As of 2022, there were 3,440 homeless people in Miami-Dade County, with 591 unsheltered homeless people on the streets of Miami.\n",
      "\u001b[0m\u001b[1;3;34mRan task: vector_tool_Chicago. Observation: Chicago has a diverse demographic makeup. The city experienced rapid growth in its early years, with the population reaching over 1 million by 1890. By the end of the 19th century, Chicago was the fifth-largest city in the world. The city saw waves of immigrants from various parts of the world, including Ireland, Southern, Central and Eastern Europe, and Africa. The city's black population doubled between 1910 and 1920, and again between 1920 and 1930. As of July 2019, the largest racial or ethnic group in Chicago is non-Hispanic White at 32.8% of the population, followed by Blacks at 30.1% and the Hispanic population at 29.0%. The city also has the third-largest LGBT population in the United States. The city's population declined in the latter half of the 20th century, but has seen a rise in the 2000 and 2020 censuses.\n",
      "\u001b[0m\u001b[1;3;34mRan task: join. Observation: None\n",
      "\u001b[0m\u001b[1;3;38;5;200m> Thought: Miami and Chicago both have diverse populations, but their compositions are different. Miami has a significant Hispanic majority, with a large portion of residents of Cuban and Central American origin. The non-Hispanic White population is growing, while the non-Hispanic Black population is declining. In contrast, Chicago's largest racial or ethnic group is non-Hispanic White, followed closely by Blacks and Hispanics. Chicago also has the third-largest LGBT population in the United States. Both cities have seen population changes over time, with Miami's population growing and Chicago's initially declining but recently increasing.\n",
      "\u001b[0m\u001b[1;3;38;5;200m> Answer: Miami and Chicago both have diverse populations, but their compositions are different. Miami has a significant Hispanic majority, with a large portion of residents of Cuban and Central American origin. The non-Hispanic White population is growing, while the non-Hispanic Black population is declining. In contrast, Chicago's largest racial or ethnic group is non-Hispanic White, followed closely by Blacks and Hispanics. Chicago also has the third-largest LGBT population in the United States. Both cities have seen population changes over time, with Miami's population growing and Chicago's initially declining but recently increasing.\n",
      "\u001b[0mMiami and Chicago both have diverse populations, but their compositions are different. Miami has a significant Hispanic majority, with a large portion of residents of Cuban and Central American origin. The non-Hispanic White population is growing, while the non-Hispanic Black population is declining. In contrast, Chicago's largest racial or ethnic group is non-Hispanic White, followed closely by Blacks and Hispanics. Chicago also has the third-largest LGBT population in the United States. Both cities have seen population changes over time, with Miami's population growing and Chicago's initially declining but recently increasing.\n"
     ]
    }
   ],
   "source": [
    "response = agent.chat(\n",
    "    \"Tell me about the demographics of Miami, and compare that with the demographics of Chicago?\"\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "e1e0c219-ce9b-435c-86e1-85f5aaf3afe8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Running step f8fdf4cb-9dde-4aba-996d-edbcee53c4c2 for task 20df27d7-cc27-4311-bb57-b1a6f4ad5799.\n",
      "> Step count: 0\n",
      "\u001b[1;3;38;5;200m> Plan: 1. vector_tool_Chicago(\"climate during wintertime\")\n",
      "2. vector_tool_Seattle(\"climate during wintertime\")\n",
      "3. join()<END_OF_PLAN>\n",
      "\u001b[0m\u001b[1;3;34mRan task: vector_tool_Seattle. Observation: During wintertime, Seattle experiences cool, wet conditions. Extreme cold temperatures, below about 15 °F or -9 °C, are rare due to the moderating influence of the adjacent Puget Sound, the greater Pacific Ocean, and Lake Washington. The city is often cloudy due to frequent storms and lows moving in from the Pacific Ocean, and it has many \"rain days\". However, the rainfall is often a light drizzle.\n",
      "\u001b[0m\u001b[1;3;34mRan task: vector_tool_Chicago. Observation: During wintertime, the city experiences relatively cold and snowy conditions. Blizzards can occur, as they did in winter 2011. The normal winter high from December through March is about 36 °F (2 °C). January and February are the coldest months. A polar vortex in January 2019 nearly broke the city's cold record of −27 °F (−33 °C), which was set on January 20, 1985. Measurable snowfall can continue through the first or second week of April. The city's proximity to Lake Michigan tends to keep the lakefront somewhat cooler in summer and less brutally cold in winter than inland parts of the city and suburbs away from the lake. Northeast winds from wintertime cyclones departing south of the region sometimes bring the city lake-effect snow.\n",
      "\u001b[0m\u001b[1;3;34mRan task: join. Observation: None\n",
      "\u001b[0m\u001b[1;3;38;5;200m> Thought: Comparing the two climates, Seattle seems to have a milder winter climate than Chicago.\n",
      "\u001b[0m\u001b[1;3;38;5;200m> Answer: Seattle\n",
      "\u001b[0mSeattle\n"
     ]
    }
   ],
   "source": [
    "response = agent.chat(\n",
    "    \"Is the climate of Chicago or Seattle better during the wintertime?\"\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcc9b12c-3c30-46f0-8d04-96fffc0d45f9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v2",
   "language": "python",
   "name": "llama_index_v2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}