{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dotenv import load_dotenv\n",
    "import os\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LangGraph Supervisor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_anthropic import ChatAnthropic\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "from langgraph_supervisor import create_supervisor\n",
    "from langgraph.prebuilt import create_react_agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model = ChatAnthropic(model=\"claude-3-7-sonnet-latest\", anthropic_api_key=os.getenv(\"ANTHROPIC_API_KEY\"))\n",
    "# model = ChatAnthropic(model=\"claude-3-5-haiku-latest\", anthropic_api_key=os.getenv(\"ANTHROPIC_API_KEY\"))\n",
    "model = ChatOpenAI(model=\"gpt-4o-mini\", openai_api_key=os.getenv(\"OPENAI_API_KEY\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 샘플 agent "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add(a: float, b: float) -> float:\n",
    "    \"\"\"Add two numbers together\"\"\"\n",
    "    return a + b\n",
    "\n",
    "def multiply(a: float, b: float) -> float:\n",
    "    \"\"\"Multiply two numbers together\"\"\"\n",
    "    return a * b\n",
    "\n",
    "def web_search(query: str) -> str:\n",
    "   \"\"\"Search the web for information.\"\"\"\n",
    "   return (\n",
    "       \"Here are the headcounts for each of the FAANG companies in 2024:\\n\"\n",
    "       \"1. **Meta**: 67,317 employees.\\n\"\n",
    "       \"2. **Apple**: 164,000 employees.\\n\"\n",
    "       \"3. **Amazon**: 1,551,000 employees.\\n\"\n",
    "       \"4. **Netflix**: 16,000 employees.\\n\"\n",
    "       \"5. **Google (Alphabet)**: 181,269 employees.\\n\"\n",
    "   )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "math_agent = create_react_agent(\n",
    "    model,\n",
    "    [add, multiply],\n",
    "    name = \"math_expert\",\n",
    "    prompt = \"You are a math expert. Always use one tool at a time.\"\n",
    ")\n",
    "\n",
    "research_agent = create_react_agent(\n",
    "    model,\n",
    "    [web_search],\n",
    "    name = \"research_expert\",\n",
    "    prompt = \"You are a world class researcher with access web search. Do not do any math.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Supervisor 선언"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow = create_supervisor(\n",
    "    [math_agent, research_agent],\n",
    "    model = model,\n",
    "    prompt = (\n",
    "        \"You are a team supervisor managing a research expert and a math expert. \"\n",
    "        \"For current events, use the research expert. \"\n",
    "        \"For math problems, use the math expert. \"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "app = workflow.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 그래프 시각화"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.graphs import visualize_graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ERROR] Visualize Graph Error: HTTPSConnectionPool(host='mermaid.ink', port=443): Read timed out. (read timeout=10)\n"
     ]
    }
   ],
   "source": [
    "visualize_graph(app)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = app.invoke(\n",
    "    {\n",
    "        \"messages\" : [\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": \"What is the combined headcount of FAANG in 2024?\"\n",
    "            }\n",
    "        ]\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content='What is the combined headcount of FAANG in 2024?', additional_kwargs={}, response_metadata={}, id='9446be73-2883-434d-af0c-dc96b6bb0097'),\n",
       "  AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_wNgOGpCBtFh6GpmirsfYNgn9', 'function': {'arguments': '{}', 'name': 'transfer_to_research_expert'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 105, 'total_tokens': 120, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_b8bc95a0ac', 'id': 'chatcmpl-BEFlmGvwI6OOTCwUkdePyRXCGjUUz', 'finish_reason': 'tool_calls', 'logprobs': None}, name='supervisor', id='run-76a21fee-3a5e-45da-9d48-edb4c53944aa-0', tool_calls=[{'name': 'transfer_to_research_expert', 'args': {}, 'id': 'call_wNgOGpCBtFh6GpmirsfYNgn9', 'type': 'tool_call'}], usage_metadata={'input_tokens': 105, 'output_tokens': 15, 'total_tokens': 120, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}),\n",
       "  ToolMessage(content='Successfully transferred to research_expert', name='transfer_to_research_expert', id='7bc10fb8-579b-46e5-b64f-7db38435fc02', tool_call_id='call_wNgOGpCBtFh6GpmirsfYNgn9'),\n",
       "  AIMessage(content='In 2024, the combined headcount of the FAANG companies is as follows:\\n\\n1. **Meta**: 67,317 employees\\n2. **Apple**: 164,000 employees\\n3. **Amazon**: 1,551,000 employees\\n4. **Netflix**: 16,000 employees\\n5. **Google (Alphabet)**: 181,269 employees\\n\\nThe total combined headcount of FAANG in 2024 is approximately 1,979,586 employees.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 106, 'prompt_tokens': 223, 'total_tokens': 329, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_b8bc95a0ac', 'id': 'chatcmpl-BEFlorTy4gcxqBl4zD97Hex7xGGGE', 'finish_reason': 'stop', 'logprobs': None}, name='research_expert', id='run-1a3771bc-c968-4c49-a4b4-32f5f7ed4f8f-0', usage_metadata={'input_tokens': 223, 'output_tokens': 106, 'total_tokens': 329, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}),\n",
       "  AIMessage(content='Transferring back to supervisor', additional_kwargs={}, response_metadata={}, name='research_expert', id='3c0d8016-a11e-4c71-8481-d6111d5d9055', tool_calls=[{'name': 'transfer_back_to_supervisor', 'args': {}, 'id': '43a7306d-953f-457a-aa2c-62c42f1a5d77', 'type': 'tool_call'}]),\n",
       "  ToolMessage(content='Successfully transferred back to supervisor', name='transfer_back_to_supervisor', id='b5f523c9-67a9-4e68-a13b-7c07e789e15a', tool_call_id='43a7306d-953f-457a-aa2c-62c42f1a5d77'),\n",
       "  AIMessage(content='The combined headcount of the FAANG companies in 2024 is approximately 1,979,586 employees. Here’s the breakdown:\\n\\n- Meta: 67,317 employees\\n- Apple: 164,000 employees\\n- Amazon: 1,551,000 employees\\n- Netflix: 16,000 employees\\n- Google (Alphabet): 181,269 employees', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 78, 'prompt_tokens': 299, 'total_tokens': 377, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_b8bc95a0ac', 'id': 'chatcmpl-BEFlql3tDiskUuSRExvqfzBz03oX3', 'finish_reason': 'stop', 'logprobs': None}, name='supervisor', id='run-f7c22ca5-0ab3-4f82-b70b-7ea0e4b909d4-0', usage_metadata={'input_tokens': 299, 'output_tokens': 78, 'total_tokens': 377, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "==============\n",
      "STEP: supervisor\n",
      "==============\n",
      "\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: transfer_to_research_expert\n",
      "\n",
      "Successfully transferred to research_expert\n",
      "\n",
      "==============\n",
      "STEP: research_expert\n",
      "==============\n",
      "\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: transfer_back_to_supervisor\n",
      "\n",
      "Successfully transferred back to supervisor\n",
      "\n",
      "==============\n",
      "STEP: supervisor\n",
      "==============\n",
      "\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Name: supervisor\n",
      "\n",
      "The combined headcount of FAANG companies in 2024 is 1,979,586 employees.\n"
     ]
    }
   ],
   "source": [
    "for event in app.stream({\"messages\": [(\"user\", \"What is the combined headcount of FAANG in 2024?\")]}):\n",
    "\n",
    "    for key, value in event.items():\n",
    "        print(f\"\\n==============\\nSTEP: {key}\\n==============\\n\")\n",
    "        \n",
    "        \n",
    "        if \"messages\" in value:\n",
    "            # 가장 최근 메시지를 출력\n",
    "            value[\"messages\"][-1].pretty_print()\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}