{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5007e5ee-b2f3-48ba-9e9e-bd2ca430c5e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#%% 0. 环境准备\n",
    "# 安装依赖（Notebook 中使用 !pip install）\n",
    "# !pip install langgraph langchain openai python-dotenv networkx matplotlib langsmith"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "4f88b9ff-ad04-4981-b892-9200e096d29f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```mermaid\n",
       "---\n",
       "config:\n",
       "  flowchart:\n",
       "    curve: linear\n",
       "---\n",
       "graph TD;\n",
       "\t__start__([<p>__start__</p>]):::first\n",
       "\tagent(agent)\n",
       "\ttool(tool)\n",
       "\tfinal(final)\n",
       "\t__end__([<p>__end__</p>]):::last\n",
       "\t__start__ --> agent;\n",
       "\tagent -.-> final;\n",
       "\tagent -.-> tool;\n",
       "\ttool --> final;\n",
       "\tfinal --> __end__;\n",
       "\tclassDef default fill:#f2f0ff,line-height:1.2\n",
       "\tclassDef first fill-opacity:0\n",
       "\tclassDef last fill:#bfb6fc\n",
       "\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[LangSmith] 开始追踪工作流执行...\n",
      "\n",
      "[Agent Decision] outcome = tool\n",
      "回复: 明天晴，气温25°C，适合出行。, 时间: 2025-05-05 13:23:00, 用户输入: 北京天气怎么样？\n",
      "\n",
      "[Traced] 对话记录:\n",
      "[human]: 北京天气怎么样？\n",
      "[ai]: 明天晴，气温25°C，适合出行。\n"
     ]
    }
   ],
   "source": [
    "#%% 1. 导入库\n",
    "import os\n",
    "from typing import TypedDict, List\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "from langchain_core.messages import BaseMessage, HumanMessage, AIMessage\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain.prompts import ChatPromptTemplate\n",
    "from langgraph.graph import StateGraph\n",
    "import networkx as nx\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import rcParams\n",
    "from datetime import datetime\n",
    "from IPython.display import display, Markdown\n",
    "from langchain_core.tracers import LangChainTracer\n",
    "from langsmith import traceable\n",
    "\n",
    "#%% 2. 环境设置与全局配置\n",
    "rcParams['font.sans-serif'] = ['Heiti TC']\n",
    "rcParams['axes.unicode_minus'] = False\n",
    "load_dotenv()\n",
    "\n",
    "#%% 3. 定义对话状态结构\n",
    "class AgentState(TypedDict):\n",
    "    messages: List[BaseMessage] # 对话历史\n",
    "    agent_outcome: str # 下一步决策：tool 或 final\n",
    "    tool_response: str # 工具调用结果\n",
    "\n",
    "#%% 4. 初始化大模型\n",
    "llm = ChatOpenAI(model_name=\"deepseek-chat\", verbose=True)\n",
    "\n",
    "#%% 5. 节点函数定义\n",
    "def agent_decision_node(state: AgentState) -> dict:\n",
    "    last_message = state[\"messages\"][-1].content\n",
    "    prompt = ChatPromptTemplate.from_template(\n",
    "        \"请判断用户的意图：\\n\"\n",
    "        \"用户输入：{input}\\n\\n\"\n",
    "        \"如果需要调用工具，请回答 tool；否则回答 final。\"\n",
    "    )\n",
    "    decision = llm.invoke(prompt.format_messages(input=last_message)).content.strip().lower()\n",
    "    print(f\"\\n[Agent Decision] outcome = {decision}\")\n",
    "    return {\"agent_outcome\": decision if decision in (\"tool\", \"final\") else \"final\"}\n",
    "\n",
    "def tool_node(_: AgentState) -> dict:\n",
    "    return {\"tool_response\": \"明天晴，气温25°C，适合出行。\"}\n",
    "\n",
    "def final_node(state: AgentState) -> dict:\n",
    "    # 获取工具的响应，或者默认消息\n",
    "    reply = state.get(\"tool_response\", \"很高兴为您服务！\")\n",
    "    \n",
    "    # 获取当前时间戳，作为元数据\n",
    "    timestamp = datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n",
    "    \n",
    "    # 可以添加额外的元数据，例如，用户的原始输入\n",
    "    user_input = state[\"messages\"][-1].content\n",
    "    \n",
    "    # 打印出回复信息及其附加内容\n",
    "    print(f\"回复: {reply}, 时间: {timestamp}, 用户输入: {user_input}\")\n",
    "    \n",
    "    # 返回包含原始消息和回复的消息记录\n",
    "    return {\n",
    "        \"messages\": state[\"messages\"] + [AIMessage(content=reply)],\n",
    "        \"metadata\": {\n",
    "            \"timestamp\": timestamp,\n",
    "            \"user_input\": user_input,\n",
    "            \"reply_length\": len(reply),\n",
    "        }\n",
    "    }\n",
    "\n",
    "#%% 6. 构建 LangGraph 工作流\n",
    "def build_graph():\n",
    "    workflow = StateGraph(AgentState)\n",
    "    workflow.add_node(\"agent\", agent_decision_node)\n",
    "    workflow.add_node(\"tool\", tool_node)\n",
    "    workflow.add_node(\"final\", final_node)\n",
    "    workflow.set_entry_point(\"agent\")\n",
    "    workflow.set_finish_point(\"final\")\n",
    "    workflow.add_conditional_edges(\"agent\", lambda st: st[\"agent_outcome\"], {\"tool\": \"tool\", \"final\": \"final\"})\n",
    "    workflow.add_edge(\"tool\", \"final\")\n",
    "    return workflow.compile()\n",
    "\n",
    "# 编译一次，供全局使用\n",
    "graph = build_graph()\n",
    "\n",
    "#%% 7. 运行测试\n",
    "def test_run():\n",
    "    init_state = {\"messages\": [HumanMessage(content=\"请问明天天气如何？\")]}\n",
    "    result = graph.invoke(init_state)\n",
    "    print(\"\\n对话记录:\")\n",
    "    for msg in result[\"messages\"]:\n",
    "        print(f\"[{msg.type}]: {msg.content}\")\n",
    "\n",
    "#%% 8. 使用 Mermaid 可视化工作流\n",
    "def visualize_workflow_with_mermaid(graph):\n",
    "    mermaid_code = graph.get_graph().draw_mermaid()  # 获取 Mermaid 格式代码\n",
    "    \n",
    "    # 通过 IPython 显示 Mermaid 图\n",
    "    display(Markdown(f\"```mermaid\\n{mermaid_code}\\n```\"))\n",
    "\n",
    "#%% 9. LangSmith 追踪（可选）\n",
    "def run_with_langsmith(input_state):\n",
    "    if \"LANGSMITH_API_KEY\" in os.environ:\n",
    "        tracer = LangChainTracer()\n",
    "\n",
    "        print(\"\\n[LangSmith] 开始追踪工作流执行...\")\n",
    "        result = graph.invoke(input_state, config={\"callbacks\": [tracer]})\n",
    "\n",
    "        print(\"\\n[Traced] 对话记录:\")\n",
    "        for msg in result[\"messages\"]:\n",
    "            print(f\"[{msg.type}]: {msg.content}\")\n",
    "    else:\n",
    "        print(\"\\n跳过 LangSmith 集成：未设置 LANGSMITH_API_KEY\")\n",
    "\n",
    "#%% 10. 主入口\n",
    "if __name__ == \"__main__\":\n",
    "    visualize_workflow_with_mermaid(graph)\n",
    "    # test_run()\n",
    "    run_with_langsmith({\"messages\": [HumanMessage(content=\"北京天气怎么样？\")]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daf80c8c-5a31-40db-b71a-1277a1a3091b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}