notebooks/llmu/multi_step_tool_use.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-Step Tool Use"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/cohere-ai/notebooks/blob/main/notebooks/llmu/multi_step_tool_use.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Read the accompanying [article here](https://cohere.ai/blog/multi-step-tool-use/).*\n",
    "\n",
    "Multi-step tool use happens when more than one tool is required, and the output of one tool is needed as the input to another tool.\n",
    "\n",
    "In other words, tool-calling needs to happen in a sequence.\n",
    "\n",
    "This is especially true in complex workflows where tasks have to be broken down into a sequence of steps.\n",
    "\n",
    "In this notebook, we’ll build a calendar assistant that can check a user’s calendar and create new appointments based on available times. This use case requires tool calls to happen in sequence, i.e., multi-step.\n",
    "\n",
    "It is also an example of using tools for performing not just a read operation (retrieving information) but also a write operation (taking actions).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let’s install the Cohere Python SDK and set up the Cohere client."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install cohere -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cohere\n",
    "import json\n",
    "co = cohere.Client(\"COHERE_API_KEY\") # Get your API key: https://dashboard.cohere.com/api-keys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s now create two tools:\n",
    "- A function called `list_calendar_events` to list the existing calendar events based on a given date. For simplicity, we are not going to connect to an actual calendar. So we’ll include a mock events record, which is a simple list of events in a day.\n",
    "- A function called `create_calendar_event` to create a new calendar event based on the provided date, time, and duration. Again, to keep things simple, we are not going to make actual changes to a database, but instead will just return a simple success message for illustration purposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the tools\n",
    "\n",
    "def list_calendar_events(date: str):\n",
    "  events = [{\"start\": \"8:00\", \"end\": \"8:59\"}, {\"start\": \"9:00\", \"end\": \"9:59\"}, {\"start\": \"11:00\", \"end\": \"11:59\"},{\"start\": \"12:00\", \"end\": \"12:59\"}]\n",
    "\n",
    "  return {\n",
    "        \"existing_events\": events\n",
    "    }\n",
    "\n",
    "def create_calendar_event(date: str, time: str, duration: int):\n",
    "  \n",
    "  return {\n",
    "        \"is_success\": True,\n",
    "        \"message\": f\"Created a {duration} hour long event at {time} on {date}\"\n",
    "    }\n",
    "\n",
    "functions_map = {\n",
    "    \"list_calendar_events\": list_calendar_events,\n",
    "    \"create_calendar_event\": create_calendar_event\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we define the tool schema for the two tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [\n",
    "    {\n",
    "      \"name\": \"list_calendar_events\",\n",
    "      \"description\": \"Returns a list of existing events for the specified date, including the start time and end time for each event.\",\n",
    "      \"parameter_definitions\": {\n",
    "        \"date\": {\n",
    "          \"description\": \"the date to list events for, formatted as mm/dd/yy\",\n",
    "          \"type\": \"str\",\n",
    "          \"required\": True\n",
    "        }\n",
    "      }\n",
    "    }, \n",
    "    {\n",
    "      \"name\": \"create_calendar_event\",\n",
    "      \"description\": \"Creates a new calendar event of the specified duration at the specified time and date. A new event cannot be created on the same time as an existing event.\",\n",
    "      \"parameter_definitions\": {\n",
    "        \"date\": {\n",
    "          \"description\": \"the date on which the event starts, formatted as mm/dd/yy\",\n",
    "          \"type\": \"str\",\n",
    "          \"required\": True\n",
    "        },\n",
    "        \"time\": {\n",
    "          \"description\": \"the time of the event, formatted using 24h military time formatting\",\n",
    "          \"type\": \"str\",\n",
    "          \"required\": True\n",
    "        },\n",
    "        \"duration\": {\n",
    "          \"description\": \"the number of hours the event lasts for\",\n",
    "          \"type\": \"float\",\n",
    "          \"required\": True\n",
    "        }\n",
    "      }\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s also create a custom preamble. The important part here is mentioning the date since the `list_calendar_events` tool is a simple enough version that it doesn’t do any date-handling. What we’re doing is nudging the model to use this date for the create_calendar_event tool call, which requires a date field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "preamble=\"\"\"## Task & Context\n",
    "You are a calendar assistant who helps people schedule events on their calendar. You must make sure that a new event does not overlap with any existing event.\n",
    "Today is Thursday, May 23, 2024\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a `run_assistant` function that does the following:\n",
    "- Get the user message (Step 1)\n",
    "- Call the Chat endpoint for tool call generation (Step 2)\n",
    "- If the response contains at least one tool call, execute the tool call(s) and get the tool results (Step 3)\n",
    "- Repeat Steps 2 and 3 until there are no more tool calls\n",
    "- Generate the final response with citations (Step 4)\n",
    "\n",
    "Here, we remove the `force_single_step` argument as it is `False` by default. This indicates to the Cohere API to enable a multi-step workflow and use the right ReAct-style prompt for running the workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = \"command-r-plus\"\n",
    "\n",
    "def run_assistant(message, chat_history=None):\n",
    "    if chat_history is None:\n",
    "        chat_history = []\n",
    "        \n",
    "    # Step 1: Get user message\n",
    "    print(f\"Question:\\n{message}\")\n",
    "    print(\"=\"*50)\n",
    "\n",
    "    # Step 2: Generate tool calls (if any)    \n",
    "    response = co.chat(\n",
    "        message=message,\n",
    "        model=model,\n",
    "        preamble=preamble,\n",
    "        tools=tools,\n",
    "        chat_history=chat_history\n",
    "    )\n",
    "\n",
    "    while response.tool_calls:\n",
    "        tool_calls = response.tool_calls\n",
    "        \n",
    "        if response.text:\n",
    "            print(\"Tool plan:\")\n",
    "            print(response.text,\"\\n\")\n",
    "        print(\"Tool calls:\")\n",
    "        for call in tool_calls:\n",
    "            print(f\"Tool name: {call.name} | Parameters: {call.parameters}\")\n",
    "        print(\"=\"*50)\n",
    "        \n",
    "        # Step 3: Get tool results\n",
    "        tool_results = []\n",
    "        for tc in tool_calls:\n",
    "            tool_call = {\"name\": tc.name, \"parameters\": tc.parameters}\n",
    "            tool_output = functions_map[tc.name](**tc.parameters)\n",
    "            tool_results.append({\"call\": tool_call, \"outputs\": [tool_output]})\n",
    "        \n",
    "        # Step 4: Generate response and citations                \n",
    "        response = co.chat(\n",
    "            message=\"\",\n",
    "            model=model,\n",
    "            preamble=preamble,\n",
    "            tools=tools,\n",
    "            tool_results=tool_results,\n",
    "            chat_history=response.chat_history\n",
    "        )\n",
    "\n",
    "        # Append the current chat turn to the chat history\n",
    "        chat_history = response.chat_history\n",
    "        \n",
    "    # Print final response\n",
    "    print(\"Final response:\")\n",
    "    print(response.text)\n",
    "    print(\"=\"*50)\n",
    "    \n",
    "    # Print citations (if any)\n",
    "    if response.citations:\n",
    "        print(\"Citations:\")\n",
    "        for citation in response.citations:\n",
    "            print(citation)\n",
    "        print(\"\\nCited Documents:\")\n",
    "        for document in response.documents:\n",
    "            print(document)\n",
    "        print(\"=\"*50)\n",
    "    \n",
    "    return chat_history"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple question that doesn't require multi-step"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s now ask the assistant the first question, starting with a simple one about the number of meetings for the day. This requires just a single step of tool calling to check existing calendar events.\n",
    "\n",
    "That’s exactly what the assistant does, and it gives the correct answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question:\n",
      "How many meetings do I have today\n",
      "==================================================\n",
      "Tool plan:\n",
      "I will use the 'list_calendar_events' tool to find out how many meetings the user has today. \n",
      "\n",
      "Tool calls:\n",
      "Tool name: list_calendar_events | Parameters: {'date': '05/23/2024'}\n",
      "==================================================\n",
      "Final response:\n",
      "You have four meetings today.\n",
      "==================================================\n",
      "Citations:\n",
      "start=9 end=22 text='four meetings' document_ids=['list_calendar_events:0:2:0']\n",
      "\n",
      "Cited Documents:\n",
      "{'existing_events': '[{\"end\":\"8:59\",\"start\":\"8:00\"},{\"end\":\"9:59\",\"start\":\"9:00\"},{\"end\":\"11:59\",\"start\":\"11:00\"},{\"end\":\"12:59\",\"start\":\"12:00\"}]', 'id': 'list_calendar_events:0:2:0', 'tool_name': 'list_calendar_events'}\n",
      "==================================================\n"
     ]
    }
   ],
   "source": [
    "chat_history = run_assistant(\"How many meetings do I have today\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multi-step"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s now try to ask a question that requires multi-step tool calling — this one asking the assistant to book an appointment. To complete this task, the assistant will have to first query the existing events and then use the information to create a new event that wouldn’t cause a conflict.\n",
    "\n",
    "In the first step, the assistant calls the `list_calendar_events` tool to get a list of existing events.\n",
    "\n",
    "In the second step, it calls the `create_calendar_event` tool to create a new event. It creates a new event at 10 a.m., which is indeed the first available slot after 9 a.m."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question:\n",
      "Create an hour-long appointment for the first available free slot after 9am\n",
      "==================================================\n",
      "Tool plan:\n",
      "I will first check the user's calendar for 23/05/2024 to see if there are any free slots after 9am. I will then create an hour-long appointment for the first available free slot. \n",
      "\n",
      "Tool calls:\n",
      "Tool name: list_calendar_events | Parameters: {'date': '05/23/2024'}\n",
      "==================================================\n",
      "Tool plan:\n",
      "The user's calendar shows that they have a free slot from 10:00 to 10:59. I will now create an hour-long appointment for this time slot. \n",
      "\n",
      "Tool calls:\n",
      "Tool name: create_calendar_event | Parameters: {'date': '05/23/2024', 'duration': 1, 'time': '10:00'}\n",
      "==================================================\n",
      "Final response:\n",
      "I've created a 1-hour long appointment for 10:00 on 23/05/2024.\n",
      "==================================================\n",
      "Citations:\n",
      "start=15 end=62 text='1-hour long appointment for 10:00 on 23/05/2024' document_ids=['create_calendar_event:0:4:0']\n",
      "\n",
      "Cited Documents:\n",
      "{'id': 'create_calendar_event:0:4:0', 'is_success': 'true', 'message': 'Created a 1 hour long event at 10:00 on 05/23/2024', 'tool_name': 'create_calendar_event'}\n",
      "==================================================\n"
     ]
    }
   ],
   "source": [
    "chat_history = run_assistant(\"Create an hour-long appointment for the first available free slot after 9am\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here’s a look at the chat history. Compared to the single-step scenario, it contains multiple pairs of CHATBOT - TOOL messages, with each pair being one step in the sequence of tool calls. In this particular case, we have two of them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "message='Create an hour-long appointment for the first available free slot after 9am' tool_calls=None role='USER' \n",
      "\n",
      "message=\"I will first check the user's calendar for 23/05/2024 to see if there are any free slots after 9am. I will then create an hour-long appointment for the first available free slot.\" tool_calls=[ToolCall(name='list_calendar_events', parameters={'date': '05/23/2024'})] role='CHATBOT' \n",
      "\n",
      "tool_results=[ToolResult(call=ToolCall(name='list_calendar_events', parameters={'date': '05/23/2024'}), outputs=[{'existing_events': [{'end': '8:59', 'start': '8:00'}, {'end': '9:59', 'start': '9:00'}, {'end': '11:59', 'start': '11:00'}, {'end': '12:59', 'start': '12:00'}]}])] role='TOOL' \n",
      "\n",
      "message=\"The user's calendar shows that they have a free slot from 10:00 to 10:59. I will now create an hour-long appointment for this time slot.\" tool_calls=[ToolCall(name='create_calendar_event', parameters={'date': '05/23/2024', 'duration': 1, 'time': '10:00'})] role='CHATBOT' \n",
      "\n",
      "tool_results=[ToolResult(call=ToolCall(name='create_calendar_event', parameters={'date': '05/23/2024', 'duration': 1, 'time': '10:00'}), outputs=[{'is_success': True, 'message': 'Created a 1 hour long event at 10:00 on 05/23/2024'}])] role='TOOL' \n",
      "\n",
      "message=\"I've created a 1-hour long appointment for 10:00 on 23/05/2024.\" tool_calls=None role='CHATBOT' \n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Print chat history\n",
    "for turn in chat_history:\n",
    "    print(turn,\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The chat history also highlights another difference between the single-step and multi-step scenarios.\n",
    "\n",
    "Notice that in the multi-step example above, the CHATBOT turns generate a textual response at each tool-calling step (look for the message) before making the actual tool calls.\n",
    "\n",
    "Let’s call this the model’s “intermediate response” to differentiate it from the model’s final response, which is the response that the user sees. These intermediate responses are the model’s internal reasoning logic, which guides its next course of action."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multi-step parallel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s now look at an example of multi-step, parallel tool use. Let’s ask the assistant to create two separate events on the same day.\n",
    "\n",
    "And here’s the assistant’s response. It follows the same two steps as in the earlier example: first, checking the existing events, and second, creating the new events.\n",
    "\n",
    "The difference this time is that since it needs to create two events, it calls `create_calendar_event` twice within the same step. Here the assistant makes the correct judgment and tool calls are independent of each other, so they can be done in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question:\n",
      "Create two hour-long appointments for any available time between 8am to 6pm\n",
      "==================================================\n",
      "Tool plan:\n",
      "I will first list the events for 23/05/2024. Then, I will create two hour-long appointments for any available time between 8am and 6pm. \n",
      "\n",
      "Tool calls:\n",
      "Tool name: list_calendar_events | Parameters: {'date': '05/23/2024'}\n",
      "==================================================\n",
      "Tool plan:\n",
      "The available times between 8am and 6pm are: 10am-11am, 1pm-6pm. I will now create two hour-long appointments for the user during these times. \n",
      "\n",
      "Tool calls:\n",
      "Tool name: create_calendar_event | Parameters: {'date': '05/23/2024', 'duration': 1, 'time': '10:00'}\n",
      "Tool name: create_calendar_event | Parameters: {'date': '05/23/2024', 'duration': 1, 'time': '13:00'}\n",
      "==================================================\n",
      "Final response:\n",
      "I have created two hour-long appointments for you today: one at 10am and another at 1pm.\n",
      "==================================================\n",
      "Citations:\n",
      "start=64 end=68 text='10am' document_ids=['create_calendar_event:0:4:0']\n",
      "start=84 end=87 text='1pm' document_ids=['create_calendar_event:1:4:0']\n",
      "\n",
      "Cited Documents:\n",
      "{'id': 'create_calendar_event:0:4:0', 'is_success': 'true', 'message': 'Created a 1 hour long event at 10:00 on 05/23/2024', 'tool_name': 'create_calendar_event'}\n",
      "{'id': 'create_calendar_event:1:4:0', 'is_success': 'true', 'message': 'Created a 1 hour long event at 13:00 on 05/23/2024', 'tool_name': 'create_calendar_event'}\n",
      "==================================================\n"
     ]
    }
   ],
   "source": [
    "chat_history = run_assistant(\"Create two hour-long appointments for any available time between 8am to 6pm\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## State management (memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The chat history for each turn consists of the following messages, and will accumulate with every new turn.\n",
    "\n",
    "- The USER message\n",
    "- Followed by the CHATBOT message with the list of tool calls\n",
    "- Followed the TOOL message with the list of tool results\n",
    "- Finally, followed by the CHATBOT message with the final response to the user\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question:\n",
      "Considering the new appointments you made, when is my next available time?\n",
      "==================================================\n",
      "Tool plan:\n",
      "I will list the events for 23/05/2024 and inform the user of their next available time. \n",
      "\n",
      "Tool calls:\n",
      "Tool name: list_calendar_events | Parameters: {'date': '05/23/2024'}\n",
      "==================================================\n",
      "Final response:\n",
      "Your next available time is 2pm.\n",
      "==================================================\n",
      "Citations:\n",
      "start=28 end=31 text='2pm' document_ids=['list_calendar_events:0:8:0']\n",
      "\n",
      "Cited Documents:\n",
      "{'existing_events': '[{\"end\":\"8:59\",\"start\":\"8:00\"},{\"end\":\"9:59\",\"start\":\"9:00\"},{\"end\":\"11:59\",\"start\":\"11:00\"},{\"end\":\"12:59\",\"start\":\"12:00\"}]', 'id': 'list_calendar_events:0:8:0', 'tool_name': 'list_calendar_events'}\n",
      "==================================================\n"
     ]
    }
   ],
   "source": [
    "chat_history = run_assistant(\"Considering the new appointments you made, when is my next available time?\", chat_history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "message='Create two hour-long appointments for any available time between 8am to 6pm' tool_calls=None role='USER' \n",
      "\n",
      "message='I will first list the events for 23/05/2024. Then, I will create two hour-long appointments for any available time between 8am and 6pm.' tool_calls=[ToolCall(name='list_calendar_events', parameters={'date': '05/23/2024'})] role='CHATBOT' \n",
      "\n",
      "tool_results=[ToolResult(call=ToolCall(name='list_calendar_events', parameters={'date': '05/23/2024'}), outputs=[{'existing_events': [{'end': '8:59', 'start': '8:00'}, {'end': '9:59', 'start': '9:00'}, {'end': '11:59', 'start': '11:00'}, {'end': '12:59', 'start': '12:00'}]}])] role='TOOL' \n",
      "\n",
      "message='The available times between 8am and 6pm are: 10am-11am, 1pm-6pm. I will now create two hour-long appointments for the user during these times.' tool_calls=[ToolCall(name='create_calendar_event', parameters={'date': '05/23/2024', 'duration': 1, 'time': '10:00'}), ToolCall(name='create_calendar_event', parameters={'date': '05/23/2024', 'duration': 1, 'time': '13:00'})] role='CHATBOT' \n",
      "\n",
      "tool_results=[ToolResult(call=ToolCall(name='create_calendar_event', parameters={'date': '05/23/2024', 'duration': 1, 'time': '10:00'}), outputs=[{'is_success': True, 'message': 'Created a 1 hour long event at 10:00 on 05/23/2024'}]), ToolResult(call=ToolCall(name='create_calendar_event', parameters={'date': '05/23/2024', 'duration': 1, 'time': '13:00'}), outputs=[{'is_success': True, 'message': 'Created a 1 hour long event at 13:00 on 05/23/2024'}])] role='TOOL' \n",
      "\n",
      "message='I have created two hour-long appointments for you today: one at 10am and another at 1pm.' tool_calls=None role='CHATBOT' \n",
      "\n",
      "message='Considering the new appointments you made, when is my next available time?' tool_calls=None role='USER' \n",
      "\n",
      "message='I will list the events for 23/05/2024 and inform the user of their next available time.' tool_calls=[ToolCall(name='list_calendar_events', parameters={'date': '05/23/2024'})] role='CHATBOT' \n",
      "\n",
      "tool_results=[ToolResult(call=ToolCall(name='list_calendar_events', parameters={'date': '05/23/2024'}), outputs=[{'existing_events': [{'end': '8:59', 'start': '8:00'}, {'end': '9:59', 'start': '9:00'}, {'end': '11:59', 'start': '11:00'}, {'end': '12:59', 'start': '12:00'}]}])] role='TOOL' \n",
      "\n",
      "message='Your next available time is 2pm.' tool_calls=None role='CHATBOT' \n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Print chat history\n",
    "for turn in chat_history:\n",
    "    print(turn,\"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}