{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a3e3ebc4-57af-4fe4-bdd3-36aff67bf276",
   "metadata": {},
   "source": [
    "# Chat Bot Evaluation as Multi-agent Simulation\n",
    "\n",
    "When building a chat bot, such as a customer support assistant, it can be hard to properly evaluate your bot's performance. It's time-consuming to have to manually interact with it intensively for each code change.\n",
    "\n",
    "One way to make the evaluation process easier and more reproducible is to simulate a user interaction.\n",
    "\n",
    "With LangGraph, it's easy to set this up. Below is an example of how to create a \"virtual user\" to simulate a conversation.\n",
    "\n",
    "The overall simulation looks something like this:\n",
    "\n",
    "![diagram](./img/virtual_user_diagram.png)\n",
    "\n",
    "## Setup\n",
    "\n",
    "First, let's install the required packages and set our API keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0d30b6f7-3bec-4d9f-af50-43dfdc81ae6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U langgraph langchain langchain_openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "30c2f3de-c730-4aec-85a6-af2c2f058803",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "\n",
    "def _set_if_undefined(var: str):\n",
    "    if not os.environ.get(var):\n",
    "        os.environ[var] = getpass.getpass(f\"Please provide your {var}\")\n",
    "\n",
    "\n",
    "_set_if_undefined(\"OPENAI_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95c9332f",
   "metadata": {},
   "source": [
    "<div class=\"admonition tip\">\n",
    "    <p class=\"admonition-title\">Set up <a href=\"https://smith.langchain.com\">LangSmith</a> for LangGraph development</p>\n",
    "    <p style=\"padding-top: 5px;\">\n",
    "        Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started <a href=\"https://docs.smith.langchain.com\">here</a>. \n",
    "    </p>\n",
    "</div>   "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ef4528d-6b2a-47c7-98b5-50f14984a304",
   "metadata": {},
   "source": [
    "## Define Chat Bot\n",
    "\n",
    "Next, we will define our chat bot. For this notebook, we assume the bot's API accepts a list of messages and responds with a message. If you want to update this, all you'll have to change is this section and the \"get_messages_for_agent\" function in \n",
    "the simulator below.\n",
    "\n",
    "The implementation within `my_chat_bot` is configurable and can even be run on another system (e.g., if your system isn't running in python)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "828479af-cf9c-4888-a365-599643a96b55",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "import openai\n",
    "\n",
    "\n",
    "# This is flexible, but you can define your agent here, or call your agent API here.\n",
    "def my_chat_bot(messages: List[dict]) -> dict:\n",
    "    system_message = {\n",
    "        \"role\": \"system\",\n",
    "        \"content\": \"You are a customer support agent for an airline.\",\n",
    "    }\n",
    "    messages = [system_message] + messages\n",
    "    completion = openai.chat.completions.create(\n",
    "        messages=messages, model=\"gpt-3.5-turbo\"\n",
    "    )\n",
    "    return completion.choices[0].message.model_dump()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f58959bf-2ab5-4330-9ac2-c00f45237e24",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'content': 'Hello! How can I assist you today?',\n",
       " 'role': 'assistant',\n",
       " 'function_call': None,\n",
       " 'tool_calls': None}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_chat_bot([{\"role\": \"user\", \"content\": \"hi!\"}])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "419340a3-5ecf-48e7-9028-4f2fad750502",
   "metadata": {},
   "source": [
    "## Define Simulated User\n",
    "\n",
    "We're now going to define the simulated user. \n",
    "This can be anything we want, but we're going to build it as a LangChain bot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "32c147df-7f90-4b0d-9a6b-671677020353",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "system_prompt_template = \"\"\"You are a customer of an airline company. \\\n",
    "You are interacting with a user who is a customer support person. \\\n",
    "\n",
    "{instructions}\n",
    "\n",
    "When you are finished with the conversation, respond with a single word 'FINISHED'\"\"\"\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system_prompt_template),\n",
    "        MessagesPlaceholder(variable_name=\"messages\"),\n",
    "    ]\n",
    ")\n",
    "instructions = \"\"\"Your name is Harrison. You are trying to get a refund for the trip you took to Alaska. \\\n",
    "You want them to give you ALL the money back. \\\n",
    "This trip happened 5 years ago.\"\"\"\n",
    "\n",
    "prompt = prompt.partial(name=\"Harrison\", instructions=instructions)\n",
    "\n",
    "model = ChatOpenAI()\n",
    "\n",
    "simulated_user = prompt | model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6f80669e-aa78-4666-b67c-a539366d5aab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='Hi, I would like to request a refund for a trip I took with your airline company to Alaska. Is it possible to get a refund for that trip?')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "messages = [HumanMessage(content=\"Hi! How can I help you?\")]\n",
    "simulated_user.invoke({\"messages\": messages})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "321312b4-a1f0-4454-a481-fdac4e37cb7d",
   "metadata": {},
   "source": [
    "## Define the Agent Simulation\n",
    "\n",
    "The code below creates a LangGraph workflow to run the simulation. The main components are:\n",
    "\n",
    "1. The two nodes: one for the simulated user, the other for the chat bot.\n",
    "2. The graph itself, with a conditional stopping criterion.\n",
    "\n",
    "Read the comments in the code below for more information.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65bc4446-462b-4ee8-b017-2862fbbdfaf5",
   "metadata": {},
   "source": [
    "### Define nodes\n",
    "\n",
    "First, we define the nodes in the graph. These should take in a list of messages and return a list of messages to ADD to the state.\n",
    "These will be thing wrappers around the chat bot and simulated user we have above.\n",
    "\n",
    "**Note:** one tricky thing here is which messages are which. Because both the chat bot AND our simulated user are both LLMs, both of them will resond with AI messages. Our state will be a list of alternating Human and AI messages. This means that for one of the nodes, there will need to be some logic that flips the AI and human roles. In this example, we will assume that HumanMessages are messages from the simulated user. This means that we need some logic in the simulated user node to swap AI and Human messages.\n",
    "\n",
    "First, let's define the chat bot node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "69e2a3a3-40f3-4223-9136-113738440be9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.adapters.openai import convert_message_to_dict\n",
    "from langchain_core.messages import AIMessage\n",
    "\n",
    "\n",
    "def chat_bot_node(state):\n",
    "    messages = state[\"messages\"]\n",
    "    # Convert from LangChain format to the OpenAI format, which our chatbot function expects.\n",
    "    messages = [convert_message_to_dict(m) for m in messages]\n",
    "    # Call the chat bot\n",
    "    chat_bot_response = my_chat_bot(messages)\n",
    "    # Respond with an AI Message\n",
    "    return {\"messages\": [AIMessage(content=chat_bot_response[\"content\"])]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "694c3c0c-56c5-4410-8fa8-ea2c0f11f506",
   "metadata": {},
   "source": [
    "Next, let's define the node for our simulated user. This will involve a little logic to swap the roles of the messages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7cad7527-ffa5-4c30-8585-b54a7a18bd98",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _swap_roles(messages):\n",
    "    new_messages = []\n",
    "    for m in messages:\n",
    "        if isinstance(m, AIMessage):\n",
    "            new_messages.append(HumanMessage(content=m.content))\n",
    "        else:\n",
    "            new_messages.append(AIMessage(content=m.content))\n",
    "    return new_messages\n",
    "\n",
    "\n",
    "def simulated_user_node(state):\n",
    "    messages = state[\"messages\"]\n",
    "    # Swap roles of messages\n",
    "    new_messages = _swap_roles(messages)\n",
    "    # Call the simulated user\n",
    "    response = simulated_user.invoke({\"messages\": new_messages})\n",
    "    # This response is an AI message - we need to flip this to be a human message\n",
    "    return {\"messages\": [HumanMessage(content=response.content)]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a48d8a3e-9171-4c43-a595-44d312722148",
   "metadata": {},
   "source": [
    "### Define edges\n",
    "\n",
    "We now need to define the logic for the edges. The main logic occurs after the simulated user goes, and it should lead to one of two outcomes:\n",
    "\n",
    "- Either we continue and call the customer support bot\n",
    "- Or we finish and the conversation is over\n",
    "\n",
    "So what is the logic for the conversation being over? We will define that as either the Human chatbot responds with `FINISHED` (see the system prompt) OR the conversation is more than 6 messages long (this is an arbitrary number just to keep this example short)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "28004fbf-a2f3-46b7-bde7-46c7adaf97fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def should_continue(state):\n",
    "    messages = state[\"messages\"]\n",
    "    if len(messages) > 6:\n",
    "        return \"end\"\n",
    "    elif messages[-1].content == \"FINISHED\":\n",
    "        return \"end\"\n",
    "    else:\n",
    "        return \"continue\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0856d4f-9334-4f28-944b-06d303e913a4",
   "metadata": {},
   "source": [
    "### Define graph\n",
    "\n",
    "We can now define the graph that sets up the simulation!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0b597e4b-4cbb-4bbc-82e5-f7e31275964c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import END, StateGraph, START\n",
    "from langgraph.graph.message import add_messages\n",
    "from typing import Annotated\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "\n",
    "class State(TypedDict):\n",
    "    messages: Annotated[list, add_messages]\n",
    "\n",
    "\n",
    "graph_builder = StateGraph(State)\n",
    "graph_builder.add_node(\"user\", simulated_user_node)\n",
    "graph_builder.add_node(\"chat_bot\", chat_bot_node)\n",
    "# Every response from  your chat bot will automatically go to the\n",
    "# simulated user\n",
    "graph_builder.add_edge(\"chat_bot\", \"user\")\n",
    "graph_builder.add_conditional_edges(\n",
    "    \"user\",\n",
    "    should_continue,\n",
    "    # If the finish criteria are met, we will stop the simulation,\n",
    "    # otherwise, the virtual user's message will be sent to your chat bot\n",
    "    {\n",
    "        \"end\": END,\n",
    "        \"continue\": \"chat_bot\",\n",
    "    },\n",
    ")\n",
    "# The input will first go to your chat bot\n",
    "graph_builder.add_edge(START, \"chat_bot\")\n",
    "simulation = graph_builder.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e0bd26e-8c1d-471d-9fef-d95dc0163491",
   "metadata": {},
   "source": [
    "## Run Simulation\n",
    "\n",
    "Now we can evaluate our chat bot! We can invoke it with empty messages (this will simulate letting the chat bot start the initial conversation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "32848c2e-be82-46f3-81db-b23fea45461c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'chat_bot': AIMessage(content='How may I assist you today regarding your flight or any other concerns?')}\n",
      "----\n",
      "{'user': HumanMessage(content='Hi, my name is Harrison. I am reaching out to request a refund for a trip I took to Alaska with your airline company. The trip occurred about 5 years ago. I would like to receive a refund for the entire amount I paid for the trip. Can you please assist me with this?')}\n",
      "----\n",
      "{'chat_bot': AIMessage(content=\"Hello, Harrison. Thank you for reaching out to us. I understand you would like to request a refund for a trip you took to Alaska five years ago. I'm afraid that our refund policy typically has a specific timeframe within which refund requests must be made. Generally, refund requests need to be submitted within 24 to 48 hours after the booking is made, or in certain cases, within a specified cancellation period.\\n\\nHowever, I will do my best to assist you. Could you please provide me with some additional information? Can you recall any specific details about the booking, such as the flight dates, booking reference or confirmation number? This will help me further look into the possibility of processing a refund for you.\")}\n",
      "----\n",
      "{'user': HumanMessage(content=\"Hello, thank you for your response. I apologize for not requesting the refund earlier. Unfortunately, I don't have the specific details such as the flight dates, booking reference, or confirmation number at the moment. Is there any other way we can proceed with the refund request without these specific details? I would greatly appreciate your assistance in finding a solution.\")}\n",
      "----\n",
      "{'chat_bot': AIMessage(content=\"I understand the situation, Harrison. Without specific details like flight dates, booking reference, or confirmation number, it becomes challenging to locate and process the refund accurately. However, I can still try to help you.\\n\\nTo proceed further, could you please provide me with any additional information you might remember? This could include the approximate date of travel, the departure and arrival airports, the names of the passengers, or any other relevant details related to the booking. The more information you can provide, the better we can investigate the possibility of processing a refund for you.\\n\\nAdditionally, do you happen to have any documentation related to your trip, such as receipts, boarding passes, or emails from our airline? These documents could assist in verifying your trip and processing the refund request.\\n\\nI apologize for any inconvenience caused, and I'll do my best to assist you further based on the information you can provide.\")}\n",
      "----\n",
      "{'user': HumanMessage(content=\"I apologize for the inconvenience caused. Unfortunately, I don't have any additional information or documentation related to the trip. It seems that I am unable to provide you with the necessary details to process the refund request. I understand that this may limit your ability to assist me further, but I appreciate your efforts in trying to help. Thank you for your time. \\n\\nFINISHED\")}\n",
      "----\n",
      "{'chat_bot': AIMessage(content=\"I understand, Harrison. I apologize for any inconvenience caused, and I appreciate your understanding. If you happen to locate any additional information or documentation in the future, please don't hesitate to reach out to us again. Our team will be more than happy to assist you with your refund request or any other travel-related inquiries. Thank you for contacting us, and have a great day!\")}\n",
      "----\n",
      "{'user': HumanMessage(content='FINISHED')}\n",
      "----\n"
     ]
    }
   ],
   "source": [
    "for chunk in simulation.stream({}):\n",
    "    # Print out all events aside from the final end chunk\n",
    "    if END not in chunk:\n",
    "        print(chunk)\n",
    "        print(\"----\")"
   ]
  }
 ],
 "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.11.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}