cookbooks/notebooks/agent/react_agent_with_query_engine.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6b0186a4",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/agent/react_agent_with_query_engine.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b50c4af8-fec3-4396-860a-1322089d76cb",
   "metadata": {},
   "source": [
    "# ReAct Agent with Query Engine (RAG) Tools\n",
    "\n",
    "In this section, we show how to setup an agent powered by the ReAct loop for financial analysis.\n",
    "\n",
    "The agent has access to two \"tools\": one to query the 2021 Lyft 10-K and the other to query the 2021 Uber 10-K.\n",
    "\n",
    "Note that you can plug in any LLM to use as a ReAct agent."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db402a8b-90d6-4e1d-8df6-347c54624f26",
   "metadata": {},
   "source": [
    "## Build Query Engine Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd31aff7",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5d5afce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a2c931d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.core import Settings\n",
    "\n",
    "Settings.llm = OpenAI(model=\"gpt-4o-mini\")\n",
    "Settings.embed_model = OpenAIEmbedding(model=\"text-embedding-3-small\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91618236-54d3-4783-86b7-7b7554efeed1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import StorageContext, load_index_from_storage\n",
    "\n",
    "try:\n",
    "    storage_context = StorageContext.from_defaults(\n",
    "        persist_dir=\"./storage/lyft\"\n",
    "    )\n",
    "    lyft_index = load_index_from_storage(storage_context)\n",
    "\n",
    "    storage_context = StorageContext.from_defaults(\n",
    "        persist_dir=\"./storage/uber\"\n",
    "    )\n",
    "    uber_index = load_index_from_storage(storage_context)\n",
    "\n",
    "    index_loaded = True\n",
    "except:\n",
    "    index_loaded = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a79cbc9",
   "metadata": {},
   "source": [
    "Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36d80144",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p 'data/10k/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf' -O 'data/10k/uber_2021.pdf'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/lyft_2021.pdf' -O 'data/10k/lyft_2021.pdf'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3d0bb8c-16c8-4946-a9d8-59528cf3952a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader, VectorStoreIndex\n",
    "\n",
    "if not index_loaded:\n",
    "    # load data\n",
    "    lyft_docs = SimpleDirectoryReader(\n",
    "        input_files=[\"./data/10k/lyft_2021.pdf\"]\n",
    "    ).load_data()\n",
    "    uber_docs = SimpleDirectoryReader(\n",
    "        input_files=[\"./data/10k/uber_2021.pdf\"]\n",
    "    ).load_data()\n",
    "\n",
    "    # build index\n",
    "    lyft_index = VectorStoreIndex.from_documents(lyft_docs)\n",
    "    uber_index = VectorStoreIndex.from_documents(uber_docs)\n",
    "\n",
    "    # persist index\n",
    "    lyft_index.storage_context.persist(persist_dir=\"./storage/lyft\")\n",
    "    uber_index.storage_context.persist(persist_dir=\"./storage/uber\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31892898-a2dc-43c8-812a-3442feb2108d",
   "metadata": {},
   "outputs": [],
   "source": [
    "lyft_engine = lyft_index.as_query_engine(similarity_top_k=3)\n",
    "uber_engine = uber_index.as_query_engine(similarity_top_k=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9f3158a-7647-4442-8de1-4db80723b4d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.tools import QueryEngineTool\n",
    "\n",
    "query_engine_tools = [\n",
    "    QueryEngineTool.from_defaults(\n",
    "        query_engine=lyft_engine,\n",
    "        name=\"lyft_10k\",\n",
    "        description=(\n",
    "            \"Provides information about Lyft financials for year 2021. \"\n",
    "            \"Use a detailed plain text question as input to the tool.\"\n",
    "        ),\n",
    "    ),\n",
    "    QueryEngineTool.from_defaults(\n",
    "        query_engine=uber_engine,\n",
    "        name=\"uber_10k\",\n",
    "        description=(\n",
    "            \"Provides information about Uber financials for year 2021. \"\n",
    "            \"Use a detailed plain text question as input to the tool.\"\n",
    "        ),\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "275c01b1-8dce-4216-9203-1e961b7fc313",
   "metadata": {},
   "source": [
    "## Setup ReAct Agent\n",
    "\n",
    "Here we setup our ReAct agent with the tools we created above.\n",
    "\n",
    "You can **optionally** specify a system prompt which will be added to the core ReAct system prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ded93297-fee8-4329-bf37-cf77e87621ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import ReActAgent\n",
    "from llama_index.core.workflow import Context\n",
    "\n",
    "agent = ReActAgent(\n",
    "    tools=query_engine_tools,\n",
    "    llm=OpenAI(model=\"gpt-4o-mini\"),\n",
    "    # system_prompt=\"...\"\n",
    ")\n",
    "\n",
    "# context to hold this session/state\n",
    "\n",
    "ctx = Context(agent)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73447845",
   "metadata": {},
   "source": [
    "## Run Some Example Queries\n",
    "\n",
    "By streaming the result, we can see the full response, including the thought process and tool calls.\n",
    "\n",
    "If we wanted to stream only the result, we can buffer the stream and start streaming once `Answer:` is in the response.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70a82471-9226-42ad-bd8a-aebde3530d95",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Thought: The current language of the user is: English. I need to use a tool to help me answer the question.\n",
      "Action: lyft_10k\n",
      "Action Input: {\"input\": \"What was Lyft's revenue growth in 2021?\"}Thought: I can answer without using any more tools. I'll use the user's language to answer.\n",
      "Answer: Lyft's revenue growth in 2021 was 36% compared to the prior year."
     ]
    }
   ],
   "source": [
    "from llama_index.core.agent.workflow import ToolCallResult, AgentStream\n",
    "\n",
    "handler = agent.run(\"What was Lyft's revenue growth in 2021?\", ctx=ctx)\n",
    "\n",
    "async for ev in handler.stream_events():\n",
    "    # if isinstance(ev, ToolCallResult):\n",
    "    #     print(f\"\\nCall {ev.tool_name} with {ev.tool_kwargs}\\nReturned: {ev.tool_output}\")\n",
    "    if isinstance(ev, AgentStream):\n",
    "        print(f\"{ev.delta}\", end=\"\", flush=True)\n",
    "\n",
    "response = await handler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f4b6967",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Lyft's revenue growth in 2021 was 36% compared to the prior year.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0873463d-4790-4c17-bfe9-2ece610fe4b3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Thought: The current language of the user is: English. I need to use a tool to gather information about Uber's revenue growth in 2021 to compare it with Lyft's.\n",
      "Action: uber_10k\n",
      "Action Input: {'input': \"What was Uber's revenue growth in 2021?\"}Thought: I now have the revenue growth information for both Uber and Lyft in 2021. Lyft's revenue growth was 36%, while Uber's was 57%. I will now provide a comparison and analysis.\n",
      "Thought: I can answer without using any more tools. I'll use the user's language to answer.\n",
      "Answer: In 2021, Uber experienced a revenue growth of 57%, increasing from $11.139 billion in 2020 to $17.455 billion. In contrast, Lyft's revenue growth was 36%. \n",
      "\n",
      "When comparing the two, Uber outperformed Lyft in terms of revenue growth, indicating a stronger recovery or expansion in its business operations during that year. This could be attributed to Uber's diversified services, including food delivery through Uber Eats, which may have contributed significantly to its revenue. Lyft, primarily focused on ride-sharing, may have faced more challenges in scaling its growth compared to Uber.\n",
      "\n",
      "Overall, while both companies showed positive growth, Uber's higher percentage suggests it was able to capitalize on market opportunities more effectively than Lyft in 2021."
     ]
    }
   ],
   "source": [
    "handler = agent.run(\n",
    "    \"Compare and contrast the revenue growth of Uber and Lyft in 2021, then give an analysis\",\n",
    "    ctx=ctx,\n",
    ")\n",
    "\n",
    "async for ev in handler.stream_events():\n",
    "    # if isinstance(ev, ToolCallResult):\n",
    "    #     print(f\"\\nCall {ev.tool_name} with {ev.tool_kwargs}\\nReturned: {ev.tool_output}\")\n",
    "    if isinstance(ev, AgentStream):\n",
    "        print(f\"{ev.delta}\", end=\"\", flush=True)\n",
    "\n",
    "response = await handler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5882851b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "In 2021, Uber experienced a revenue growth of 57%, increasing from $11.139 billion in 2020 to $17.455 billion. In contrast, Lyft's revenue growth was 36%. \n",
      "\n",
      "When comparing the two, Uber outperformed Lyft in terms of revenue growth, indicating a stronger recovery or expansion in its business operations during that year. This could be attributed to Uber's diversified services, including food delivery through Uber Eats, which may have contributed significantly to its revenue. Lyft, primarily focused on ride-sharing, may have faced more challenges in scaling its growth compared to Uber.\n",
      "\n",
      "Overall, while both companies showed positive growth, Uber's higher percentage suggests it was able to capitalize on market opportunities more effectively than Lyft in 2021.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-index-caVs7DDe-py3.10",
   "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
}