1_agentic-design-ptn/03_multi-agent/LangGraph/02.1_multi-agent-collabration.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d2b29a4b",
   "metadata": {},
   "source": [
    "# Multi-Agent Collaboration \n",
    "---\n",
    "\n",
    "### What is Multi-Agent Collaboration?\n",
    "Multi-Agent Collaboration refers to the process where multiple autonomous agents—each capable of independent decision-making—work together to achieve common or complementary objectives. This concept is widely used in fields like artificial intelligence, robotics, distributed computing, and simulation, and it involves several key aspects:\n",
    "\n",
    "- **Effective Communication and Coordination**:\n",
    "Agents exchange information and align their actions to collectively achieve a goal, ensuring that tasks are organized and synchronized.\n",
    "\n",
    "- **Autonomous, Distributed Decision-Making**:\n",
    "Each agent operates independently, making local decisions while contributing to a broader strategy, which enhances flexibility and fault tolerance.\n",
    "\n",
    "- **Adaptive Task Specialization**:\n",
    "Agents focus on specific roles or subtasks based on their capabilities, and they adjust their strategies through iterative feedback, leading to improved overall performance.\n",
    "\n",
    "\n",
    "### Key Advantages\n",
    "- **Efficiency Through Task Specialization**:\n",
    "By assigning specific roles to different agents, the system can handle complex tasks in parallel. This specialization allows each agent to focus on its area of expertise, resulting in faster and more effective problem-solving.\n",
    "\n",
    "- **Scalability and Flexibility**:\n",
    "AutoGen's structured communication and dynamic task allocation enable the system to scale easily. It can adapt to varying project complexities by adding or reassigning agents as needed, ensuring that the collaboration remains robust even as demands change.\n",
    "\n",
    "- **Enhanced Iterative Refinement**:\n",
    "The framework’s built-in feedback loops and iterative dialogue facilitate continuous improvement. Agents can refine their outputs based on real-time feedback, leading to more accurate and cohesive final results.\n",
    "\n",
    "**Reference**\n",
    "- [AutoGen paper: Enabling Next-Gen LLM Applications via Multi-Agent Conversation](https://arxiv.org/abs/2308.08155)\n",
    "- [Multi-Agent Collabration Concept](https://langchain-ai.github.io/langgraph/concepts/multi_agent/#network) \n",
    "- [LangChain `create_react_agent` built-in function](https://langchain-ai.github.io/langgraph/reference/prebuilt/#langgraph.prebuilt.chat_agent_executor.create_react_agent) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7862a931",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from azure_genai_utils.tracer import get_langchain_api_key, set_langsmith\n",
    "\n",
    "load_dotenv(override=True)\n",
    "\n",
    "# If you want to trace your RAG API calls, please set the tracing=True. You need to have a valid Langchain API key.\n",
    "langchain_key, has_langchain_key = get_langchain_api_key()\n",
    "set_langsmith(\"[RAG Innv Lab] 1_Agentic-Design-Pattern\", tracing=False)\n",
    "\n",
    "azure_openai_chat_deployment_name = os.getenv(\"AZURE_OPENAI_CHAT_DEPLOYMENT_NAME\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12699669",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 🧪 Step 1. Test and Construct each module\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc3c08d7",
   "metadata": {},
   "source": [
    "### Define your LLM\n",
    "This hands-on only uses the `gpt-4o-mini`, but you can utilize multiple models in the pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38fc3653",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import AzureChatOpenAI\n",
    "\n",
    "llm = AzureChatOpenAI(model=azure_openai_chat_deployment_name, temperature=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66eb4c57",
   "metadata": {},
   "source": [
    "### Tools\n",
    "\n",
    "Before building the entire the graph pipeline, we will test and construct each module separately.\n",
    "\n",
    "- **Researcher**\n",
    "- **Coder**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "731dec64",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure_genai_utils.tools import BingSearch\n",
    "from langchain_experimental.tools import PythonREPLTool\n",
    "\n",
    "WEB_SEARCH_FORMAT_OUTPUT = True\n",
    "\n",
    "web_search_tool = BingSearch(\n",
    "    max_results=3,\n",
    "    locale=\"en-US\",\n",
    "    include_news=False,\n",
    "    include_entity=False,\n",
    "    format_output=WEB_SEARCH_FORMAT_OUTPUT,\n",
    ")\n",
    "\n",
    "python_repl_tool = PythonREPLTool()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bff5cda9",
   "metadata": {},
   "outputs": [],
   "source": [
    "web_search_tool.invoke(\"Where is Seoul?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fe113f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "python_repl_tool.invoke(\"print('Hello, World!')\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21f87a5f",
   "metadata": {},
   "source": [
    "### ReAct agent test (Not required. Just for testing)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf6a0103",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "research_agent = create_react_agent(llm, tools=[web_search_tool])\n",
    "research_agent.invoke({\"messages\": \"Where is Seoul?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79fe861c",
   "metadata": {},
   "outputs": [],
   "source": [
    "coder_agent = create_react_agent(llm, tools=[python_repl_tool], prompt=None)\n",
    "coder_agent.invoke({\"messages\": [(\"user\", \"print 'Hello, World!'\")]})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99e92c04",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 🧪 Step 2. Define the Graph\n",
    "---\n",
    "\n",
    "### State Definition\n",
    "\n",
    "- `messages`: Messages to be passed between agents\n",
    "- `sender`: The sender of the next message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23606dcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "from typing import Sequence, Annotated\n",
    "from typing_extensions import TypedDict\n",
    "from langchain_core.messages import BaseMessage\n",
    "\n",
    "\n",
    "class AgentState(TypedDict):\n",
    "    messages: Annotated[Sequence[BaseMessage], operator.add]\n",
    "    sender: Annotated[str, \"The sender of the last message\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ce2e2dc",
   "metadata": {},
   "source": [
    "### Create Agents\n",
    "\n",
    "Create agents that perform sub-tasks.\n",
    "- `Researcher`: Researches the topic\n",
    "- `Coder`: Codes the solution\n",
    "\n",
    "#### Researcher Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "400e5479",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_system_prompt(suffix: str) -> str:\n",
    "    return (\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 or any of the other assistants have the final answer or deliverable,\"\n",
    "        \" prefix your response with FINAL ANSWER so the team knows to stop.\"\n",
    "        f\"\\n\\n{suffix}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89c72d9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "# Create a Research Agent\n",
    "research_agent = create_react_agent(\n",
    "    llm,\n",
    "    tools=[web_search_tool],\n",
    "    prompt=make_system_prompt(\n",
    "        \"You can only do research. You are working with a chart generator colleague.\"\n",
    "    ),\n",
    ")\n",
    "\n",
    "\n",
    "# Research Agent Node definition\n",
    "def research_node(state: AgentState) -> AgentState:\n",
    "\n",
    "    result = research_agent.invoke(state)\n",
    "\n",
    "    # Convert the last message to HumanMessage\n",
    "    last_message = HumanMessage(\n",
    "        content=result[\"messages\"][-1].content, name=\"researcher\"\n",
    "    )\n",
    "\n",
    "    return {\"messages\": [last_message], \"sender\": \"researcher\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d70a6215",
   "metadata": {},
   "source": [
    "Not required, but it is good to test the agent before creating the graph workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a06f15e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "research_node({\"messages\": \"Where is Seoul?\", \"sender\": \"\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d9e53f7",
   "metadata": {},
   "source": [
    "#### Coder (Chart Generator) Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c0ccbe8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import load_prompt\n",
    "\n",
    "code_system_prompt = load_prompt(\"../../../prompts/code-system-prompt-kr.yaml\").format()\n",
    "chart_agent = create_react_agent(\n",
    "    llm,\n",
    "    tools=[python_repl_tool],\n",
    "    prompt=make_system_prompt(code_system_prompt),\n",
    ")\n",
    "\n",
    "\n",
    "def chart_node(state: AgentState) -> AgentState:\n",
    "\n",
    "    result = chart_agent.invoke(state)\n",
    "\n",
    "    last_message = HumanMessage(\n",
    "        content=result[\"messages\"][-1].content, name=\"chart_generator\"\n",
    "    )\n",
    "    return {\n",
    "        # share internal message history of chart agent with other agents\n",
    "        \"messages\": [last_message],\n",
    "        \"sender\": \"chart_generator\",\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc49d497",
   "metadata": {},
   "source": [
    "### Construct the Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01fc6ca0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import HumanMessage, ToolMessage\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "\n",
    "def router(state: AgentState):\n",
    "    messages = state[\"messages\"]\n",
    "    last_message = messages[-1]\n",
    "    if \"FINAL ANSWER\" in last_message.content:\n",
    "        # Any agent decided the work is done\n",
    "        return END\n",
    "    return \"continue\"\n",
    "\n",
    "\n",
    "workflow = StateGraph(AgentState)\n",
    "\n",
    "# Node definition\n",
    "workflow.add_node(\"researcher\", research_node)\n",
    "workflow.add_node(\"chart_generator\", chart_node)\n",
    "\n",
    "# Edge definition\n",
    "workflow.add_conditional_edges(\n",
    "    \"researcher\",\n",
    "    router,\n",
    "    {\"continue\": \"chart_generator\", END: END},\n",
    ")\n",
    "workflow.add_conditional_edges(\n",
    "    \"chart_generator\",\n",
    "    router,\n",
    "    {\"continue\": \"researcher\", END: END},\n",
    ")\n",
    "\n",
    "workflow.add_edge(START, \"researcher\")\n",
    "\n",
    "app = workflow.compile(checkpointer=MemorySaver())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "361a3bd1",
   "metadata": {},
   "source": [
    "### Visualize the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb35a842",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure_genai_utils.graphs import visualize_langgraph\n",
    "\n",
    "visualize_langgraph(app, xray=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ecb01f9",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 🧪 Step 3. Execute the Graph\n",
    "---\n",
    "\n",
    "### Execute the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27f19ddf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables import RunnableConfig\n",
    "from azure_genai_utils.messages import stream_graph, invoke_graph, random_uuid\n",
    "\n",
    "config = RunnableConfig(recursion_limit=10, configurable={\"thread_id\": random_uuid()})\n",
    "\n",
    "inputs = {\n",
    "    \"messages\": [\n",
    "        HumanMessage(\n",
    "            content=\"Please draw a graph of Korea's population growth rate over the past 10 years.\"\n",
    "        )\n",
    "    ],\n",
    "}\n",
    "\n",
    "invoke_graph(app, inputs, config, node_names=[\"researcher\", \"chart_generator\", \"agent\"])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv_agent",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}