{ "cells": [ { "attachments": {}, "cell_type": "markdown", "id": "26ad48bf", "metadata": {}, "source": [ "\"Open" ] }, { "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 }