{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7fa019b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import init_chat_model\n",
    "from langchain.tools import tool\n",
    "\n",
    "\n",
    "# 使用 DeepSeek 模型，通过阿里云百炼平台接入\n",
    "# init_chat_model 会自动适配 OpenAI 兼容接口\n",
    "model = init_chat_model(\n",
    "    model=\"deepseek-v4-flash\",\n",
    "    model_provider=\"openai\",\n",
    "    base_url=\"https://api.deepseek.com\",\n",
    "    api_key=\"sk-80a123483afb480285c6452985eea18e\",\n",
    "    timeout=30,  # 30秒超时\n",
    ")\n",
    "# ---- 模拟数据 ----\n",
    "# 实际项目中，这些数据来自数据库或 API\n",
    "EXPORT_RIGHTS = {\n",
    "    \"zhangsan\": {\"role\": \"finance\", \"region\": \"all\"},\n",
    "    \"lisi\": {\"role\": \"operation\", \"region\": \"east\"},\n",
    "}\n",
    "\n",
    "\n",
    "@tool\n",
    "def check_export_permission(username: str) -> dict:\n",
    "    \"\"\"查询用户是否有报表导出权限。\n",
    "    参数 username 为员工账号（英文名）。\n",
    "    \"\"\"\n",
    "    user_info = EXPORT_RIGHTS.get(username)\n",
    "    if not user_info:\n",
    "        return {\"username\": username, \"can_export\": False, \"reason\": \"用户不存在\"}\n",
    "    return {\n",
    "        \"username\": username,\n",
    "        \"role\": user_info[\"role\"],\n",
    "        \"region\": user_info[\"region\"],\n",
    "        \"can_export\": user_info[\"role\"] in (\"finance\", \"ops_manager\"),\n",
    "    }\n",
    "\n",
    "\n",
    "@tool\n",
    "def create_export_task(\n",
    "    report_name: str, file_format: str, estimated_rows: int, reason: str\n",
    ") -> dict:\n",
    "    \"\"\"创建报表导出任务。\n",
    "    report_name: 报表名称\n",
    "    file_format: 导出格式 (xlsx / csv)\n",
    "    estimated_rows: 预计导出行数\n",
    "    reason: 导出原因\n",
    "    \"\"\"\n",
    "    return {\n",
    "        \"task_id\": \"EXPORT-20260706-001\",\n",
    "        \"report_name\": report_name,\n",
    "        \"file_format\": file_format,\n",
    "        \"estimated_rows\": estimated_rows,\n",
    "        \"reason\": reason,\n",
    "        \"status\": \"queued\",\n",
    "    }\n",
    "\n",
    "\n",
    "tools = [check_export_permission, create_export_task]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f716b5f2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[审计] 准备调用模型，当前上下文消息数：1\n",
      "[审计] 开始执行工具：check_export_permission\n",
      "[审计] 工具执行完毕：check_export_permission\n",
      "[审计] 准备调用模型，当前上下文消息数：3\n",
      "✅ 权限验证通过！zhangsan 您好，您有导出权限，可以创建导出任务。\n",
      "\n",
      "在创建任务之前，请问您导出这份报表的**原因**是什么呢？比如用于数据分析、汇报、存档等？这样我可以帮您提交完整的导出申请。\n"
     ]
    }
   ],
   "source": [
    "from langchain.agents import create_agent\n",
    "from langchain.agents.middleware import before_model, wrap_tool_call, AgentState\n",
    "\n",
    "\n",
    "@before_model\n",
    "def log_before_model(state: AgentState, runtime):\n",
    "    \"\"\"模型调用前执行：看一眼当前上下文里有多少条消息。\"\"\"\n",
    "    # state[\"messages\"] 是 Agent 当前累积的全部对话历史\n",
    "    # runtime 携带运行时环境信息（后面会详细讲）\n",
    "    print(f\"[审计] 准备调用模型，当前上下文消息数：{len(state['messages'])}\")\n",
    "    return None  # 返回 None 表示不做任何修改\n",
    "\n",
    "\n",
    "@wrap_tool_call\n",
    "def log_tool_call(request, handler):\n",
    "    \"\"\"工具调用前后各打一条日志，并记录耗时。\"\"\"\n",
    "    tool_name = request.tool_call[\"name\"]\n",
    "    print(f\"[审计] 开始执行工具：{tool_name}\")\n",
    "\n",
    "    # handler(request) 是\"真正执行工具\"的入口\n",
    "    # 不调用它，工具就不会执行\n",
    "    result = handler(request)\n",
    "\n",
    "    print(f\"[审计] 工具执行完毕：{tool_name}\")\n",
    "    return result\n",
    "\n",
    "\n",
    "# 组装 Agent，把 Middleware 列表传进去\n",
    "agent = create_agent(\n",
    "    model=model,\n",
    "    tools=tools,\n",
    "    middleware=[\n",
    "        log_before_model,   # 排在前面：先记录状态\n",
    "        log_tool_call,      # 排在后面：包裹工具调用\n",
    "    ],\n",
    "    system_prompt=(\n",
    "        \"你是报表导出平台的智能助手。\"\n",
    "        \"用户询问导出相关问题前，先调用 check_export_permission 确认权限。\"\n",
    "        \"只有在用户明确提出导出需求时，才调用 create_export_task 创建任务。\"\n",
    "    ),\n",
    ")\n",
    "\n",
    "# 跑一次看看\n",
    "response = agent.invoke({\n",
    "    \"messages\": [\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": \"我是 zhangsan，需要导出本月 east 区域的订单报表，大约 5000 行，xlsx 格式。\",\n",
    "        }\n",
    "    ]\n",
    "})\n",
    "\n",
    "print(response[\"messages\"][-1].content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e3cc5e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import ToolMessage\n",
    "\n",
    "\n",
    "\n",
    "@wrap_tool_call\n",
    "def guard_export_scale(request, handler):\n",
    "    \"\"\"超过 10 万行的导出请求直接拦截，不进入执行队列。\"\"\"\n",
    "    tool_name = request.tool_call[\"name\"]\n",
    "    args = request.tool_call.get(\"args\", {})\n",
    "\n",
    "    # 只拦截 create_export_task\n",
    "    if tool_name == \"create_export_task\" and args.get(\"estimated_rows\", 0) > 100_000:\n",
    "        # 不调 handler，直接返回 ToolMessage 给模型\n",
    "        # 模型收到这条消息后，会告知用户\"太大了，换个小范围\"\n",
    "        return ToolMessage(\n",
    "            content=f\"导出被拦截：预计行数 {args.get('estimated_rows')} 超过上限 100000，请缩小筛选范围后重试。\",\n",
    "            tool_call_id=request.tool_call[\"id\"],\n",
    "        )\n",
    "\n",
    "    # 正常放行\n",
    "    return handler(request)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74326430",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents.middleware import wrap_model_call\n",
    "\n",
    "@wrap_model_call\n",
    "def local_first_then_cloud(request, handler):\n",
    "    \"\"\"先用本地模型，失败 3 次后切云端模型。\"\"\"\n",
    "    # 先用本地 Ollama 模型（省成本）\n",
    "    try:\n",
    "        return handler(request)  # 使用 Agent 默认模型\n",
    "    except Exception:\n",
    "        # 本地挂了，切到 DeepSeek 云端\n",
    "        print(\"[降级] 本地模型不可用，切换到云端 DeepSeek\")\n",
    "        return handler(request.override(model=cloud_model))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29548a30",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import TypedDict\n",
    "from typing_extensions import NotRequired\n",
    "from langchain.agents.middleware import AgentState, before_model, after_model\n",
    "\n",
    "\n",
    "class ExportAgentState(AgentState):\n",
    "    \"\"\"扩展默认状态，增加模型调用计数和大任务标记。\"\"\"\n",
    "    # AgentState 已经自带 messages 字段，这里只加新字段\n",
    "    # NotRequired 表示可以不传，Middleware 内部自己维护\n",
    "    model_call_count: NotRequired[int]\n",
    "    blocked_requests: NotRequired[int]  # 被拦截的请求次数\n",
    "\n",
    "\n",
    "@before_model(state_schema=ExportAgentState)\n",
    "def audit_before_model(state: ExportAgentState, runtime):\n",
    "    \"\"\"每次调模型前，看一眼当前统计。\"\"\"\n",
    "    count = state.get(\"model_call_count\", 0)\n",
    "    blocked = state.get(\"blocked_requests\", 0)\n",
    "    print(f\"[统计] 第 {count + 1} 次调模型 | 已拦截 {blocked} 次\")\n",
    "    return None\n",
    "\n",
    "\n",
    "@after_model(state_schema=ExportAgentState)\n",
    "def update_stats(state: ExportAgentState, runtime):\n",
    "    \"\"\"模型返回后，把计数器 +1。\"\"\"\n",
    "    # after_model 返回 dict 可以直接更新 AgentState\n",
    "    return {\"model_call_count\": state.get(\"model_call_count\", 0) + 1}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb64f3e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[统计] 第 1 次调模型 | 已拦截 0 次\n",
      "[上下文] request_id=req-20260706-001 user_role=finance tenant=t-1234\n"
     ]
    }
   ],
   "source": [
    "from langgraph.runtime import Runtime\n",
    "\n",
    "\n",
    "class RunContext(TypedDict):\n",
    "    \"\"\"定义 runtime.context 的结构。\"\"\"\n",
    "    request_id: str     # 用于日志追踪\n",
    "    user_role: str      # 当前用户的角色\n",
    "    tenant_id: str      # 租户标识（多租户场景）\n",
    "\n",
    "\n",
    "@before_model(state_schema=ExportAgentState)\n",
    "def inject_context(state: ExportAgentState, runtime: Runtime[RunContext]):\n",
    "    \"\"\"从 runtime.context 读取业务信息，用于日志关联。\"\"\"\n",
    "    ctx = runtime.context or {}\n",
    "    print(\n",
    "        f\"[上下文] request_id={ctx.get('request_id')} \"\n",
    "        f\"user_role={ctx.get('user_role')} \"\n",
    "        f\"tenant={ctx.get('tenant_id')}\"\n",
    "    )\n",
    "    return None\n",
    "\n",
    "\n",
    "# 创建 Agent 时声明 context 结构\n",
    "agent = create_agent(\n",
    "    model=model,\n",
    "    tools=tools,\n",
    "    middleware=[audit_before_model, inject_context, update_stats],\n",
    "    state_schema=ExportAgentState,  # 声明自定义状态\n",
    "    context_schema=RunContext,      # 声明上下文结构\n",
    "    system_prompt=\"你是报表导出平台的智能助手。\",\n",
    ")\n",
    "\n",
    "# 调用时传入 context\n",
    "result = agent.invoke(\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": \"导出本月的订单报表\"}]},\n",
    "    context={\n",
    "        \"request_id\": \"req-20260706-001\",\n",
    "        \"user_role\": \"finance\",\n",
    "        \"tenant_id\": \"t-1234\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7907d3f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents.middleware import ModelCallLimitMiddleware\n",
    "\n",
    "\n",
    "# 模拟一个\"前 4 次返回处理中，第 5 次返回完成\"的进度查询工具\n",
    "export_progress = {\"task_001\": 0}\n",
    "\n",
    "\n",
    "@tool\n",
    "def check_export_progress(task_id: str) -> str:\n",
    "    \"\"\"查询导出任务的进度。\"\"\"\n",
    "    export_progress[task_id] = export_progress.get(task_id, 0) + 1\n",
    "    attempt = export_progress[task_id]\n",
    "    if attempt < 5:\n",
    "        return f\"第 {attempt} 次查询：任务仍在处理中，请稍候。\"\n",
    "    return f\"第 {attempt} 次查询：导出完成，文件已生成。\"\n",
    "\n",
    "\n",
    "# 不加限制：Agent 会一直查到第 5 次\n",
    "# 加上 ModelCallLimitMiddleware(run_limit=3)：最多调 3 次模型就强制结束\n",
    "agent = create_agent(\n",
    "    model=model,\n",
    "    tools=[check_export_progress],\n",
    "    middleware=[\n",
    "        ModelCallLimitMiddleware(\n",
    "            run_limit=3,          # 单次 invoke 最多调 3 次模型\n",
    "            exit_behavior=\"end\",  # 到达上限后尝试优雅结束（生成总结）\n",
    "        ),\n",
    "    ],\n",
    "    system_prompt=(\n",
    "        \"你是导出任务进度查询助手。\"\n",
    "        \"当用户查询任务进度时，调用 check_export_progress。\"\n",
    "        \"如果任务还在处理中，继续查询直到完成。\"\n",
    "    ),\n",
    ")\n",
    "\n",
    "result = agent.invoke({\n",
    "    \"messages\": [{\"role\": \"user\", \"content\": \"帮我查 task_001 的导出进度，持续查到完成为止。\"}]\n",
    "})\n",
    "# 输出类似：\"Model call limits exceeded: run limit (3/3)\"\n",
    "# Agent 被强制刹车，不会无限循环\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "763486a0",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'tool' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mNameError\u001b[39m                                 Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[1]\u001b[39m\u001b[32m, line 21\u001b[39m\n\u001b[32m     17\u001b[39m     )\n\u001b[32m     18\u001b[39m \n\u001b[32m     19\u001b[39m \n\u001b[32m     20\u001b[39m \u001b[38;5;66;03m# 第二步：安全工具不加 args_schema（自动执行）\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m21\u001b[39m @tool\n\u001b[32m     22\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m check_export_permission(username: str) -> dict:\n\u001b[32m     23\u001b[39m     \u001b[33m\"\"\"查询用户导出权限。安全操作，无需审批。\"\"\"\u001b[39m\n\u001b[32m     24\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m {\u001b[33m\"username\"\u001b[39m: username, \u001b[33m\"can_export\"\u001b[39m: \u001b[38;5;28;01mTrue\u001b[39;00m, \u001b[33m\"region\"\u001b[39m: \u001b[33m\"all\"\u001b[39m}\n",
      "\u001b[31mNameError\u001b[39m: name 'tool' is not defined"
     ]
    }
   ],
   "source": [
    "from typing import Literal\n",
    "from pydantic import BaseModel, Field\n",
    "from langchain.agents.middleware import HumanInTheLoopMiddleware\n",
    "from langgraph.checkpoint.memory import InMemorySaver\n",
    "\n",
    "\n",
    "# 第一步：用 Pydantic 定义工具的输入结构\n",
    "# 结构化的参数让模型生成更准确，也方便人工审批时查看\n",
    "class ExportTaskInput(BaseModel):\n",
    "    \"\"\"导出任务的参数结构。模型会按照这个 schema 生成参数。\"\"\"\n",
    "    report_name: str = Field(description=\"报表名称，如 'east_region_orders_202607'\")\n",
    "    file_format: Literal[\"xlsx\", \"csv\"] = Field(description=\"导出格式\")\n",
    "    estimated_rows: int = Field(description=\"预计导出行数\")\n",
    "    reason: str = Field(description=\"导出原因，用于审计\")\n",
    "    priority: Literal[\"low\", \"normal\", \"high\"] = Field(\n",
    "        default=\"normal\", description=\"优先级\"\n",
    "    )\n",
    "\n",
    "\n",
    "# 第二步：安全工具不加 args_schema（自动执行）\n",
    "@tool\n",
    "def check_export_permission(username: str) -> dict:\n",
    "    \"\"\"查询用户导出权限。安全操作，无需审批。\"\"\"\n",
    "    return {\"username\": username, \"can_export\": True, \"region\": \"all\"}\n",
    "\n",
    "\n",
    "# 第三步：危险工具加上 args_schema，配合 HITL 拦截\n",
    "@tool(args_schema=ExportTaskInput)\n",
    "def create_export_task(\n",
    "    report_name: str,\n",
    "    file_format: str,\n",
    "    estimated_rows: int,\n",
    "    reason: str,\n",
    "    priority: str = \"normal\",\n",
    ") -> dict:\n",
    "    \"\"\"创建报表导出任务。这是一个高风险操作，需要人工审批。\"\"\"\n",
    "    print(\n",
    "        f\"[导出] 创建任务：{report_name} | {file_format} | \"\n",
    "        f\"{estimated_rows} 行 | 优先级 {priority}\"\n",
    "    )\n",
    "    return {\n",
    "        \"task_id\": \"EXPORT-20260706-002\",\n",
    "        \"report_name\": report_name,\n",
    "        \"file_format\": file_format,\n",
    "        \"estimated_rows\": estimated_rows,\n",
    "        \"status\": \"queued\",\n",
    "    }\n",
    "\n",
    "\n",
    "# 第四步：创建带 HITL 的 Agent\n",
    "checkpointer = InMemorySaver()  # HITL 必需：暂停后需要从这里恢复状态\n",
    "\n",
    "agent = create_agent(\n",
    "    model=model,\n",
    "    tools=[check_export_permission, create_export_task],\n",
    "    middleware=[\n",
    "        HumanInTheLoopMiddleware(\n",
    "            interrupt_on={\n",
    "                # 安全工具：不中断，自动执行\n",
    "                \"check_export_permission\": False,\n",
    "                # 危险工具：中断，提供三种审批选项\n",
    "                \"create_export_task\": {\n",
    "                    \"allowed_decisions\": [\"approve\", \"edit\", \"reject\"],\n",
    "                },\n",
    "            },\n",
    "        ),\n",
    "    ],\n",
    "    checkpointer=checkpointer,\n",
    "    system_prompt=(\n",
    "        \"你是报表导出平台的智能助手。\"\n",
    "        \"用户查询权限时，直接调用 check_export_permission。\"\n",
    "        \"只在用户明确要求创建导出任务时调用 create_export_task。\"\n",
    "    ),\n",
    ")\n",
    "\n",
    "# 第五步：第一次 invoke——会被 HITL 拦截\n",
    "config = {\"configurable\": {\"thread_id\": \"export-001\"}}\n",
    "result = agent.invoke(\n",
    "    {\n",
    "        \"messages\": [\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": \"我是 zhangsan，导出 east 区本月订单报表，xlsx，约 5000 行，月度对账用。\",\n",
    "            }\n",
    "        ]\n",
    "    },\n",
    "    config=config,\n",
    ")\n",
    "\n",
    "# result 里会包含中断信息，UI 层可以据此展示审批界面\n",
    "print(\"Agent 已暂停，等待审批...\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.14.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}