{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "39fd1948-b5c3-48c4-b10e-2ae7e8c83334",
   "metadata": {},
   "source": [
    "# Multi-agent Collaboration Intro\n",
    "\n",
    "A single agent can usually operate effectively using a handful of tools and a scoped domain, but even using powerful models like `gpt-4`, it can be less effective at using many tools. One way to improve your system's performance is through multi-agent collaboration. Each agent can specialize in a task or domain, and LangGraph can effectively orchestrate them to accomplish a larger goal. \n",
    "\n",
    "This notebook is inspired by the paper [AutoGen: Enabling Next-Gen LLM Applications via Multi-Agent Conversation](https://arxiv.org/abs/2308.08155), by Wu, et. al."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d7b6dcc-c985-46e2-8457-7e6b0298b950",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -U langchain langchain_openai langsmith pandas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "743c19df-6da9-4d1e-b2d2-ea40080b9fdc",
   "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(f\"Please provide your {var}\")\n",
    "\n",
    "\n",
    "_set_if_undefined(\"OPENAI_API_KEY\")\n",
    "_set_if_undefined(\"LANGCHAIN_API_KEY\")\n",
    "_set_if_undefined(\"TAVILY_API_KEY\")\n",
    "\n",
    "# Optional, add tracing in LangSmith\n",
    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
    "os.environ[\"LANGCHAIN_PROJECT\"] = \"Multi-agent Collaboration\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d5fc0f3-5d9e-4e72-a281-177f101c2a7d",
   "metadata": {},
   "source": [
    "## Example 1: 2 Agents\n",
    "\n",
    "Below is an example of 2 agents collaborating to accomplish a single task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "075c91c3-c249-471d-b259-41975faa83fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Tuple, Union\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from langchain_community.tools.tavily_search import TavilySearchResults\n",
    "from langchain_core.tools import tool\n",
    "\n",
    "tavily_tool = TavilySearchResults(max_results=5)\n",
    "\n",
    "\n",
    "@tool\n",
    "def create_plot(\n",
    "    data: Union[List[float], List[int]],\n",
    "    labels: Union[List[str], None] = None,\n",
    "    title: str = \"Plot\",\n",
    "    xlabel: str = \"X\",\n",
    "    ylabel: str = \"Y\",\n",
    "    color: Union[str, List[str]] = \"blue\",\n",
    "    plot_type: str = \"bar\",\n",
    ") -> Tuple[plt.Figure, plt.Axes]:\n",
    "    \"\"\"\n",
    "    Generates a bar or line plot from the provided data and returns the figure and axis objects.\n",
    "\n",
    "    :param data: A list of numerical values for the bar heights or line points.\n",
    "    :param labels: A list of strings for the bar or point labels. Default is None.\n",
    "    :param title: Title of the plot. Default is 'Plot'.\n",
    "    :param xlabel: Label for the X-axis. Default is 'X'.\n",
    "    :param ylabel: Label for the Y-axis. Default is 'Y'.\n",
    "    :param color: Color of the bars or line. Can be a single color or a list of colors. Default is 'blue'.\n",
    "    :param figsize: Size of the figure as a tuple (width, height). Default is (10, 6).\n",
    "    :param plot_type: Type of plot ('bar' or 'line'). Default is 'bar'.\n",
    "    :return: Tuple containing the figure and axes objects.\n",
    "    \"\"\"\n",
    "    if plot_type not in [\"bar\", \"line\"]:\n",
    "        raise ValueError(\"Invalid plot_type. Expected 'bar' or 'line'.\")\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(10, 6))\n",
    "    x_positions = range(len(data))\n",
    "\n",
    "    if labels and len(labels) == len(data):\n",
    "        plt.xticks(x_positions, labels)\n",
    "\n",
    "    if plot_type == \"bar\":\n",
    "        ax.bar(x_positions, data, color=color)\n",
    "    elif plot_type == \"line\":\n",
    "        ax.plot(x_positions, data, color=color, marker=\"o\")  # 'o' for circular markers\n",
    "\n",
    "    ax.set_title(title)\n",
    "    ax.set_xlabel(xlabel)\n",
    "    ax.set_ylabel(ylabel)\n",
    "\n",
    "    return fig, ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0ca7d80-31e4-4394-bfce-ffac314bac7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import operator\n",
    "from typing import Annotated, Sequence, TypedDict\n",
    "\n",
    "from langchain.agents import create_openai_functions_agent\n",
    "from langchain.tools.render import format_tool_to_openai_function\n",
    "from langchain_core.messages import (\n",
    "    AIMessage,\n",
    "    BaseMessage,\n",
    "    ChatMessage,\n",
    "    FunctionMessage,\n",
    "    HumanMessage,\n",
    ")\n",
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain_openai import ChatOpenAI\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "from langgraph.graph import END, StateGraph\n",
    "from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation\n",
    "\n",
    "\n",
    "class AgentState(TypedDict):\n",
    "    messages: Annotated[Sequence[BaseMessage], operator.add]\n",
    "    sender: str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dce3901-6ad5-4df5-8528-6e865cf96cb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow = StateGraph(AgentState)\n",
    "\n",
    "# This a helper class we have that is useful for running tools\n",
    "# It takes in an agent action and calls that tool and returns the result\n",
    "tools = [tavily_tool, create_plot]\n",
    "tool_executor = ToolExecutor(tools)\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You are a helpful AI assistant, collaborating with other assistants.\"\n",
    "            \" Use the provided tools to progress towards answering the question.\"\n",
    "            \" If you are unable to fully answer, that's OK, another assistant with different tools \"\n",
    "            \" will help where you left off. Execute what you can to make progress.\"\n",
    "            \" If you have the final answer, prefix your response with FINAL ANSWER.\"\n",
    "            \" You have access to the following tools: {tool_names}.\",\n",
    "        ),\n",
    "        MessagesPlaceholder(variable_name=\"messages\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "def add_agent_node(name: str, llm, tools, prompt):\n",
    "    functions = [format_tool_to_openai_function(t) for t in tools]\n",
    "\n",
    "    def _update_state(ai_message) -> dict:\n",
    "        if isinstance(ai_message, FunctionMessage):\n",
    "            result = ai_message\n",
    "        else:\n",
    "            message = ai_message.dict(exclude={\"type\"})\n",
    "            message[\"name\"] = name\n",
    "            result = HumanMessage(**message)\n",
    "        return {\n",
    "            \"messages\": [result],\n",
    "            \"sender\": name,\n",
    "        }\n",
    "\n",
    "    chain = (\n",
    "        (lambda x: {**x, \"intermediate_steps\": []})\n",
    "        | prompt.partial(tool_names=\", \".join([tool.name for tool in tools]))\n",
    "        | llm.bind_functions(functions)\n",
    "        | _update_state\n",
    "    )\n",
    "    workflow.add_node(name, chain)\n",
    "\n",
    "\n",
    "def call_tool(state):\n",
    "    messages = state[\"messages\"]\n",
    "    # Based on the continue condition\n",
    "    # we know the last message involves a function call\n",
    "    last_message = messages[-1]\n",
    "    # We construct an ToolInvocation from the function_call\n",
    "    action = ToolInvocation(\n",
    "        tool=last_message.additional_kwargs[\"function_call\"][\"name\"],\n",
    "        tool_input=json.loads(\n",
    "            last_message.additional_kwargs[\"function_call\"][\"arguments\"]\n",
    "        ),\n",
    "    )\n",
    "    # We call the tool_executor and get back a response\n",
    "    response = tool_executor.invoke(action)\n",
    "    # We use the response to create a FunctionMessage\n",
    "    function_message = FunctionMessage(content=str(response), name=action.tool)\n",
    "    # We return a list, because this will get added to the existing list\n",
    "    return {\"messages\": [function_message]}\n",
    "\n",
    "\n",
    "def should_continue(state):\n",
    "    # This is the router\n",
    "    messages = state[\"messages\"]\n",
    "    last_message = messages[-1]\n",
    "    if \"function_call\" in last_message.additional_kwargs:\n",
    "        return \"call_tool\"\n",
    "    if \"FINAL ANSWER\" in last_message.content:\n",
    "        return \"end\"\n",
    "    return \"continue\"\n",
    "\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4\")\n",
    "add_agent_node(\n",
    "    \"Researcher\",\n",
    "    llm,\n",
    "    [tavily_tool],\n",
    "    prompt,\n",
    ")\n",
    "add_agent_node(\"Chart Generator\", llm, [create_plot], prompt)\n",
    "workflow.add_node(\"call_tool\", call_tool)\n",
    "workflow.add_conditional_edges(\n",
    "    \"Researcher\",\n",
    "    should_continue,\n",
    "    {\"continue\": \"Chart Generator\", \"call_tool\": \"call_tool\", \"end\": END},\n",
    ")\n",
    "workflow.add_conditional_edges(\n",
    "    \"Chart Generator\",\n",
    "    should_continue,\n",
    "    {\"continue\": \"Researcher\", \"call_tool\": \"call_tool\", \"end\": END},\n",
    ")\n",
    "# We will let the user_info tool decide when to respond.\n",
    "workflow.add_edge(\"call_tool\", \"Researcher\")\n",
    "workflow.set_entry_point(\"Researcher\")\n",
    "graph = workflow.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "176a99b0-b457-45cf-8901-90facaa852da",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = graph.invoke(\n",
    "    {\n",
    "        \"messages\": [\n",
    "            HumanMessage(\n",
    "                content=\"Fetch the UK's GDP over the past 5 years, then draw a line graph of it.\"\n",
    "            )\n",
    "        ]\n",
    "    },\n",
    "    # Maximum number of steps to take in the graph\n",
    "    {\"recursion_limit\": 150},\n",
    ")\n",
    "result[\"messages\"][-1]"
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}