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

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "26ad48bf",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/agent/openai_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": [
    "# Agent with Query Engine Tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db402a8b-90d6-4e1d-8df6-347c54624f26",
   "metadata": {},
   "source": [
    "## Build Query Engine Tools"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "30e2aa77",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0ed4104",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02160804-64a2-4ef3-8a0d-8c16b06fd205",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4968218b",
   "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"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0a875e00",
   "metadata": {},
   "source": [
    "Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2ae5855",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p 'data/10k/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/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/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 Agent\n",
    "\n",
    "For LLMs like OpenAI that have a function calling API, we should use the `FunctionAgent`.\n",
    "\n",
    "For other LLMs, we can use the `ReActAgent`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32f71a46-bdf6-4365-b1f1-e23a0d913a3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import FunctionAgent, ReActAgent\n",
    "from llama_index.core.workflow import Context\n",
    "\n",
    "agent = FunctionAgent(tools=query_engine_tools, llm=OpenAI(model=\"gpt-4o\"))\n",
    "\n",
    "# context to hold the session/state\n",
    "ctx = Context(agent)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22716961-11c3-4ac4-82a8-419f787bc36a",
   "metadata": {},
   "source": [
    "## Let's Try It Out!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42a1bce0-b398-4937-9008-6cee04368ac4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Call lyft_10k with args {'input': \"What was Lyft's revenue for the year 2021?\"}\n",
      "Returned: Lyft's revenue for the year 2021 was $3,208,323,000.\n",
      "Call uber_10k with args {'input': \"What was Uber's revenue for the year 2021?\"}\n",
      "Returned: Uber's revenue for the year 2021 was $17.455 billion.\n",
      "In 2021, Lyft's revenue was approximately $3.21 billion, while Uber's revenue was significantly higher at $17.455 billion."
     ]
    }
   ],
   "source": [
    "from llama_index.core.agent.workflow import ToolCallResult, AgentStream\n",
    "\n",
    "handler = agent.run(\"What's the revenue for Lyft in 2021 vs Uber?\", ctx=ctx)\n",
    "\n",
    "async for ev in handler.stream_events():\n",
    "    if isinstance(ev, ToolCallResult):\n",
    "        print(\n",
    "            f\"Call {ev.tool_name} with args {ev.tool_kwargs}\\nReturned: {ev.tool_output}\"\n",
    "        )\n",
    "    elif isinstance(ev, AgentStream):\n",
    "        print(ev.delta, end=\"\", flush=True)\n",
    "\n",
    "response = await handler"
   ]
  }
 ],
 "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
}