wehub-resource-sync
a0c8464e58
Build Package / build (ubuntu-latest) (push) Failing after 1s
CodeQL / Analyze (python) (push) Failing after 1s
Core Typecheck / core-typecheck (push) Failing after 1s
Linting / lint (push) Failing after 1s
llama-dev tests / test-llama-dev (push) Failing after 1s
Publish Sub-Package to PyPI if Needed / publish_subpackage_if_needed (push) Has been skipped
Sync Docs to Developer Hub / sync-docs (push) Failing after 0s
Build Package / build (windows-latest) (push) Has been cancelled
504 lines
15 KiB
Plaintext
504 lines
15 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "9e34586b",
|
|
"metadata": {},
|
|
"source": [
|
|
"<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/agent/openai_agent_context_retrieval.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "99cea58c-48bc-4af6-8358-df9695659983",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Context-Augmented Function Calling Agent"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "673df1fe-eb6c-46ea-9a73-a96e7ae7942e",
|
|
"metadata": {},
|
|
"source": [
|
|
"In this tutorial, we show you how to to make your agent context-aware.\n",
|
|
"\n",
|
|
"Our indexing/retrieval modules help to remove the complexity of having too many functions to fit in the prompt."
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "54b7bc2e-606f-411a-9490-fcfab9236dfc",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Initial Setup "
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "23e80e5b-aaee-4f23-b338-7ae62b08141f",
|
|
"metadata": {},
|
|
"source": [
|
|
"Here we setup a normal FunctionAgent, and then augment it with context. This agent will perform retrieval first before calling any tools. This can help ground the agent's tool picking and answering capabilities in context."
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "d6d2e0ae",
|
|
"metadata": {},
|
|
"source": [
|
|
"If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "c40cadb6",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install llama-index"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "36f37aa5",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import os\n",
|
|
"\n",
|
|
"os.environ[\"OPENAI_API_KEY\"] = \"sk-...\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "e0f35c0d",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.llms.openai import OpenAI\n",
|
|
"from llama_index.embeddings.openai import OpenAIEmbedding\n",
|
|
"from llama_index.core.settings 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": "9d47283b-025e-4874-88ed-76245b22f82e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import json\n",
|
|
"from typing import Sequence\n",
|
|
"\n",
|
|
"from llama_index.core import (\n",
|
|
" SimpleDirectoryReader,\n",
|
|
" VectorStoreIndex,\n",
|
|
" StorageContext,\n",
|
|
" load_index_from_storage,\n",
|
|
")\n",
|
|
"from llama_index.core.tools import QueryEngineTool"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "9f034df0-f6f1-4ffb-9c4b-d68c2202051c",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"try:\n",
|
|
" storage_context = StorageContext.from_defaults(\n",
|
|
" persist_dir=\"./storage/march\"\n",
|
|
" )\n",
|
|
" march_index = load_index_from_storage(storage_context)\n",
|
|
"\n",
|
|
" storage_context = StorageContext.from_defaults(\n",
|
|
" persist_dir=\"./storage/june\"\n",
|
|
" )\n",
|
|
" june_index = load_index_from_storage(storage_context)\n",
|
|
"\n",
|
|
" storage_context = StorageContext.from_defaults(\n",
|
|
" persist_dir=\"./storage/sept\"\n",
|
|
" )\n",
|
|
" sept_index = load_index_from_storage(storage_context)\n",
|
|
"\n",
|
|
" index_loaded = True\n",
|
|
"except:\n",
|
|
" index_loaded = False"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "f9010d88",
|
|
"metadata": {},
|
|
"source": [
|
|
"Download Data"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "75aa7b79",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!mkdir -p 'data/10q/'\n",
|
|
"!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10q/uber_10q_march_2022.pdf' -O 'data/10q/uber_10q_march_2022.pdf'\n",
|
|
"!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10q/uber_10q_june_2022.pdf' -O 'data/10q/uber_10q_june_2022.pdf'\n",
|
|
"!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10q/uber_10q_sept_2022.pdf' -O 'data/10q/uber_10q_sept_2022.pdf'"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "2a90e4d4-589b-4349-a134-a2ef931d8e89",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# build indexes across the three data sources\n",
|
|
"if not index_loaded:\n",
|
|
" # load data\n",
|
|
" march_docs = SimpleDirectoryReader(\n",
|
|
" input_files=[\"./data/10q/uber_10q_march_2022.pdf\"]\n",
|
|
" ).load_data()\n",
|
|
" june_docs = SimpleDirectoryReader(\n",
|
|
" input_files=[\"./data/10q/uber_10q_june_2022.pdf\"]\n",
|
|
" ).load_data()\n",
|
|
" sept_docs = SimpleDirectoryReader(\n",
|
|
" input_files=[\"./data/10q/uber_10q_sept_2022.pdf\"]\n",
|
|
" ).load_data()\n",
|
|
"\n",
|
|
" # build index\n",
|
|
" march_index = VectorStoreIndex.from_documents(march_docs)\n",
|
|
" june_index = VectorStoreIndex.from_documents(june_docs)\n",
|
|
" sept_index = VectorStoreIndex.from_documents(sept_docs)\n",
|
|
"\n",
|
|
" # persist index\n",
|
|
" march_index.storage_context.persist(persist_dir=\"./storage/march\")\n",
|
|
" june_index.storage_context.persist(persist_dir=\"./storage/june\")\n",
|
|
" sept_index.storage_context.persist(persist_dir=\"./storage/sept\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "00f82f26-83a6-4a37-8a9f-55bf49d8b247",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"march_engine = march_index.as_query_engine(similarity_top_k=3)\n",
|
|
"june_engine = june_index.as_query_engine(similarity_top_k=3)\n",
|
|
"sept_engine = sept_index.as_query_engine(similarity_top_k=3)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "1edb4379-75b8-4b83-8a2d-16170fa6cb67",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"query_engine_tools = [\n",
|
|
" QueryEngineTool.from_defaults(\n",
|
|
" query_engine=march_engine,\n",
|
|
" name=\"uber_march_10q\",\n",
|
|
" description=(\n",
|
|
" \"Provides information about Uber 10Q filings for March 2022. \"\n",
|
|
" \"Use a detailed plain text question as input to the tool.\"\n",
|
|
" ),\n",
|
|
" ),\n",
|
|
" QueryEngineTool.from_defaults(\n",
|
|
" query_engine=june_engine,\n",
|
|
" name=\"uber_june_10q\",\n",
|
|
" description=(\n",
|
|
" \"Provides information about Uber financials for June 2021. \"\n",
|
|
" \"Use a detailed plain text question as input to the tool.\"\n",
|
|
" ),\n",
|
|
" ),\n",
|
|
" QueryEngineTool.from_defaults(\n",
|
|
" query_engine=sept_engine,\n",
|
|
" name=\"uber_sept_10q\",\n",
|
|
" description=(\n",
|
|
" \"Provides information about Uber financials for Sept 2021. \"\n",
|
|
" \"Use a detailed plain text question as input to the tool.\"\n",
|
|
" ),\n",
|
|
" ),\n",
|
|
"]"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "b08efb96-ce44-4706-a22d-b0c670b23a60",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Try Context-Augmented Agent\n",
|
|
"\n",
|
|
"Here we augment our agent with context in different settings:\n",
|
|
"- toy context: we define some abbreviations that map to financial terms (e.g. R=Revenue). We supply this as context to the agent"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "72709da5-785e-4b9d-9e8f-231a2d2fbb53",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core import Document\n",
|
|
"from llama_index.core.agent.workflow import FunctionAgent"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "d2b235ed-d2a0-46cb-830b-d1a3affeb0c2",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# toy index - stores a list of abbreviations\n",
|
|
"texts = [\n",
|
|
" \"Abbreviation: 'Y' = Revenue\",\n",
|
|
" \"Abbreviation: 'X' = Risk Factors\",\n",
|
|
" \"Abbreviation: 'Z' = Costs\",\n",
|
|
"]\n",
|
|
"docs = [Document(text=t) for t in texts]\n",
|
|
"context_index = VectorStoreIndex.from_documents(docs)\n",
|
|
"context_retriever = context_index.as_retriever(similarity_top_k=2)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "dbd2710f",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core.tools import BaseTool\n",
|
|
"\n",
|
|
"system_prompt_template = \"\"\"You are a helpful assistant. \n",
|
|
"Here is some context that you can use to answer the user's question and for help with picking the right tool:\n",
|
|
"\n",
|
|
"{context}\n",
|
|
"\"\"\"\n",
|
|
"\n",
|
|
"\n",
|
|
"async def get_agent_with_context_awareness(\n",
|
|
" query: str, context_retriever, tools: list[BaseTool]\n",
|
|
") -> FunctionAgent:\n",
|
|
" context_nodes = await context_retriever.aretrieve(query)\n",
|
|
" context_text = \"\\n----\\n\".join([n.node.text for n in context_nodes])\n",
|
|
"\n",
|
|
" return FunctionAgent(\n",
|
|
" tools=tools,\n",
|
|
" llm=OpenAI(model=\"gpt-4o\"),\n",
|
|
" system_prompt=system_prompt_template.format(context=context_text),\n",
|
|
" )"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "b44e9819-69a5-4085-9957-27d8eb940d24",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"query = \"What is the 'X' of March 2022?\"\n",
|
|
"agent = await get_agent_with_context_awareness(\n",
|
|
" query, context_retriever, query_engine_tools\n",
|
|
")\n",
|
|
"\n",
|
|
"response = await agent.run(query)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "9969fceb-d9b2-491f-92f2-d5e0fe5a8ed2",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"The risk factors mentioned in Uber's 10-Q filing for March 2022 include uncertainties related to the COVID-19 pandemic, such as the severity and duration of the outbreak, potential future waves or variants of the virus, the administration and efficacy of vaccines, and the impact of governmental actions. There are also concerns regarding the effects on drivers, merchants, consumers, and business partners, as well as other factors that may affect the company's business, results of operations, financial position, and cash flows.\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"print(str(response))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "c407dd42-39a3-4bda-8294-27ac043a5adb",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"query = \"What is the 'Y' and 'Z' in September 2022?\"\n",
|
|
"agent = await get_agent_with_context_awareness(\n",
|
|
" query, context_retriever, query_engine_tools\n",
|
|
")\n",
|
|
"\n",
|
|
"response = await agent.run(query)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "ef4e84af",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"In September 2022, Uber's revenue (Y) was $8,343 million, and the total costs (Z) were $8,839 million.\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"print(str(response))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "9e46a631",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Managing Context/Memory\n",
|
|
"\n",
|
|
"By default, each `.run()` call is stateless. We can manage context by using a serializable `Context` object."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "edbadfe2",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"You asked for the revenue ('Y') and costs ('Z') for Uber in September 2022.\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"from llama_index.core.workflow import Context\n",
|
|
"\n",
|
|
"ctx = Context(agent)\n",
|
|
"\n",
|
|
"query = \"What is the 'Y' and 'Z' in September 2022?\"\n",
|
|
"agent = await get_agent_with_context_awareness(\n",
|
|
" query, context_retriever, query_engine_tools\n",
|
|
")\n",
|
|
"response = await agent.run(query, ctx=ctx)\n",
|
|
"\n",
|
|
"query = \"What did I just ask?\"\n",
|
|
"agent = await get_agent_with_context_awareness(\n",
|
|
" query, context_retriever, query_engine_tools\n",
|
|
")\n",
|
|
"response = await agent.run(query, ctx=ctx)\n",
|
|
"print(str(response))"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "ad81c4e1-4ecb-405d-bb03-a4c3549816e7",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Use Uber 10-Q as context, use Calculator as Tool"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "f8dc898e-5fe0-45a2-8e04-debdaeb2c1bb",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core.tools import FunctionTool\n",
|
|
"\n",
|
|
"\n",
|
|
"def magic_formula(revenue: int, cost: int) -> int:\n",
|
|
" \"\"\"Runs MAGIC_FORMULA on revenue and cost.\"\"\"\n",
|
|
" return revenue - cost\n",
|
|
"\n",
|
|
"\n",
|
|
"magic_tool = FunctionTool.from_defaults(magic_formula)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "4dc2e5b7-3b41-43ea-91db-847cf28fc6a8",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"context_retriever = sept_index.as_retriever(similarity_top_k=3)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "99b9a0f8-0029-495c-a44b-913d1e0556e9",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"The result of running MAGIC_FORMULA on Uber's revenue of $8,343 million and cost of $5,173 million is 3,170.\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"query = \"Can you run MAGIC_FORMULA on Uber's revenue and cost?\"\n",
|
|
"agent = await get_agent_with_context_awareness(\n",
|
|
" query, context_retriever, [magic_tool]\n",
|
|
")\n",
|
|
"response = await agent.run(query)\n",
|
|
"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
|
|
}
|