{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Function Calling NVIDIA Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook shows you how to use our NVIDIA agent, powered by function calling capabilities."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initial Setup "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start by importing some simple building blocks.  \n",
    "\n",
    "The main thing we need is:\n",
    "1. the NVIDIA NIM Endpoint (using our own `llama_index` LLM class)\n",
    "2. a place to keep conversation history \n",
    "3. a definition for tools that our agent can use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade --quiet llama-index-llms-nvidia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Valid NVIDIA_API_KEY already in environment. Delete to reset\n"
     ]
    }
   ],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "# del os.environ['NVIDIA_API_KEY']  ## delete key and reset\n",
    "if os.environ.get(\"NVIDIA_API_KEY\", \"\").startswith(\"nvapi-\"):\n",
    "    print(\"Valid NVIDIA_API_KEY already in environment. Delete to reset\")\n",
    "else:\n",
    "    nvapi_key = getpass.getpass(\"NVAPI Key (starts with nvapi-): \")\n",
    "    assert nvapi_key.startswith(\n",
    "        \"nvapi-\"\n",
    "    ), f\"{nvapi_key[:5]}... is not a valid key\"\n",
    "    os.environ[\"NVIDIA_API_KEY\"] = nvapi_key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.nvidia import NVIDIA\n",
    "from llama_index.core.tools import FunctionTool\n",
    "from llama_index.embeddings.nvidia import NVIDIAEmbedding"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define some very simple calculator tools for our agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
    "def add(a: int, b: int) -> int:\n",
    "    \"\"\"Add two integers and returns the result integer\"\"\"\n",
    "    return a + b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we initialize a simple NVIDIA agent with calculator functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = NVIDIA(\"meta/llama-3.1-70b-instruct\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import FunctionAgent\n",
    "\n",
    "agent = FunctionAgent(\n",
    "    tools=[multiply, add],\n",
    "    llm=llm,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Chat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await agent.run(\"What is (121 * 3) + 42?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# inspect sources\n",
    "print(response.tool_calls)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Managing Context/Memory\n",
    "\n",
    "By default, `.run()` is stateless. If you want to maintain state, you can pass in a `context` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import Context\n",
    "\n",
    "ctx = Context(agent)\n",
    "\n",
    "response = await agent.run(\"Hello, my name is John Doe.\", ctx=ctx)\n",
    "print(str(response))\n",
    "\n",
    "response = await agent.run(\"What is my name?\", ctx=ctx)\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Agent with Personality"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can specify a system prompt to give the agent additional instruction or personality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = FunctionAgent(\n",
    "    tools=[multiply, add],\n",
    "    llm=llm,\n",
    "    system_prompt=\"Talk like a pirate in every response.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await agent.run(\"Hi\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await agent.run(\"Tell me a story\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NVIDIA Agent with RAG/Query Engine Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.tools import QueryEngineTool\n",
    "from llama_index.core import SimpleDirectoryReader, VectorStoreIndex\n",
    "\n",
    "embed_model = NVIDIAEmbedding(model=\"NV-Embed-QA\", truncate=\"END\")\n",
    "\n",
    "# load data\n",
    "uber_docs = SimpleDirectoryReader(\n",
    "    input_files=[\"./data/10k/uber_2021.pdf\"]\n",
    ").load_data()\n",
    "\n",
    "# build index\n",
    "uber_index = VectorStoreIndex.from_documents(\n",
    "    uber_docs, embed_model=embed_model\n",
    ")\n",
    "uber_engine = uber_index.as_query_engine(similarity_top_k=3, llm=llm)\n",
    "query_engine_tool = 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",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = FunctionAgent(tools=[query_engine_tool], llm=llm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await agent.run(\n",
    "    \"Tell me both the risk factors and tailwinds for Uber? Do two parallel tool calls.\"\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ReAct Agent "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import ReActAgent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = ReActAgent([multiply_tool, add_tool], llm=llm, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the `stream_events()` method, we can stream the response as it is generated to see the agent's thought process.\n",
    "\n",
    "The final response will have only the final answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import AgentStream\n",
    "\n",
    "handler = agent.run(\"What is 20+(2*4)? Calculate step by step \")\n",
    "async for ev in handler.stream_events():\n",
    "    if isinstance(ev, AgentStream):\n",
    "        print(ev.delta, end=\"\", flush=True)\n",
    "\n",
    "response = await handler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(response.tool_calls)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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": 4
}