{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example_SummaryOfConversation\n",
    "\n",
    "对话历史的管理, 处理\n",
    "\n",
    "参考了: https://langchain-ai.github.io/langgraph/how-tos/memory/manage-conversation-history/#build-the-agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "Author: yangfh\n",
    "Date: 2024-10-24 14\n",
    "LastEditors: yangfh\n",
    "LastEditTime: 2024-10-25 10\n",
    "Description: \n",
    "\n",
    "Copyright (c) 2024 by www.simae.cn, All Rights Reserved. \n",
    "'''\n",
    "import os\n",
    "\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "from typing import Annotated, Literal, TypedDict\n",
    "from langchain_core.messages import SystemMessage,HumanMessage, RemoveMessage\n",
    "from langchain_core.tools import tool\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph import END, START, StateGraph, MessagesState\n",
    "from langgraph.prebuilt import ToolNode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# langsmith  # langsmith 调试时使用 (在线且收费的)\n",
    "\n",
    "# from langsmith.wrappers import wrap_openai\n",
    "# from langsmith import traceable\n",
    "# os.environ[\"LANGSMITH_TRACING\"]=\"true\"\n",
    "# os.environ[\"LANGSMITH_API_KEY\"]=\"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\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph 流转的状态对象\n",
    "\n",
    "添加一个 summary 属性作为总结的对话概要"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We will add a `summary` attribute (in addition to `messages` key,  which MessagesState already has)\n",
    "class MyGraphState(MessagesState):\n",
    "    summary: str\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LLM定义"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = 'sk-4c48484ff93f452d8e3bb41192452e01'\n",
    "llm_model = ChatOpenAI(model=\"qwen-turbo\",base_url=\"https://dashscope.aliyuncs.com/compatible-mode/v1\", streaming=True)\n",
    "# 本地\n",
    "# os.environ[\"OPENAI_API_KEY\"] = 'EMPTY'\n",
    "# llm_model = ChatOpenAI(model=\"chatglm3-6b\",base_url=\"http://172.16.21.155:8000/v1/chat/completions\", streaming=False)\n",
    "\n",
    "bound_model = llm_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 路由节点\n",
    "\n",
    "若是消息长度大于6 则路由到 summarize_conversation 节点, 去总结概要"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We now define the logic for determining whether to end or summarize the conversation\n",
    "def should_continue(state: MyGraphState) -> Literal[\"summarize_conversation\", END]:\n",
    "    \"\"\"Return the next node to execute.\"\"\"\n",
    "    messages = state[\"messages\"]\n",
    "    # If there are more than six messages, then we summarize the conversation\n",
    "    if len(messages) > 6:\n",
    "        return \"summarize_conversation\"\n",
    "    # Otherwise we can just end\n",
    "    return END"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 总结消息概要"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def summarize_conversation(state: MyGraphState):\n",
    "    # First, we summarize the conversation\n",
    "    summary = state.get(\"summary\", \"\")\n",
    "    if summary:\n",
    "        # If a summary already exists, we use a different system prompt\n",
    "        # to summarize it than if one didn't\n",
    "        summary_message = (\n",
    "            f\"这是此前对话摘要: {summary}\\n\\n\"\n",
    "            \"请考虑到此前的对话摘要加上述对话记录, 创建为一个新对话摘要\"\n",
    "        )\n",
    "    else:\n",
    "        summary_message = \"请将上述的对话创建为摘要\"\n",
    "    # 注意, 这里是插到最后面\n",
    "    messages = state[\"messages\"] + [HumanMessage(content=summary_message)]\n",
    "    response = llm_model.invoke(messages)\n",
    "    # 保留最新的2条消息, 删除其余的所有消息\n",
    "    delete_messages = [RemoveMessage(id=m.id) for m in state[\"messages\"][:-2]]\n",
    "    return {\"summary\": response.content, \"messages\": delete_messages} # 这个 messages(delete message 由langchain处理)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LLM调用节点\n",
    "\n",
    "调用模型时如果状态数据中, 若是有summary概要数据, 则在调用模型时, 替换为一个概要数据的系统消息 + 一个用户消息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the logic to call the model\n",
    "def call_model(state: MyGraphState):\n",
    "    # If a summary exists, we add this in as a system message\n",
    "    summary = state.get(\"summary\", \"\")\n",
    "    if summary:\n",
    "        system_message = f\"此前的对话摘要: {summary}\"\n",
    "        messages = [SystemMessage(content=system_message)] + state[\"messages\"]\n",
    "    else:\n",
    "        messages = state[\"messages\"]\n",
    "    response = bound_model.invoke(messages)\n",
    "    # We return a list, because this will get added to the existing list\n",
    "    return {\"messages\": [response]}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 图(Graph)定义"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "memory = MemorySaver()\n",
    "\n",
    "# Define a new graph\n",
    "workflow = StateGraph(MyGraphState)\n",
    "\n",
    "# Define the conversation node and the summarize node\n",
    "workflow.add_node(\"conversation\", call_model)\n",
    "workflow.add_node(summarize_conversation)\n",
    "\n",
    "# Set the entrypoint as conversation\n",
    "workflow.add_edge(START, \"conversation\")\n",
    "\n",
    "# We now add a conditional edge\n",
    "workflow.add_conditional_edges(\n",
    "    # First, we define the start node. We use `conversation`.\n",
    "    # This means these are the edges taken after the `conversation` node is called.\n",
    "    \"conversation\",\n",
    "    # Next, we pass in the function that will determine which node is called next.\n",
    "    should_continue,\n",
    ")\n",
    "\n",
    "# We now add a normal edge from `summarize_conversation` to END.\n",
    "# This means that after `summarize_conversation` is called, we end.\n",
    "workflow.add_edge(\"summarize_conversation\", END)\n",
    "\n",
    "# Finally, we compile it!\n",
    "app = workflow.compile(checkpointer=memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 调用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "你好, 我叫张三\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "你好，张三！很高兴认识你。有什么我可以帮你的吗？\n",
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "你叫什么名字?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "我是Qwen，是阿里云开发的大型语言模型。你可以问我任何问题或者告诉我需要帮助的地方。\n",
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "我今天和昨天晒了渔网, 接下来的三天我就要去打鱼了\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "听起来你已经做好了充分的准备去打鱼！晒渔网是为了确保它们在使用时更加耐用和有效。希望你接下来的捕鱼之旅能够收获满满。如果在捕鱼过程中有任何技巧或设备方面的问题，我很乐意提供帮助。祝你好运！\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "def print_update(update):\n",
    "    for k, v in update.items():\n",
    "        for m in v[\"messages\"]:\n",
    "            m.pretty_print()\n",
    "        if \"summary\" in v:\n",
    "            print(v[\"summary\"])\n",
    "            \n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": \"4\"}}\n",
    "input_message = HumanMessage(content=\"你好, 我叫张三\")\n",
    "input_message.pretty_print()\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
    "    print_update(event)\n",
    "\n",
    "input_message = HumanMessage(content=\"你叫什么名字?\")\n",
    "input_message.pretty_print()\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
    "    print_update(event)\n",
    "\n",
    "input_message = HumanMessage(content=\"我今天和昨天晒了渔网, 接下来的三天我就要去打鱼了\")\n",
    "input_message.pretty_print()\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
    "    print_update(event)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "我昨天干了什么?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "根据你之前的信息，昨天你晒了渔网。希望你的渔网现在已经晒干并且准备好用于接下来的捕鱼活动。如果你还有其他问题或需要进一步的帮助，请告诉我！\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "================================\u001b[1m Remove Message \u001b[0m================================\n",
      "\n",
      "\n",
      "好的，以下是你们上述对话的摘要：\n",
      "\n",
      "- 用户名：张三\n",
      "- 张三提到他昨天和前天晒了渔网，并且接下来的三天要去打鱼。\n",
      "- Qwen提醒张三晒渔网的重要性，并祝愿他捕鱼成功。\n",
      "- 张三询问Qwen的名字，得知Qwen是阿里云开发的大型语言模型。\n",
      "- 最后，张三询问昨天做了什么，Qwen回复说他昨天晒了渔网。\n",
      "\n",
      "如果有其他需求或更多信息需要添加到摘要中，请告诉我！\n"
     ]
    }
   ],
   "source": [
    "# 上面刚刚好 6条消息\n",
    "\n",
    "# 再执行会触发 消息总结\n",
    "input_message = HumanMessage(content=\"我昨天干了什么?\")\n",
    "input_message.pretty_print()\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
    "    print_update(event)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "我昨天干了什么?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "你提到昨天和前天都晒了渔网。因此，昨天你也应该是去晒渔网了。\n"
     ]
    }
   ],
   "source": [
    "# 再执行会触发 消息总结\n",
    "input_message = HumanMessage(content=\"我昨天干了什么?\")\n",
    "input_message.pretty_print()\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"updates\"):\n",
    "    print_update(event)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain0.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.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}