{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "01e26d34",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "用户: 北京今天天气怎么样？\n",
      "助手: 您好！很抱歉，我无法获取实时的天气数据，因此无法告诉您北京今天的具体天气情况。\n",
      "\n",
      "建议您可以通过以下方式查询北京今天的天气：\n",
      "\n",
      "1. **手机天气应用** - 如墨迹天气、天气通等\n",
      "2. **搜索引擎...\n",
      "--------------------------------------------------\n",
      "\n",
      "用户: 帮我写一个 Python 快速排序\n",
      "助手: # Python 快速排序实现\n",
      "\n",
      "## 方法一：原地排序（经典版本）\n",
      "\n",
      "```python\n",
      "def quick_sort(arr, left, right):\n",
      "    \"\"\"原地快速排序\"\"\"\n",
      "   ...\n",
      "--------------------------------------------------\n",
      "\n",
      "用户: 你好，介绍一下你自己\n",
      "助手: 你好！我是 MiniMax AI，一个由 MiniMax 公司开发的人工智能助手。\n",
      "\n",
      "我可以帮助你：\n",
      "\n",
      "- 回答各种问题\n",
      "- 进行对话和交流\n",
      "- 协助写作、编程、翻译等任务\n",
      "- 提供信息和知识\n",
      "\n",
      "有什...\n",
      "--------------------------------------------------\n",
      "\n",
      "用户: 再见啦！\n",
      "助手: 再见！期待下次与你交流。...\n",
      "--------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from langgraph.graph import StateGraph, START, END, MessagesState\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.messages import SystemMessage, HumanMessage\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "# 初始化 LLM\n",
    "llm = ChatOpenAI(\n",
    "    model=\"MiniMax-M2.7-highspeed\",\n",
    "    openai_api_key=\"sk-8d657b8b7efe0cb6c141a30d9cee97f726efb9b18ea72bb1e8cfb080b42c140d\",\n",
    "    openai_api_base=\"https://console.pivotbak.cfd/v1\",\n",
    "    temperature=0.5\n",
    ")\n",
    "\n",
    "# 路由函数\n",
    "def classify_intent(state: MessagesState) -> str:\n",
    "    \"\"\"根据用户意图路由到不同的 Agent\"\"\"\n",
    "    last_message = state[\"messages\"][-1]\n",
    "    content = last_message.content.lower()\n",
    "    \n",
    "    if \"天气\" in content or \"温度\" in content:\n",
    "        return \"weather_agent\"\n",
    "    elif \"代码\" in content or \"编程\" in content:\n",
    "        return \"code_agent\"\n",
    "    elif \"再见\" in content or \"退出\" in content:\n",
    "        return \"farewell\"\n",
    "    else:\n",
    "        return \"general_agent\"\n",
    "\n",
    "# 定义各个 Agent 节点\n",
    "def router_node(state: MessagesState) -> dict:\n",
    "    \"\"\"路由节点：不做处理，只用于触发路由判断\"\"\"\n",
    "    return {}\n",
    "\n",
    "def weather_node(state: MessagesState) -> dict:\n",
    "    \"\"\"天气 Agent\"\"\"\n",
    "    response = llm.invoke([\n",
    "        SystemMessage(content=\"你是一个天气助手，友好地回答天气相关问题。如果没有实时数据，可以给出一般性建议。\"),\n",
    "        *state[\"messages\"]\n",
    "    ])\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "def code_node(state: MessagesState) -> dict:\n",
    "    \"\"\"代码 Agent\"\"\"\n",
    "    response = llm.invoke([\n",
    "        SystemMessage(content=\"你是一个编程助手，擅长解答代码问题并给出清晰的代码示例。\"),\n",
    "        *state[\"messages\"]\n",
    "    ])\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "def general_node(state: MessagesState) -> dict:\n",
    "    \"\"\"通用 Agent\"\"\"\n",
    "    response = llm.invoke([\n",
    "        SystemMessage(content=\"你是一个友善的 AI 助手，可以回答各种问题。\"),\n",
    "        *state[\"messages\"]\n",
    "    ])\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "def farewell_node(state: MessagesState) -> dict:\n",
    "    \"\"\"告别节点\"\"\"\n",
    "    return {\"messages\": [{\"role\": \"assistant\", \"content\": \"再见！期待下次与你交流。\"}]}\n",
    "\n",
    "# 构建图\n",
    "builder = StateGraph(MessagesState)\n",
    "\n",
    "# 添加节点\n",
    "builder.add_node(\"router\", router_node)\n",
    "builder.add_node(\"weather_agent\", weather_node)\n",
    "builder.add_node(\"code_agent\", code_node)\n",
    "builder.add_node(\"general_agent\", general_node)\n",
    "builder.add_node(\"farewell\", farewell_node)\n",
    "\n",
    "# 添加边\n",
    "builder.add_edge(START, \"router\")\n",
    "builder.add_conditional_edges(\n",
    "    \"router\",\n",
    "    classify_intent,\n",
    "    {\n",
    "        \"weather_agent\": \"weather_agent\",\n",
    "        \"code_agent\": \"code_agent\",\n",
    "        \"general_agent\": \"general_agent\",\n",
    "        \"farewell\": \"farewell\",\n",
    "    }\n",
    ")\n",
    "\n",
    "# 所有 agent 节点处理完后结束\n",
    "for node in [\"weather_agent\", \"code_agent\", \"general_agent\", \"farewell\"]:\n",
    "    builder.add_edge(node, END)\n",
    "\n",
    "# 编译图\n",
    "graph = builder.compile()\n",
    "\n",
    "# 测试不同意图\n",
    "test_inputs = [\n",
    "    \"北京今天天气怎么样？\",\n",
    "    \"帮我写一个 Python 快速排序\",\n",
    "    \"你好，介绍一下你自己\",\n",
    "    \"再见啦！\"\n",
    "]\n",
    "\n",
    "for user_input in test_inputs:\n",
    "    print(f\"\\n用户: {user_input}\")\n",
    "    result = graph.invoke({\"messages\": [HumanMessage(content=user_input)]})\n",
    "    print(f\"助手: {result['messages'][-1].content[:100]}...\")\n",
    "    print(\"-\" * 50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fbee2a39",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28808 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Image\n",
    "Image(graph.get_graph().draw_mermaid_png())"
   ]
  }
 ],
 "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
