{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build a tool-calling model with mlflow.pyfunc.ChatModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Welcome to the notebook tutorial on building a simple tool calling model using the [mlflow.pyfunc.ChatModel](https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#mlflow.pyfunc.ChatModel) wrapper. ChatModel is a subclass of MLflow's highly customizable [PythonModel](https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#mlflow.pyfunc.PythonModel), which was specifically designed to make creating GenAI workflows easier.\n",
    "\n",
    "Briefly, here are some of the benefits of using ChatModel:\n",
    "\n",
    "1. No need to define a complex signature! Chat models often accept complex inputs with many levels of nesting, and this can be cumbersome to define yourself.\n",
    "2. Support for JSON / dict inputs (no need to wrap inputs or convert to Pandas DataFrame)\n",
    "3. Includes the use of Dataclasses for defining expected inputs / outputs for a simplified development experience\n",
    "\n",
    "For a more in-depth exploration of ChatModel, please check out the [detailed guide](https://mlflow.org/docs/latest/llms/chat-model-guide/index.html).\n",
    "\n",
    "In this tutorial, we'll be building a simple OpenAI wrapper that makes use of the tool calling support (released in MLflow 2.17.0)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Environment setup\n",
    "\n",
    "First, let's set up the environment. We'll need the OpenAI Python SDK, as well as MLflow >= 2.17.0. We'll also need to set our OpenAI API key in order to use the SDK."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "%pip install 'mlflow>=2.17.0' 'openai>=1.0' -qq"
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass(\"Enter your OpenAI API key: \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1: Creating the tool definition\n",
    "\n",
    "Let's begin to define our model! As mentioned in the introduction, we'll be subclassing `mlflow.pyfunc.ChatModel`. For this example, we'll build a toy model that uses a tool to retrieve the weather for a given city.\n",
    "\n",
    "The first step is to create a tool definition that we can pass to OpenAI. We do this by using [mlflow.types.llm.FunctionToolDefinition](https://mlflow.org/docs/latest/python_api/mlflow.types.html#mlflow.types.llm.FunctionToolDefinition) to describe the parameters that our tool accepts. The format of this dataclass is aligned with the OpenAI spec:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlflow\n",
    "from mlflow.types.llm import (\n",
    "    FunctionToolDefinition,\n",
    "    ParamProperty,\n",
    "    ToolParamsSchema,\n",
    ")\n",
    "\n",
    "\n",
    "class WeatherModel(mlflow.pyfunc.ChatModel):\n",
    "    def __init__(self):\n",
    "        # a sample tool definition. we use the `FunctionToolDefinition`\n",
    "        # class to describe the name and expected params for the tool.\n",
    "        # for this example, we're defining a simple tool that returns\n",
    "        # the weather for a given city.\n",
    "        weather_tool = FunctionToolDefinition(\n",
    "            name=\"get_weather\",\n",
    "            description=\"Get weather information\",\n",
    "            parameters=ToolParamsSchema({\n",
    "                \"city\": ParamProperty(\n",
    "                    type=\"string\",\n",
    "                    description=\"City name to get weather information for\",\n",
    "                ),\n",
    "            }),\n",
    "            # make sure to call `to_tool_definition()` to convert the `FunctionToolDefinition`\n",
    "            # to a `ToolDefinition` object. this step is necessary to normalize the data format,\n",
    "            # as multiple types of tools (besides just functions) might be available in the future.\n",
    "        ).to_tool_definition()\n",
    "\n",
    "        # OpenAI expects tools to be provided as a list of dictionaries\n",
    "        self.tools = [weather_tool.to_dict()]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Implementing the tool\n",
    "\n",
    "Now that we have a definition for the tool, we need to actually implement it. For the purposes of this tutorial, we're just going to mock a response, but the implementation can be arbitrary—you might make an API call to an actual weather service, for example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WeatherModel(mlflow.pyfunc.ChatModel):\n",
    "    def __init__(self):\n",
    "        weather_tool = FunctionToolDefinition(\n",
    "            name=\"get_weather\",\n",
    "            description=\"Get weather information\",\n",
    "            parameters=ToolParamsSchema({\n",
    "                \"city\": ParamProperty(\n",
    "                    type=\"string\",\n",
    "                    description=\"City name to get weather information for\",\n",
    "                ),\n",
    "            }),\n",
    "        ).to_tool_definition()\n",
    "\n",
    "        self.tools = [weather_tool.to_dict()]\n",
    "\n",
    "        def get_weather(self, city: str) -> str:\n",
    "            # in a real-world scenario, the implementation might be more complex\n",
    "            return f\"It's sunny in {city}, with a temperature of 20C\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Implementing the `predict` method\n",
    "\n",
    "The next thing we need to do is define a `predict()` function that accepts the following arguments:\n",
    "\n",
    "1. `context`: [PythonModelContext](https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#mlflow.pyfunc.PythonModelContext) (not used in this tutorial)\n",
    "2. `messages`: List\\[[ChatMessage](https://mlflow.org/docs/latest/python_api/mlflow.types.html#mlflow.types.llm.ChatMessage)\\]. This is the chat input that the model uses for generation.\n",
    "3. `params`: [ChatParams](https://mlflow.org/docs/latest/python_api/mlflow.types.html#mlflow.types.llm.ChatParams). These are commonly used params used to configure the chat model, e.g. `temperature`, `max_tokens`, etc. This is where the tool specifications can be found.\n",
    "\n",
    "This is the function that will ultimately be called during inference.\n",
    "\n",
    "For the implementation, we'll simply forward the user's input to OpenAI, and provide the `get_weather` tool as an option for the LLM to use if it chooses to do so. If we receive a tool call request, we'll call the `get_weather()` function and return the response back to OpenAI. We'll need to use what we've defined in the previous two steps in order to do this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "from openai import OpenAI\n",
    "\n",
    "import mlflow\n",
    "from mlflow.types.llm import (\n",
    "    ChatMessage,\n",
    "    ChatParams,\n",
    "    ChatResponse,\n",
    ")\n",
    "\n",
    "\n",
    "class WeatherModel(mlflow.pyfunc.ChatModel):\n",
    "    def __init__(self):\n",
    "        weather_tool = FunctionToolDefinition(\n",
    "            name=\"get_weather\",\n",
    "            description=\"Get weather information\",\n",
    "            parameters=ToolParamsSchema({\n",
    "                \"city\": ParamProperty(\n",
    "                    type=\"string\",\n",
    "                    description=\"City name to get weather information for\",\n",
    "                ),\n",
    "            }),\n",
    "        ).to_tool_definition()\n",
    "\n",
    "        self.tools = [weather_tool.to_dict()]\n",
    "\n",
    "    def get_weather(self, city: str) -> str:\n",
    "        return \"It's sunny in {}, with a temperature of 20C\".format(city)\n",
    "\n",
    "    # the core method that needs to be implemented. this function\n",
    "    # will be called every time a user sends messages to our model\n",
    "    def predict(self, context, messages: list[ChatMessage], params: ChatParams):\n",
    "        # instantiate the OpenAI client\n",
    "        client = OpenAI()\n",
    "\n",
    "        # convert the messages to a format that the OpenAI API expects\n",
    "        messages = [m.to_dict() for m in messages]\n",
    "\n",
    "        # call the OpenAI API\n",
    "        response = client.chat.completions.create(\n",
    "            model=\"gpt-4o-mini\",\n",
    "            messages=messages,\n",
    "            # pass the tools in the request\n",
    "            tools=self.tools,\n",
    "        )\n",
    "\n",
    "        # if OpenAI returns a tool_calling response, then we call\n",
    "        # our tool. otherwise, we just return the response as is\n",
    "        if tool_calls := response.choices[0].message.tool_calls:\n",
    "            print(\"Received a tool call, calling the weather tool...\")\n",
    "\n",
    "            # for this example, we only provide the model with one tool,\n",
    "            # so we can assume the tool call is for the weather tool. if\n",
    "            # we had more, we'd need to check the name of the tool that\n",
    "            # was called\n",
    "            city = json.loads(tool_calls[0].function.arguments)[\"city\"]\n",
    "            tool_call_id = tool_calls[0].id\n",
    "\n",
    "            # call the tool and construct a new chat message\n",
    "            tool_response = ChatMessage(\n",
    "                role=\"tool\", content=self.get_weather(city), tool_call_id=tool_call_id\n",
    "            ).to_dict()\n",
    "\n",
    "            # send another request to the API, making sure to append\n",
    "            # the assistant's tool call along with the tool response.\n",
    "            messages.append(response.choices[0].message)\n",
    "            messages.append(tool_response)\n",
    "            response = client.chat.completions.create(\n",
    "                model=\"gpt-4o-mini\",\n",
    "                messages=messages,\n",
    "                tools=self.tools,\n",
    "            )\n",
    "\n",
    "        # return the result as a ChatResponse, as this\n",
    "        # is the expected output of the predict method\n",
    "        return ChatResponse.from_dict(response.to_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4 (optional, but recommended): Enable tracing for the model\n",
    "\n",
    "This step is optional, but highly recommended to improve observability in your app. We'll be using [MLflow Tracing](https://mlflow.org/docs/latest/llms/tracing/index.html) to log the inputs and outputs of our model's internal functions, so we can easily debug when things go wrong. Agent-style tool calling models can make many layers of function calls during the lifespan of a single request, so tracing is invaluable in helping us understand what's going on at each step.\n",
    "\n",
    "Integrating tracing is easy, we simply decorate the functions we're interested in (`get_weather()` and `predict()`) with `@mlflow.trace`! MLflow Tracing also has integrations with many popular GenAI frameworks, such as LangChain, OpenAI, LlamaIndex, and more. For the full list, check out this [documentation page](https://mlflow.org/docs/latest/llms/tracing/index.html#automatic-tracing). In this tutorial, we're using the OpenAI SDK to make API calls, so we can enable tracing for this by calling `mlflow.openai.autolog()`.\n",
    "\n",
    "To view the traces in the UI, run `mlflow server` in a separate terminal shell, and navigate to the `Traces` tab after using the model for inference below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlflow.entities.span import (\n",
    "    SpanType,\n",
    ")\n",
    "\n",
    "# automatically trace OpenAI SDK calls\n",
    "mlflow.openai.autolog()\n",
    "\n",
    "\n",
    "class WeatherModel(mlflow.pyfunc.ChatModel):\n",
    "    def __init__(self):\n",
    "        weather_tool = FunctionToolDefinition(\n",
    "            name=\"get_weather\",\n",
    "            description=\"Get weather information\",\n",
    "            parameters=ToolParamsSchema({\n",
    "                \"city\": ParamProperty(\n",
    "                    type=\"string\",\n",
    "                    description=\"City name to get weather information for\",\n",
    "                ),\n",
    "            }),\n",
    "        ).to_tool_definition()\n",
    "\n",
    "        self.tools = [weather_tool.to_dict()]\n",
    "\n",
    "    @mlflow.trace(span_type=SpanType.TOOL)\n",
    "    def get_weather(self, city: str) -> str:\n",
    "        return \"It's sunny in {}, with a temperature of 20C\".format(city)\n",
    "\n",
    "    @mlflow.trace(span_type=SpanType.AGENT)\n",
    "    def predict(self, context, messages: list[ChatMessage], params: ChatParams):\n",
    "        client = OpenAI()\n",
    "\n",
    "        messages = [m.to_dict() for m in messages]\n",
    "\n",
    "        response = client.chat.completions.create(\n",
    "            model=\"gpt-4o-mini\",\n",
    "            messages=messages,\n",
    "            tools=self.tools,\n",
    "        )\n",
    "\n",
    "        if tool_calls := response.choices[0].message.tool_calls:\n",
    "            print(\"Received a tool call, calling the weather tool...\")\n",
    "\n",
    "            city = json.loads(tool_calls[0].function.arguments)[\"city\"]\n",
    "            tool_call_id = tool_calls[0].id\n",
    "\n",
    "            tool_response = ChatMessage(\n",
    "                role=\"tool\", content=self.get_weather(city), tool_call_id=tool_call_id\n",
    "            ).to_dict()\n",
    "\n",
    "            messages.append(response.choices[0].message)\n",
    "            messages.append(tool_response)\n",
    "            response = client.chat.completions.create(\n",
    "                model=\"gpt-4o-mini\",\n",
    "                messages=messages,\n",
    "                tools=self.tools,\n",
    "            )\n",
    "\n",
    "        return ChatResponse.from_dict(response.to_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 5: Logging the model\n",
    "\n",
    "Finally, we need to log the model. This saves the model as an artifact in MLflow Tracking, and allows us to load and serve it later on.\n",
    "\n",
    "(Note: this is a fundamental pattern in MLflow. To learn more, check out the [Quickstart guide](https://mlflow.org/docs/latest/getting-started/intro-quickstart/index.html)!)\n",
    "\n",
    "In order to do this, we need to do a few things:\n",
    "\n",
    "1. Define an input example to inform users about the input we expect\n",
    "2. Instantiate the model\n",
    "3. Call `mlflow.pyfunc.log_model()` with the above as arguments\n",
    "\n",
    "Take note of the Model URI printed out at the end of the cell—we'll need it when serving the model later!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/10/29 09:30:14 INFO mlflow.pyfunc: Predicting on input example to validate output\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Received a tool call, calling the weather tool...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7efce7f1f8e64673ab381052e5b02499",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading artifacts:   0%|          | 0/7 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Received a tool call, calling the weather tool...\n",
      "Successfully logged the model at the following URI:  runs:/8051850efa194a3b8b2450c4c9f4d42f/weather-model\n"
     ]
    }
   ],
   "source": [
    "# messages to use as input examples\n",
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": \"Please use the provided tools to answer user queries.\"},\n",
    "    {\"role\": \"user\", \"content\": \"What's the weather in Singapore?\"},\n",
    "]\n",
    "\n",
    "input_example = {\n",
    "    \"messages\": messages,\n",
    "}\n",
    "\n",
    "# instantiate the model\n",
    "model = WeatherModel()\n",
    "\n",
    "# log the model\n",
    "with mlflow.start_run():\n",
    "    model_info = mlflow.pyfunc.log_model(\n",
    "        name=\"weather-model\",\n",
    "        python_model=model,\n",
    "        input_example=input_example,\n",
    "    )\n",
    "\n",
    "    print(\"Successfully logged the model at the following URI: \", model_info.model_uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using the model for inference\n",
    "\n",
    "Now that the model is logged, our work is more or less done! In order to use the model for inference, let's load it back using `mlflow.pyfunc.load_model()`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/10/29 09:30:27 WARNING mlflow.tracing.processor.mlflow: Creating a trace within the default experiment with id '0'. It is strongly recommended to not use the default experiment to log traces due to ambiguous search results and probable performance issues over time due to directory table listing performance degradation with high volumes of directories within a specific path. To avoid performance and disambiguation issues, set the experiment for your environment using `mlflow.set_experiment()` API.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Received a tool call, calling the weather tool...\n",
      "The weather in Singapore is sunny, with a temperature of 20°C.\n",
      "Received a tool call, calling the weather tool...\n",
      "The weather in San Francisco is sunny, with a temperature of 20°C.\n"
     ]
    }
   ],
   "source": [
    "import mlflow\n",
    "\n",
    "# Load the previously logged ChatModel\n",
    "tool_model = mlflow.pyfunc.load_model(model_info.model_uri)\n",
    "\n",
    "system_prompt = {\n",
    "    \"role\": \"system\",\n",
    "    \"content\": \"Please use the provided tools to answer user queries.\",\n",
    "}\n",
    "\n",
    "messages = [\n",
    "    system_prompt,\n",
    "    {\"role\": \"user\", \"content\": \"What's the weather in Singapore?\"},\n",
    "]\n",
    "\n",
    "# Call the model's predict method\n",
    "response = tool_model.predict({\"messages\": messages})\n",
    "print(response[\"choices\"][0][\"message\"][\"content\"])\n",
    "\n",
    "messages = [\n",
    "    system_prompt,\n",
    "    {\"role\": \"user\", \"content\": \"What's the weather in San Francisco?\"},\n",
    "]\n",
    "\n",
    "# Generating another response\n",
    "response = tool_model.predict({\"messages\": messages})\n",
    "print(response[\"choices\"][0][\"message\"][\"content\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Serving the model\n",
    "\n",
    "MLflow also allows you to serve models, using the `mlflow models serve` CLI tool. In another terminal shell, run the following from the same folder as this notebook:\n",
    "\n",
    "```sh\n",
    "$ export OPENAI_API_KEY=<YOUR OPENAI API KEY>\n",
    "$ mlflow models serve -m <MODEL_URI>\n",
    "```\n",
    "\n",
    "This will start serving the model on `http://127.0.0.1:5000`, and the model can be queried via POST request to the `/invocations` route."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'choices': [{'index': 0,\n",
       "   'message': {'role': 'assistant',\n",
       "    'content': 'The weather in Tokyo is sunny, with a temperature of 20°C.'},\n",
       "   'finish_reason': 'stop'}],\n",
       " 'usage': {'prompt_tokens': 100, 'completion_tokens': 16, 'total_tokens': 116},\n",
       " 'id': 'chatcmpl-ANVOhWssEiyYNFwrBPxp1gmQvZKsy',\n",
       " 'model': 'gpt-4o-mini-2024-07-18',\n",
       " 'object': 'chat.completion',\n",
       " 'created': 1730165599}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import requests\n",
    "\n",
    "messages = [\n",
    "    system_prompt,\n",
    "    {\"role\": \"user\", \"content\": \"What's the weather in Tokyo?\"},\n",
    "]\n",
    "\n",
    "response = requests.post(\"http://127.0.0.1:5000/invocations\", json={\"messages\": messages})\n",
    "response.raise_for_status()\n",
    "response.json()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conclusion\n",
    "\n",
    "In this tutorial, we covered how to use MLflow's `ChatModel` class to create a convenient OpenAI wrapper that supports tool calling. Though the use-case was simple, the concepts covered here can be easily extended to support more complex functionality.\n",
    "\n",
    "If you're looking to dive deeper into building quality GenAI apps, you might be also be interested in checking out [MLflow Tracing](https://mlflow.org/docs/latest/llms/tracing/index.html), an observability tool you can use to trace the execution of arbitrary functions (such as your tool calls, for example)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dev",
   "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}