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

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "24103c51",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/agent/mistral_agent.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99cea58c-48bc-4af6-8358-df9695659983",
   "metadata": {},
   "source": [
    "# Function Calling Mistral Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "673df1fe-eb6c-46ea-9a73-a96e7ae7942e",
   "metadata": {},
   "source": [
    "This notebook shows you how to use our Mistral agent, powered by function calling capabilities."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54b7bc2e-606f-411a-9490-fcfab9236dfc",
   "metadata": {},
   "source": [
    "## Initial Setup "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23e80e5b-aaee-4f23-b338-7ae62b08141f",
   "metadata": {},
   "source": [
    "Let's start by importing some simple building blocks.  \n",
    "\n",
    "The main thing we need is:\n",
    "1. the OpenAI API (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": "markdown",
   "id": "41101795",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4985c578",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index\n",
    "%pip install llama-index-llms-mistralai\n",
    "%pip install llama-index-embeddings-mistralai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fe08eb1-e638-4c00-9103-5c305bfacccf",
   "metadata": {},
   "source": [
    "Let's define some very simple calculator tools for our agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dd3c4a6-f3e0-46f9-ad3b-7ba57d1bc992",
   "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",
   "id": "eeac7d4c-58fd-42a5-9da9-c258375c61a0",
   "metadata": {},
   "source": [
    "Make sure your MISTRAL_API_KEY is set. Otherwise explicitly specify the `api_key` parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4becf171-6632-42e5-bdec-918a00934696",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.mistralai import MistralAI\n",
    "\n",
    "llm = MistralAI(model=\"mistral-large-latest\", api_key=\"...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "707d30b8-6405-4187-a9ed-6146dcc42167",
   "metadata": {},
   "source": [
    "## Initialize Mistral Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "798ca3fd-6711-4c0c-a853-d868dd14b484",
   "metadata": {},
   "source": [
    "Here we initialize a simple Mistral agent with calculator functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38ab3938-1138-43ea-b085-f430b42f5377",
   "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",
   "id": "500cbee4",
   "metadata": {},
   "source": [
    "### Chat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9450401d-769f-46e8-8bab-0f27f7362f5d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Added user message to memory: What is (121 + 2) * 5?\n",
      "=== Calling Function ===\n",
      "Calling function: add with args: {\"a\": 121, \"b\": 2}\n",
      "=== Calling Function ===\n",
      "Calling function: multiply with args: {\"a\": 123, \"b\": 5}\n",
      "assistant: The result of (121 + 2) * 5 is 615.\n"
     ]
    }
   ],
   "source": [
    "response = await agent.run(\"What is (121 + 2) * 5?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "538bf32f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# inspect sources\n",
    "print(response.tool_calls)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7e29a33",
   "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,
   "id": "f6d755ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import Context\n",
    "\n",
    "ctx = Context(agent)\n",
    "\n",
    "response = await agent.run(\"My name is John Doe\", ctx=ctx)\n",
    "response = await agent.run(\"What is my name?\", ctx=ctx)\n",
    "\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cabfdf01-8d63-43ff-b06e-a3059ede2ddf",
   "metadata": {},
   "source": [
    "## Mistral Agent over RAG Pipeline\n",
    "\n",
    "Build a Mistral agent over a simple 10K document. We use both Mistral embeddings and mistral-medium to construct the RAG pipeline, and pass it to the Mistral agent as a tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48120dd4-7f50-426f-bc7e-a903e090d32e",
   "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,
   "id": "48c0cf98-3f10-4599-8437-d88dc89cefad",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.tools import QueryEngineTool\n",
    "from llama_index.core import SimpleDirectoryReader, VectorStoreIndex\n",
    "from llama_index.embeddings.mistralai import MistralAIEmbedding\n",
    "from llama_index.llms.mistralai import MistralAI\n",
    "\n",
    "embed_model = MistralAIEmbedding(api_key=\"...\")\n",
    "query_llm = MistralAI(model=\"mistral-medium\", api_key=\"...\")\n",
    "\n",
    "# load data\n",
    "uber_docs = SimpleDirectoryReader(\n",
    "    input_files=[\"./data/10k/uber_2021.pdf\"]\n",
    ").load_data()\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=query_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,
   "id": "ebfdaf80-e5e1-4c60-b556-20558da3d5e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import FunctionAgent\n",
    "\n",
    "agent = FunctionAgent(tools=[query_engine_tool], llm=llm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58c53f2a-0a3f-4abe-b8b6-97a974ec7546",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Added user message to memory: Tell me both the risk factors and tailwinds for Uber? Do two parallel tool calls.\n",
      "=== Calling Function ===\n",
      "Calling function: uber_10k with args: {\"input\": \"What are the risk factors for Uber in 2021?\"}\n",
      "=== Calling Function ===\n",
      "Calling function: uber_10k with args: {\"input\": \"What are the tailwinds for Uber in 2021?\"}\n",
      "assistant: Based on the information provided, here are the risk factors for Uber in 2021:\n",
      "\n",
      "1. Failure to offer or develop autonomous vehicle technologies, which could result in inferior performance or safety concerns compared to competitors.\n",
      "2. Dependence on high-quality personnel and the potential impact of attrition or unsuccessful succession planning on the business.\n",
      "3. Security or data privacy breaches, unauthorized access, or destruction of proprietary, employee, or user data.\n",
      "4. Cyberattacks, such as malware, ransomware, viruses, spamming, and phishing attacks, which could harm the company's reputation and operations.\n",
      "5. Climate change risks, including physical and transitional risks, that may adversely impact the business if not managed effectively.\n",
      "6. Reliance on third parties to maintain open marketplaces for distributing products and providing software, which could negatively affect the business if interfered with.\n",
      "7. The need for additional capital to support business growth, which may not be available on reasonable terms or at all.\n",
      "8. Difficulties in identifying, acquiring, and integrating suitable businesses, which could harm operating results and prospects.\n",
      "9. Legal and regulatory risks, including extensive government regulation and oversight related to payment and financial services.\n",
      "10. Intellectual property risks, such as the inability to protect intellectual property or claims of misappropriation by third parties.\n",
      "11. Volatility in the market price of common stock, which could result in steep declines and loss of investment for shareholders.\n",
      "12. Economic risks related to the COVID-19 pandemic, which has adversely impacted and could continue to adversely impact the business, financial condition, and results of operations.\n",
      "13. The potential reclassification of Drivers as employees, workers, or quasi-employees, which could result in material costs associated with defending, settling, or resolving lawsuits and demands for arbitration.\n",
      "\n",
      "On the other hand, here are some tailwinds for Uber in 2021:\n",
      "\n",
      "1. Launch of Uber One, a single cross-platform membership program in the United States, which offers discounts, special pricing, priority service, and exclusive perks across rides, delivery, and grocery offerings.\n",
      "2. Introduction of a \"Super App\" view on iOS\n"
     ]
    }
   ],
   "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))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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
}