{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "43c445ae-9c30-4661-ad86-65be497f792c",
   "metadata": {},
   "source": [
    "# Function Calling Rest API with Mistral7Bv3 using Ollama\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 a Pet Store's Pets and User info. We can use these two tools to provide answers for pet-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 an OpenAPI spec end-points consisting of Pet store information. When users ask questions about this API, they can use certain tools to answer questions about this data. This is just an example to emulate an external database via API that the LLM cannot directly access."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bdc3e6c-70c2-4ea5-b526-a70e4403828f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --upgrade ollama mistral-common pandas\n",
    "!pip install --upgrade  prance openapi-spec-validator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4950dee7-2bd2-4997-8018-1179882f4d7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import prance\n",
    "from typing import List\n",
    "from mistral_common.tokens.tokenizers.mistral import MistralTokenizer\n",
    "from mistral_common.protocol.instruct.request import ChatCompletionRequest\n",
    "from mistral_common.protocol.instruct.tool_calls import Function, Tool\n",
    "import ollama\n",
    "from mistral_common.protocol.instruct.messages import UserMessage\n",
    "import json\n",
    "import requests\n",
    "import functools\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a327f736-51b1-41c8-800c-49a2df55a083",
   "metadata": {},
   "source": [
    "Setup functions to make REST API call. We take example of pet store from [Swagger Editor](https://editor.swagger.io/)  \n",
    "We download the openapi.json specification.\n",
    "\n",
    "Example curl query to get information of a Pet by PetID\n",
    "\n",
    "`\n",
    "curl -X 'GET' \\\n",
    "  'https://petstore3.swagger.io/api/v3/pet/1' \\\n",
    "  -H 'accept: application/json'\n",
    "`\n",
    "\n",
    "Example curl query to get information of a User by username\n",
    "`\n",
    "curl -X 'GET' \\\n",
    "  'https://petstore3.swagger.io/api/v3/user/user1' \\\n",
    "  -H 'accept: application/json'\n",
    "`  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "094ee113-45f1-4748-84b2-99c1ea31e52d",
   "metadata": {},
   "source": [
    "# Function Calling for REST API\n",
    "\n",
    "## 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_pet_info` and `retreive_user_info` to retrieve pet and user info given `petID` and `username`.\n",
    "- 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": 2,
   "id": "762814ec-9d1c-4f40-bd2c-4f1546f3e856",
   "metadata": {},
   "outputs": [],
   "source": [
    "def getPetById(petId: int) -> str:\n",
    "    try:\n",
    "        method = 'GET'\n",
    "        headers=None\n",
    "        data=None\n",
    "        url =  'https://petstore3.swagger.io/api/v3/pet/' + str(petId)\n",
    "        response = requests.request(method, url, headers=headers, data=data)\n",
    "        # Raise an exception if the response was unsuccessful\n",
    "        response.raise_for_status()\n",
    "        #response = make_api_call('GET', url + str(petId))\n",
    "        if response.ok :\n",
    "            json_response = response.json()\n",
    "            if petId == json_response['id']:\n",
    "                return json_response\n",
    "        return json.dumps({'error': 'Pet id not found.'})\n",
    "    except requests.exceptions.HTTPError as e:\n",
    "        if response.status_code == 404:\n",
    "            return json.dumps({'error': 'Pet id not found.'})\n",
    "        else:\n",
    "            return json.dumps({'error': 'Error with API.'})\n",
    "\n",
    "def getUserByName(username: str) -> str:\n",
    "    try:\n",
    "        url = 'https://petstore3.swagger.io/api/v3/user/' + username\n",
    "        response = requests.get(url)\n",
    "        # Raise an exception if the response was unsuccessful\n",
    "        response.raise_for_status()\n",
    "        if response.ok :\n",
    "            json_response = response.json()\n",
    "            if username == json_response['username']:\n",
    "                return json_response\n",
    "        return json.dumps({'error': 'Username id not found.'})\n",
    "    except requests.exceptions.HTTPError as e:\n",
    "        if response.status_code == 404:\n",
    "            return json.dumps({'error': 'Username not found.'})\n",
    "        else:\n",
    "            return json.dumps({'error': 'Error with API.'})\n",
    "\n",
    "names_to_functions = {\n",
    "  'getPetById': functools.partial(getPetById, petId=''),\n",
    "  'getUserByName': functools.partial(getUserByName, username='')  \n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3157604d-81a5-49d7-b393-2eae96d3a789",
   "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.\n",
    "- Tool Generator -\n",
    "parse open api spec for dynamic tool definition creation. Download openai.json from https://editor.swagger.io/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "070c0b61-3ea5-455b-95f7-78494db8b4f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_tools(objs, function_end_point)-> List[Tool]:\n",
    "    params = ['operationId', 'description',  'parameters']\n",
    "    parser = prance.ResolvingParser(function_end_point, backend='openapi-spec-validator')\n",
    "    spec = parser.specification\n",
    "    \n",
    "    user_tools = []\n",
    "    for obj in objs:\n",
    "        resource, field = obj\n",
    "        path = '/' + resource + '/{' + field + '}'\n",
    "        function_name=spec['paths'][path]['get'][params[0]]\n",
    "        function_description=spec['paths'][path]['get'][params[1]]\n",
    "        function_parameters=spec['paths'][path]['get'][params[2]]\n",
    "        func_parameters = {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            function_parameters[0]['name']: {\n",
    "                \"type\": function_parameters[0]['schema']['type'],\n",
    "                \"description\": function_parameters[0]['description']\n",
    "            }\n",
    "        },\n",
    "        \"required\": [function_parameters[0]['name']]\n",
    "    }\n",
    "        user_function= Function(name = function_name, description = function_description, parameters = func_parameters, )\n",
    "        user_tool = Tool(function = user_function)\n",
    "        user_tools.append(user_tool)\n",
    "    return user_tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34bdfba6-4784-40bd-84a6-eade3f0dcd44",
   "metadata": {},
   "source": [
    "### User query\n",
    "\n",
    "Suppose a user asks the following question: “What’s the status of my Pet 1?” 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": "d067b520-1b4d-441f-83c6-aebd2eeb1177",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_user_messages(queries: List[str]) -> List[UserMessage]:\n",
    "    user_messages=[]\n",
    "    for query in queries:\n",
    "        user_message = UserMessage(content=query)\n",
    "        user_messages.append(user_message)\n",
    "    return user_messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c7c64d5-24dd-4cc4-9db3-05d68b3a3c4c",
   "metadata": {},
   "source": [
    "For external ollama endpoint, set the environment variable \"OLLAMA_ENDPOINT\"\n",
    "\n",
    "export OLLAMA_ENDPOINT=\"YOUR-Ollama-IP:Port\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d6f585c6-7f72-4617-925b-930beafbaa71",
   "metadata": {},
   "outputs": [],
   "source": [
    "def execute_generator():\n",
    "    queries = [\"What's the status of my Pet 1?\", \"Find information of user user1?\" ,  \"What's the status of my Store Order 3?\"]\n",
    "    return_objs = [['pet','petId'], ['user', 'username'], ['store/order','orderId']]\n",
    "    function_end_point\n",
    "    user_messages=get_user_messages(queries)\n",
    "    user_tools = generate_tools(return_objs, function_end_point)\n",
    "\n",
    "    #create tokens for message and tools prompt\n",
    "    tokenizer = MistralTokenizer.v3()\n",
    "    completion_request = ChatCompletionRequest(tools=user_tools, messages=user_messages,)\n",
    "    tokenized = tokenizer.encode_chat_completion(completion_request)\n",
    "    _, text = tokenized.tokens, tokenized.text\n",
    "\n",
    "    ollama_endpoint_env = os.environ.get('OLLAMA_ENDPOINT')\n",
    "    model = \"mistral:7b\"\n",
    "    prompt = text \n",
    "\n",
    "    if ollama_endpoint_env is None:\n",
    "        ollama_endpoint_env = 'http://localhost:11434'\n",
    "    ollama_endpoint = ollama_endpoint_env +  \"/api/generate\"  # replace with localhost\n",
    "\n",
    "    response = requests.post(ollama_endpoint,\n",
    "                      json={\n",
    "                          'model': model,\n",
    "                          'prompt': prompt,\n",
    "                          'stream':False,\n",
    "                          'raw': True\n",
    "                      }, stream=False\n",
    "                      )\n",
    "    \n",
    "    response.raise_for_status()\n",
    "    result = response.json()\n",
    "\n",
    "    process_results(result, user_messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ae83be6-0637-45f0-ae00-981657b569be",
   "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 `getPetId` with the parameter `petId` set to 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "86474b15-1ebd-482b-9d1f-52bce07f99b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_results(result, messages):\n",
    "\n",
    "    result_format = result['response'].split(\"\\n\\n\")\n",
    "    result_tool_calls = result_format[0].replace(\"[TOOL_CALLS] \",\"\")\n",
    "\n",
    "    tool_calls = json.loads(result_tool_calls)\n",
    "    index = 0 \n",
    "    try:\n",
    "        for tool_call in tool_calls:\n",
    "            function_name = tool_call[\"name\"]\n",
    "            function_params = (tool_call[\"arguments\"]) \n",
    "            print(messages[index].content)\n",
    "            function_result = names_to_functions[function_name](**function_params)\n",
    "            print(function_result)\n",
    "            index = index + 1\n",
    "    except:\n",
    "        print(function_name + \" is not defined\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6f2ccfc3-3b32-40cb-9bb7-604664eca37d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "What's the status of my Pet 1?\n",
      "{'id': 1, 'category': {'id': -2, 'name': 'Кошка'}, 'name': 'Фей-Фей', 'photoUrls': ['string'], 'tags': [{'id': 3, 'name': 'Русская Голубая'}], 'status': 'pending'}\n",
      "Find information of user user1?\n",
      "{'id': 1, 'username': 'user1', 'firstName': 'first name 1', 'lastName': 'last name 1', 'email': 'email1@test.com', 'password': 'XXXXXXXXXXX', 'phone': '123-456-7890', 'userStatus': 1}\n",
      "What's the status of my Store Order 3?\n",
      "getOrderById is not defined\n"
     ]
    }
   ],
   "source": [
    "execute_generator()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38d64866-595d-4ea3-9313-5bcca5ba36a4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}