{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "2cba639c-0ed3-45ea-9688-06f5e956672f",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "# Building a Tool-calling Agent with LlamaIndex Workflow and MLflow\n",
    "\n",
    "Welcome to this interactive tutorial designed to introduce you to LlamaIndex Workflow and its integration with MLflow. This tutorial is structured as a notebook to provide a hands-on, practical learning experience with **Workflow**, LlamaIndex's novel approach to design LLM applications, and managing the development process with MLflow.\n",
    "\n",
    "![LlamaIndex Workflow Graph](/images/llms/llama-index/llama_index_workflow_graph.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "b399e48b-3509-4a24-b227-eb8a2365a1f8",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## What you will learn\n",
    "By the end of this tutorial you will have:\n",
    "\n",
    "* Created an MVP agentic application with tool calling functionality in a LlamaIndex Workflow.\n",
    "* Observed the agent actions with MLflow Tracing.\n",
    "* Logged that workflow to the MLflow Experiment.\n",
    "* Loaded the model back and performed inference.\n",
    "* Explored the MLflow UI to learn about logged artifacts."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "17a2dc91-310f-4920-9a5c-bfd087068c9f",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Installation\n",
    "\n",
    "MLflow's integration with LlamaIndex's Workflow API is available in MLflow >= 2.17.0 and LlamaIndex (core) >= 0.11.16. After installing the packages, you may need to restart the Python kernel to correctly load modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "dd6951b8-8c39-452b-a2bf-04d3f0b16756",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": "%pip install 'mlflow>=2.17.0' llama-index>=0.11.16 -qqqU\n# Workflow util is required for rendering Workflow as HTML\n%pip install llama-index-utils-workflow -qqqU"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "4481237e-342b-4614-ba03-eb9ab36f2a70",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Choose your favorite LLM\n",
    "\n",
    "By default, LlamaIndex uses OpenAI as the source for LLms and embedding models. If you are signing up with different LLM providers or using a local model, configure them for use by using the ``Settings`` object."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "70b5aa9b-80d2-47e6-9abc-b10721f81ccf",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Option 1: OpenAI (default)\n",
    "\n",
    "LlamaIndex by default uses OpenAI APIs for LLMs and embeddings models. To proceed with this setting, you just need to set the API key in the environment variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "664f349f-b60a-4f74-af6b-d5b514f3459f",
     "showTitle": true,
     "title": "%md #### Option 1: OpenAI (default)  LlamaIndex by default uses OpenAI APIs for LLMs and embeddings models. To proceed with this setting, you just need to set the API key in the environment variable."
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"<YOUR_OPENAI_API_KEY>\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "fdd9402d-4f22-4c34-a07d-7e254efbf5b8",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Option 2: Other Hosted LLMs\n",
    "\n",
    "If you want to use other hosted LLMs, \n",
    "\n",
    "1. Download the integration package for the model provider of your choice.\n",
    "2. Set up required environment variables as specified in the integration documentation.\n",
    "3. Instantiate the LLM instance and set it to the global `Settings` object.\n",
    "\n",
    "The following cells show an example for using Databricks hosted LLMs (Llama3.1 70B instruct)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "7e909334-24a0-4a3e-9d2e-afe0849bc8a8",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-databricks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "894b20a0-e611-48fc-aab6-d957502083fd",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"DATABRICKS_TOKEN\"] = \"<YOUR_DATABRICKS_API_TOKEN>\"\n",
    "os.environ[\"DATABRICKS_SERVING_ENDPOINT\"] = \"https://YOUR_DATABRICKS_HOST/serving-endpoints/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "93515bbd-f648-4096-8030-1239a9180594",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from llama_index.core import Settings\n",
    "from llama_index.llms.databricks import Databricks\n",
    "\n",
    "llm = Databricks(model=\"databricks-meta-llama-3-1-70b-instruct\")\n",
    "Settings.llm = llm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "0534d0cb-c15e-4048-b32a-ef89aaa798d0",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Option 3: Local LLM\n",
    "\n",
    "LlamaIndex also support locally hosted LLMs. Please refer to the [Starter Tutorial (Local Models)](https://docs.llamaindex.ai/en/stable/getting_started/starter_example_local/) for how to set them up."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "1b9bc91b-d9e2-42f2-92be-8bbbd73c2062",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Create an MLflow Experiemnt\n",
    "\n",
    "*Skip this step if you are running this tutorial on a Databricks Notebook. An MLflow experiment is automatically set up when you created any notebook.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "387ef167-c2f4-496f-9197-3d62cb20f4eb",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import mlflow\n",
    "\n",
    "mlflow.set_experiment(\"MLflow LlamaIndex Workflow Tutorial\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "c2506f73-6975-48ae-9f5e-0789362dd48b",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Define tools\n",
    "\n",
    "The agents access with various functions and resources via `tool` objects. In this example, we define the simplest possible math tools `add` and `multiply` based on Python functions. For a real-world application, you can create arbitrary tools such as vector search retrieval, web search, or even calling another agent as a tool. Please refer to the [Tools documentation](https://docs.llamaindex.ai/en/stable/module_guides/deploying/agents/tools/) for more details.\n",
    "\n",
    "\n",
    "Please ignore the `### [USE IN MODEL]` comment at the beginning of some cells like below. This will be used in later steps in this tutorial!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "1de5d2a5-945a-4931-8d09-b79a7481e52d",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# [USE IN MODEL]\n",
    "from llama_index.core.tools import FunctionTool\n",
    "\n",
    "\n",
    "def add(x: int, y: int) -> int:\n",
    "    \"\"\"Useful function to add two numbers.\"\"\"\n",
    "    return x + y\n",
    "\n",
    "\n",
    "def multiply(x: int, y: int) -> int:\n",
    "    \"\"\"Useful function to multiply two numbers.\"\"\"\n",
    "    return x * y\n",
    "\n",
    "\n",
    "tools = [\n",
    "    FunctionTool.from_defaults(add),\n",
    "    FunctionTool.from_defaults(multiply),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "5b530b12-a63b-41fc-a9f4-4987c4addbad",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Define Workflow\n",
    "\n",
    "\n",
    "#### Workflow Primer\n",
    "\n",
    "LlamaIndex Workflow is an event-driven orchestration framework. At its core, a workflow consists of two fundamental components: **Steps** and **Events**.\n",
    "\n",
    "\n",
    "- **Steps**: Units of execution within the workflow. Steps are defined as methods marked with the `@step` decorator in a class that implements the `Workflow` base class.\n",
    "- **Events**: Custom objects that trigger steps. Two special events, `StartEvent` and `EndEvent`, are reserved for dispatch at the beginning and end of the workflow.\n",
    "\n",
    "\n",
    "Each step specifies its input and output events through its function signature.\n",
    "\n",
    "```python\n",
    "@step\n",
    "async def my_step(self, event: StartEvent) -> FooEvent:\n",
    "    # This method triggers when a StartEvent is emitted at the workflow's start,\n",
    "    # and then dispatches a FooEvent.\n",
    "```\n",
    "\n",
    "Based on each step's signature and defined events, LlamaIndex automatically constructs the workflow's execution flow.\n",
    "\n",
    "You may notice that the `my_step` function is defined as an async function. LlamaIndex Workflow makes asynchronous operations a first-class feature, enabling easy parallel execution and scalable workflows.\n",
    "\n",
    "Another essential component of the workflow is the **Context** object. This global registry, accessible from any step, allows shared information to be defined without the need to pass it through multiple events."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "489a2ac1-887b-4c57-ae35-c3a7b4dfc1f8",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "#### Define a ReAct Agent as a Workflow\n",
    "\n",
    "The Workflow definition below models a ReAct Agent that utilizes the simple math tools we defined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "0d380e54-2df8-4ee2-b583-652297cab6ea",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# [USE IN MODEL]\n",
    "\n",
    "# Event definitions\n",
    "from llama_index.core.llms import ChatMessage, ChatResponse\n",
    "from llama_index.core.tools import ToolOutput, ToolSelection\n",
    "from llama_index.core.workflow import Event\n",
    "\n",
    "\n",
    "class PrepEvent(Event):\n",
    "    \"\"\"An event to handle new messages and prepare the chat history\"\"\"\n",
    "\n",
    "\n",
    "class LLMInputEvent(Event):\n",
    "    \"\"\"An event to prmopt the LLM with the react prompt (chat history)\"\"\"\n",
    "\n",
    "    input: list[ChatMessage]\n",
    "\n",
    "\n",
    "class LLMOutputEvent(Event):\n",
    "    \"\"\"An event represents LLM generation\"\"\"\n",
    "\n",
    "    response: ChatResponse\n",
    "\n",
    "\n",
    "class ToolCallEvent(Event):\n",
    "    \"\"\"An event to trigger tool calls, if any\"\"\"\n",
    "\n",
    "    tool_calls: list[ToolSelection]\n",
    "\n",
    "\n",
    "class ToolOutputEvent(Event):\n",
    "    \"\"\"An event to handle the results of tool calls, if any\"\"\"\n",
    "\n",
    "    output: ToolOutput"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "59a99eb5-e727-4917-97ea-b5af5a351794",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# [USE IN MODEL]\n",
    "\n",
    "# Workflow definition\n",
    "from llama_index.core import Settings\n",
    "from llama_index.core.agent.react import ReActChatFormatter, ReActOutputParser\n",
    "from llama_index.core.agent.react.types import ActionReasoningStep, ObservationReasoningStep\n",
    "from llama_index.core.memory import ChatMemoryBuffer\n",
    "from llama_index.core.workflow import (\n",
    "    Context,\n",
    "    StartEvent,\n",
    "    StopEvent,\n",
    "    Workflow,\n",
    "    step,\n",
    ")\n",
    "\n",
    "\n",
    "class ReActAgent(Workflow):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.tools = tools\n",
    "        # Store the chat history in memory so the agent can handle multiple interactions with users.\n",
    "        self.memory = ChatMemoryBuffer.from_defaults(llm=Settings.llm)\n",
    "\n",
    "    @step\n",
    "    async def new_user_msg(self, ctx: Context, ev: StartEvent) -> PrepEvent:\n",
    "        \"\"\"Start workflow with the new user messsage\"\"\"\n",
    "        # StartEvent carries whatever keys passed to the workflow's run() method as attributes.\n",
    "        user_input = ev.input\n",
    "        user_msg = ChatMessage(role=\"user\", content=user_input)\n",
    "        self.memory.put(user_msg)\n",
    "\n",
    "        # We store the executed reasoning steps in the context. Clear it at the start.\n",
    "        await ctx.set(\"steps\", [])\n",
    "\n",
    "        return PrepEvent()\n",
    "\n",
    "    @step\n",
    "    async def prepare_llm_prompt(self, ctx: Context, ev: PrepEvent) -> LLMInputEvent:\n",
    "        \"\"\"Prepares the react prompt, using the chat history, tools, and current reasoning (if any)\"\"\"\n",
    "        steps = await ctx.get(\"steps\", default=[])\n",
    "        chat_history = self.memory.get()\n",
    "\n",
    "        # Construct an LLM from the chat history, tools, and current reasoning, using the\n",
    "        # built-in prompt template.\n",
    "        llm_input = ReActChatFormatter().format(self.tools, chat_history, current_reasoning=steps)\n",
    "        return LLMInputEvent(input=llm_input)\n",
    "\n",
    "    @step\n",
    "    async def invoke_llm(self, ev: LLMInputEvent) -> LLMOutputEvent:\n",
    "        \"\"\"Call the LLM with the react prompt\"\"\"\n",
    "        response = await Settings.llm.achat(ev.input)\n",
    "        return LLMOutputEvent(response=response)\n",
    "\n",
    "    @step\n",
    "    async def handle_llm_response(\n",
    "        self, ctx: Context, ev: LLMOutputEvent\n",
    "    ) -> ToolCallEvent | PrepEvent | StopEvent:\n",
    "        \"\"\"\n",
    "        Parse the LLM response to extract any tool calls requested.\n",
    "        If theere is no tool call, we can stop and emit a StopEvent. Otherwise, we emit a ToolCallEvent to handle tool calls.\n",
    "        \"\"\"\n",
    "        try:\n",
    "            step = ReActOutputParser().parse(ev.response.message.content)\n",
    "            (await ctx.get(\"steps\", default=[])).append(step)\n",
    "\n",
    "            if step.is_done:\n",
    "                # No additional tool call is required. Ending the workflow by emitting StopEvent.\n",
    "                return StopEvent(result=step.response)\n",
    "            elif isinstance(step, ActionReasoningStep):\n",
    "                # Tool calls are returned from LLM, trigger the tool call event.\n",
    "                return ToolCallEvent(\n",
    "                    tool_calls=[\n",
    "                        ToolSelection(\n",
    "                            tool_id=\"fake\",\n",
    "                            tool_name=step.action,\n",
    "                            tool_kwargs=step.action_input,\n",
    "                        )\n",
    "                    ]\n",
    "                )\n",
    "        except Exception as e:\n",
    "            error_step = ObservationReasoningStep(\n",
    "                observation=f\"There was an error in parsing my reasoning: {e}\"\n",
    "            )\n",
    "            (await ctx.get(\"steps\", default=[])).append(error_step)\n",
    "\n",
    "        # if no tool calls or final response, iterate again\n",
    "        return PrepEvent()\n",
    "\n",
    "    @step\n",
    "    async def handle_tool_calls(self, ctx: Context, ev: ToolCallEvent) -> PrepEvent:\n",
    "        \"\"\"\n",
    "        Safely calls tools with error handling, adding the tool outputs to the current reasoning. Then, by emitting a PrepEvent, we loop around for another round of ReAct prompting and parsing.\n",
    "        \"\"\"\n",
    "        tool_calls = ev.tool_calls\n",
    "        tools_by_name = {tool.metadata.get_name(): tool for tool in self.tools}\n",
    "\n",
    "        # call tools -- safely!\n",
    "        for tool_call in tool_calls:\n",
    "            if tool := tools_by_name.get(tool_call.tool_name):\n",
    "                try:\n",
    "                    tool_output = tool(**tool_call.tool_kwargs)\n",
    "                    step = ObservationReasoningStep(observation=tool_output.content)\n",
    "                except Exception as e:\n",
    "                    step = ObservationReasoningStep(\n",
    "                        observation=f\"Error calling tool {tool.metadata.get_name()}: {e}\"\n",
    "                    )\n",
    "            else:\n",
    "                step = ObservationReasoningStep(\n",
    "                    observation=f\"Tool {tool_call.tool_name} does not exist\"\n",
    "                )\n",
    "            (await ctx.get(\"steps\", default=[])).append(step)\n",
    "\n",
    "        # prep the next iteration\n",
    "        return PrepEvent()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "5978d41e-e4bf-40a2-af8d-b078ef6fa928",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "\n",
    "#### Check the Workflow Visually\n",
    "\n",
    "Before instantiating the agent object, let's pause and validate if the workflow is constructed as we expect.\n",
    "\n",
    "To check that, we can render the graphical representation of the workflow by using the `draw_all_possible_flows` utility function.\n",
    "\n",
    "(Note: If the rendered HTML is blank, it might be due to the safety feature in Jupyter. In that case, you can trust the notebook by `!jupyter trust llama_index_workflow_tutorial.ipynb`. See [Jupyter documentation](https://jupyter-notebook.readthedocs.io/en/latest/notebook.html#signing-notebooks) for more details.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "bf6ecb27-c219-485f-a13a-8eb27175f950",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "from llama_index.utils.workflow import draw_all_possible_flows\n",
    "\n",
    "draw_all_possible_flows(ReActAgent, filename=\"workflow.html\")\n",
    "\n",
    "with open(\"workflow.html\") as file:\n",
    "    html_content = file.read()\n",
    "HTML(html_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "d29e2728-17c9-4401-b24a-ca7a0bb87e6e",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# [USE IN MODEL]\n",
    "agent = ReActAgent(timeout=180)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "9430c2b4-4cb4-480a-9c12-c51e3d9425be",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Run the Workflow (with Trace)\n",
    "\n",
    "Now your workflow is all set! But before running that, let's not forget to turn on [MLflow Tracing](https://mlflow.org/docs/latest/llms/tracing/index.html), so you get observability into each step during the agent run, and record it for the review later.\n",
    "\n",
    "Mlflow supports automatic tracing for LlamaIndex Workflow. To enable it, you just need to call the `mlflow.llama_index.autolog()` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "0d145fb2-de22-4ece-b07c-134c38bcfd9f",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import mlflow\n",
    "\n",
    "mlflow.llama_index.autolog()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "8eb90e36-a9cc-487a-9e8c-132875a4e3e3",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The result of (123 + 456) * 789 is 579,027.'"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Run the workflow\n",
    "await agent.run(input=\"What is (123 + 456) * 789?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "ca4d2bdb-670f-4778-ba45-1410391263ee",
     "showTitle": false,
     "title": ""
    },
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Review the Trace\n",
    "\n",
    "The generated traces are automatically recorded to your MLflow Experiment.\n",
    "\n",
    "1. Open a terminal, run `mlflow server --port 5000` within the current directory (and keep it running).\n",
    "2. Navigate to `http://127.0.0.1:5000` in your browser.\n",
    "3. Open the experiment \"MLflow LlamaIndex Workflow Tutorial\".\n",
    "4. Navigate to the \"Trace\" tab below the experiment name header.\n",
    "\n",
    "![LlamaIndex Workflow Trace](/images/llms/llama-index/llama_index_workflow_trace.png)\n",
    "\n",
    "The Trace records the individual steps inside the workflow execution with its inputs, outputs, and additional metadata such as latency. Let's do a quick exercise to find the following information on the Trace UI.\n",
    "\n",
    "<html>\n",
    "    <details>\n",
    "        <summary>1. Token count used for the first LLM invocation</summary>\n",
    "        <p>You can find token counts for LLm call in the <strong>Attribtues</strong> section of the LLM call span, inside the <code>usage</code> field.</p>\n",
    "    </details>\n",
    "    <details>\n",
    "        <summary>2. Input numbers for the \"add\" tool call.</summary>\n",
    "        <p>You can find input numbers <code>x=123</code> and <code>y=456</code> in the <code>Inputs</code> field of the span named <code>FunctionTool.call</code>. That span is located under the <code>ReActAgent.handle_tool_calls</code> step span.</p>\n",
    "    </details>\n",
    "</html>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Log the Workflow to an MLflow Experiment\n",
    "\n",
    "Now that you've built your first ReAct Agent using LlamaIndex Workflow, it's essential to iteratively refine and optimize for better performance. An **MLflow Experiment** is the ideal place to record and manage these improvements"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Prepare a Model script\n",
    "\n",
    "MLflow supports logging LlamaIndex workflows using the **Models from Code** method, allowing models to be defined and logged directly from a standalone Python script. This approach bypasses the need for risky and brittle serialization methods like `pickle`, using code as the single source of truth for the model definition. Combined with MLflow's environment-freezing capability, this provides a reliable way to persist the model.\n",
    "\n",
    "For more details, see the [MLflow documentation](https://mlflow.org/docs/latest/models.html#models-from-code)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You could manually create a separate Python file by copying the code from this notebook. However, for convenience, we define a utility function to generate a model script automatically from this notebook's content in one step. Running the cell below will create this script in the current directory, ready for MLflow logging."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "fdc40724-5d70-4425-96e8-79b6231556dc",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model code saved to react_agent.py\n"
     ]
    }
   ],
   "source": [
    "def generate_model_script(output_path, notebook_path=\"llama_index_workflow_tutorial.ipynb\"):\n",
    "    \"\"\"\n",
    "    A utility function to generate a ready-to-log .py script that\n",
    "    contains necessary library imports and model definitions.\n",
    "\n",
    "    Args:\n",
    "       output_path: The path to write the .py file to.\n",
    "       notebook_path: The path to the tutorial notebook.\n",
    "    \"\"\"\n",
    "    import nbformat\n",
    "\n",
    "    with open(notebook_path, encoding=\"utf-8\") as f:\n",
    "        notebook = nbformat.read(f, as_version=4)\n",
    "\n",
    "    # Filter cells that are code cells and contain the specified marker\n",
    "    merged_code = (\n",
    "        \"\\n\\n\".join([\n",
    "            cell.source\n",
    "            for cell in notebook.cells\n",
    "            if cell.cell_type == \"code\" and cell.source.startswith(\"# [USE IN MODEL]\")\n",
    "        ])\n",
    "        + \"\\n\\nimport mlflow\\n\\nmlflow.models.set_model(agent)\"\n",
    "    )\n",
    "\n",
    "    # Write to the output .py file\n",
    "    with open(output_path, \"w\", encoding=\"utf-8\") as f:\n",
    "        f.write(merged_code)\n",
    "\n",
    "    print(f\"Model code saved to {output_path}\")\n",
    "\n",
    "\n",
    "# Pass `notebook_path` argument if you changed the notebook name\n",
    "generate_model_script(output_path=\"react_agent.py\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Logging the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "bf6c1ac1-3a8d-4b6e-8799-52240707fc1f",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import mlflow\n",
    "\n",
    "with mlflow.start_run(run_name=\"react-agent-workflow\"):\n",
    "    model_info = mlflow.llama_index.log_model(\n",
    "        \"react_agent.py\",\n",
    "        name=\"model\",\n",
    "        # Logging with an input example help MLflow to record dependency and signature information accurately.\n",
    "        input_example={\"input\": \"What is (123 + 456) * 789?\"},\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "77cb3bcd-0b40-41dd-9af1-19636dac0dfa",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Explore the MLflow UI\n",
    "\n",
    "Let's open the MLflow UI again to see which information is being tracked in the experiment.\n",
    "\n",
    "1. Access the MLflow UI like we did for reviewing traces.\n",
    "2. Open the experiment \"MLflow LlamaIndex Workflow Tutorial\".\n",
    "3. The `Runs` tab in the experiment should contain a run named \"react-agent-workflow\". Open it.\n",
    "4. On the run page, navigate to the `\"Artifacts\"` tab.\n",
    "\n",
    "The artifacts tab shows various files saved by MLflow in the Run. See the below image and open the annotated files to check which information is stored in each file.\n",
    "\n",
    "![LlamaIndex Workflow Artifacts](/images/llms/llama-index/llama_index_workflow_artifacts.png)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "74347217-2b5d-45d0-ba92-1cbb6a3cf136",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Load the Model Back for Inference\n",
    "\n",
    "With all necessary metadata logged to MLflow, you can load the model in a different notebook or deploy it for inference without concerns about environment inconsistencies. Let's do a quick exercise to demonstrate how this helps in reproducing experiment results.\n",
    "\n",
    "To simulate a different environment, we'll remove the `llm` configuration from the global `Settings` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'text'"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from llama_index.core.llms import MockLLM\n",
    "\n",
    "Settings.llm = MockLLM(max_tokens=1)\n",
    "\n",
    "await agent.run(input=\"What is (123 + 456) * 789?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "409c9b95-64b0-41e5-acc3-db64b366cafb",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "Since the dummy LLM is configured, the workflow could not generate the correct output but just returns \"text\".\n",
    "\n",
    "Now try loading the model back from the MLflow Experiment by calling `mlflow.llama_index.load_model()` API and run the workflow again. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "57fbf242-a067-4768-acfd-ab876ed40a19",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2f9d4e32ea5748dd8708942f3194f755",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading artifacts:   0%|          | 0/12 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "'(123 + 456) * 789 = 456831'"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loaded_model = mlflow.llama_index.load_model(\"runs:/f8e0a0d2dd5546d5ac93ce126358c444/model\")\n",
    "await loaded_model.run(input=\"What is (123 + 456) * 789?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This time, the output is computed correctly, because MLflow automatically restores the original LLM setting at the time of logging."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning More\n",
    "\n",
    "Congratulations! 🎉 You've successfully learned how to build a tool-calling agent using LlamaIndex Workflow and MLflow.\n",
    "\n",
    "Continue your journey with these advanced resources:\n",
    "\n",
    "* **Improve Workflow Quality**: Evaluate your workflow to enhance performance with [MLflow LLM Evaluation](https://mlflow.org/docs/latest/llms/llm-evaluate/index.html).\n",
    "* **Deploy Your Model**: Deploy your MLflow model to a serving endpoint with [MLflow Deployment](https://mlflow.org/docs/latest/deployment/index.html).\n",
    "* **Explore More Examples**: Discover additional examples of LlamaIndex Workflow in the [official documentation](https://docs.llamaindex.ai/en/stable/module_guides/workflow/#examples).\n"
   ]
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "dashboards": [],
   "environmentMetadata": {
    "base_environment": "",
    "client": "1"
   },
   "language": "python",
   "notebookMetadata": {
    "pythonIndentUnit": 2
   },
   "notebookName": "llama_index_workflow_tutorial",
   "widgets": {}
  },
  "kernelspec": {
   "display_name": "llama",
   "language": "python",
   "name": "llama"
  },
  "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}