{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a9226af6-32fa-4c3c-966a-e54099fcd09d",
   "metadata": {},
   "source": [
    "# Function Calling\n",
    "\n",
    "Function calling allows Mistral models to connect to external tools. By integrating Mistral models with external tools such as user defined functions or APIs, users can easily build applications catering to specific use cases and practical problems. In this guide, for instance, we wrote two functions for tracking payment status and payment date. We can use these two tools to provide answers for payment-related queries.\n",
    "\n",
    "At a glance, there are four steps with function calling:\n",
    "\n",
    "- User: specify tools and query\n",
    "- Model: Generate function arguments if applicable\n",
    "- User: Execute function to obtain tool results\n",
    "- Model: Generate final answer\n",
    "\n",
    "In this guide, we will walk through a simple example to demonstrate how function calling works with Mistral models in these four steps.\n",
    "\n",
    "Before we get started, let’s assume we have a dataframe consisting of payment transactions. When users ask questions about this dataframe, they can use certain tools to answer questions about this data. This is just an example to emulate an external database that the LLM cannot directly access."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b6a555d-75d1-4241-bee3-a4ffaab4ecd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pandas mistralai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8b4c0408-94a8-4c3f-b518-253b718f436e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# Assuming we have the following data\n",
    "data = {\n",
    "    'transaction_id': ['T1001', 'T1002', 'T1003', 'T1004', 'T1005'],\n",
    "    'customer_id': ['C001', 'C002', 'C003', 'C002', 'C001'],\n",
    "    'payment_amount': [125.50, 89.99, 120.00, 54.30, 210.20],\n",
    "    'payment_date': ['2021-10-05', '2021-10-06', '2021-10-07', '2021-10-05', '2021-10-08'],\n",
    "    'payment_status': ['Paid', 'Unpaid', 'Paid', 'Paid', 'Pending']\n",
    "}\n",
    "\n",
    "# Create DataFrame\n",
    "df = pd.DataFrame(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e0ef041-a23a-437a-b36a-58a1b69be5db",
   "metadata": {},
   "source": [
    "## Step 1. User: specify tools and query\n",
    "\n",
    "### Tools\n",
    "\n",
    "Users can define all the necessary tools for their use cases.\n",
    "\n",
    "- In many cases, we might have multiple tools at our disposal. For example, let’s consider we have two functions as our two tools: `retrieve_payment_status` and `retrieve_payment_date` to retrieve payment status and payment date given transaction ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d9e153c1-6ff6-4878-a8c2-e02bac749210",
   "metadata": {},
   "outputs": [],
   "source": [
    "def retrieve_payment_status(df: data, transaction_id: str) -> str:\n",
    "    if transaction_id in df.transaction_id.values: \n",
    "        return json.dumps({'status': df[df.transaction_id == transaction_id].payment_status.item()})\n",
    "    return json.dumps({'error': 'transaction id not found.'})\n",
    "\n",
    "def retrieve_payment_date(df: data, transaction_id: str) -> str:\n",
    "    if transaction_id in df.transaction_id.values: \n",
    "        return json.dumps({'date': df[df.transaction_id == transaction_id].payment_date.item()})\n",
    "    return json.dumps({'error': 'transaction id not found.'})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5728a885-ec7c-4d21-b461-ebb5e29b0a1e",
   "metadata": {},
   "source": [
    "- In order for Mistral models to understand the functions, we need to outline the function specifications with a JSON schema. Specifically, we need to describe the type, function name, function description, function parameters, and the required parameter for the function. Since we have two functions here, let’s list two function specifications in a list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "756006e4-ae7e-4023-89d5-7de970f2efe7",
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [\n",
    "    {\n",
    "        \"type\": \"function\",\n",
    "        \"function\": {\n",
    "            \"name\": \"retrieve_payment_status\",\n",
    "            \"description\": \"Get payment status of a transaction\",\n",
    "            \"parameters\": {\n",
    "                \"type\": \"object\",\n",
    "                \"properties\": {\n",
    "                    \"transaction_id\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"The transaction id.\",\n",
    "                    }\n",
    "                },\n",
    "                \"required\": [\"transaction_id\"],\n",
    "            },\n",
    "        },\n",
    "    },\n",
    "    {\n",
    "        \"type\": \"function\",\n",
    "        \"function\": {\n",
    "            \"name\": \"retrieve_payment_date\",\n",
    "            \"description\": \"Get payment date of a transaction\",\n",
    "            \"parameters\": {\n",
    "                \"type\": \"object\",\n",
    "                \"properties\": {\n",
    "                    \"transaction_id\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"The transaction id.\",\n",
    "                    }\n",
    "                },\n",
    "                \"required\": [\"transaction_id\"],\n",
    "            },\n",
    "        },\n",
    "    }\n",
    "]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34144aca-5621-4190-9e7e-a05e4266b6d0",
   "metadata": {},
   "source": [
    "- Then we organize the two functions into a dictionary where keys represent the function name, and values are the function with the df defined. This allows us to call each function based on its function name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "88d3fda1-72bb-4ee3-a80f-340716a22a01",
   "metadata": {},
   "outputs": [],
   "source": [
    "import functools\n",
    "\n",
    "names_to_functions = {\n",
    "  'retrieve_payment_status': functools.partial(retrieve_payment_status, df=df),\n",
    "  'retrieve_payment_date': functools.partial(retrieve_payment_date, df=df)\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0a05bf6-c244-4124-8196-f9a780fef95d",
   "metadata": {},
   "source": [
    "### User query\n",
    "\n",
    "Suppose a user asks the following question: “What’s the status of my transaction?” A standalone LLM would not be able to answer this question, as it needs to query the business logic backend to access the necessary data. But what if we have an exact tool we can use to answer this question? We could potentially provide an answer!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6777e7a6-f7e2-42bd-a141-5afe8fbd3c74",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mistralai.models.chat_completion import ChatMessage\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(role=\"user\", content=\"What's the status of my transaction T1001?\")\n",
    "]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d926683-3cc1-4de6-ad53-adefe8d5cc0b",
   "metadata": {},
   "source": [
    "## Step 2. Model: Generate function arguments \n",
    "\n",
    "How do Mistral models know about these functions and know which function to use? We provide both the user query and the tools specifications to Mistral models. The goal in this step is not for the Mistral model to run the function directly. It’s to 1) determine the appropriate function to use , 2) identify if there is any essential information missing for a function, and 3) generate necessary arguments for the chosen function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "59493530-b310-447f-9f77-32d9013ffd30",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ChatCompletionResponse(id='a26fd9e27f5646f5a73a2a98df1605b9', object='chat.completion', created=1716472309, model='mistral-small-latest', choices=[ChatCompletionResponseChoice(index=0, message=ChatMessage(role='assistant', content='', name=None, tool_calls=[ToolCall(id='cn0Be32Nw', type=<ToolType.function: 'function'>, function=FunctionCall(name='retrieve_payment_status', arguments='{\"transaction_id\": \"T1001\"}'))], tool_call_id=None), finish_reason=<FinishReason.tool_calls: 'tool_calls'>)], usage=UsageInfo(prompt_tokens=166, total_tokens=196, completion_tokens=30))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from mistralai.client import MistralClient\n",
    "\n",
    "model = \"mistral-small-latest\"\n",
    "api_key=\"TYPE YOUR API KEY\"\n",
    "\n",
    "client = MistralClient(api_key=api_key)\n",
    "\n",
    "response = client.chat(\n",
    "    model=model,\n",
    "    messages=messages,\n",
    "    tools=tools,\n",
    "    tool_choice=\"auto\"\n",
    ")\n",
    "\n",
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a7af5b4d-2128-4833-bd94-b4d924495783",
   "metadata": {},
   "outputs": [],
   "source": [
    "messages.append(response.choices[0].message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4a31eeaa-6ef2-4af8-8fbf-e5648d12a263",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[ChatMessage(role='user', content=\"What's the status of my transaction T1001?\", name=None, tool_calls=None, tool_call_id=None),\n",
       " ChatMessage(role='assistant', content='', name=None, tool_calls=[ToolCall(id='cn0Be32Nw', type=<ToolType.function: 'function'>, function=FunctionCall(name='retrieve_payment_status', arguments='{\"transaction_id\": \"T1001\"}'))], tool_call_id=None)]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffab43e6-3225-45b4-bc85-3689c2eee7c3",
   "metadata": {},
   "source": [
    "## Step 3. User: Execute function to obtain tool results\n",
    "\n",
    "How do we execute the function? Currently, it is the user’s responsibility to execute these functions and the function execution lies on the user side. In the future, we may introduce some helpful functions that can be executed server-side.\n",
    "\n",
    "Let’s extract some useful function information from model response including function_name and function_params. It’s clear here that our Mistral model has chosen to use the function `retrieve_payment_status` with the parameter `transaction_id` set to T1001."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "db259833-75b0-4538-a78e-8a1566326ec9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "function_name:  retrieve_payment_status \n",
      "function_params:  {'transaction_id': 'T1001'}\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "tool_call = response.choices[0].message.tool_calls[0]\n",
    "function_name = tool_call.function.name\n",
    "function_params = json.loads(tool_call.function.arguments)\n",
    "print(\"\\nfunction_name: \", function_name, \"\\nfunction_params: \", function_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5eae109b-41c3-463d-98da-0e5499ea8a37",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'{\"status\": \"Paid\"}'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function_result = names_to_functions[function_name](**function_params)\n",
    "function_result\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "81a14939-b6a9-45ba-92cb-b190969f6889",
   "metadata": {},
   "outputs": [],
   "source": [
    "messages.append(ChatMessage(role=\"tool\", name=function_name, content=function_result, tool_call_id=tool_call.id))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "fafc8f24-f1a8-4d8a-ad7f-8465da97939f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[ChatMessage(role='user', content=\"What's the status of my transaction T1001?\", name=None, tool_calls=None, tool_call_id=None),\n",
       " ChatMessage(role='assistant', content='', name=None, tool_calls=[ToolCall(id='cn0Be32Nw', type=<ToolType.function: 'function'>, function=FunctionCall(name='retrieve_payment_status', arguments='{\"transaction_id\": \"T1001\"}'))], tool_call_id=None),\n",
       " ChatMessage(role='tool', content='{\"status\": \"Paid\"}', name='retrieve_payment_status', tool_calls=None, tool_call_id='cn0Be32Nw')]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2aa2e194-1e97-4446-817d-9b57dcfab371",
   "metadata": {},
   "source": [
    "## Step 4. Model: Generate final answer\n",
    "\n",
    "We can now provide the output from the tools to Mistral models, and in return, the Mistral model can produce a customised final response for the specific user."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "36e1cbb5-96bf-49b6-bada-ff6338d67d86",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Your transaction T1001 has been successfully paid.'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = client.chat(\n",
    "    model=model,\n",
    "    messages=messages\n",
    ")\n",
    "\n",
    "response.choices[0].message.content"
   ]
  }
 ],
 "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",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}