State & Memory
Module Summary

# ── 依赖导入 ──────────────────────────────────────────────
from typing import Annotated, Optional, Literal
from typing_extensions import TypedDict
from pydantic import BaseModel, field_validator
from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, SystemMessage, RemoveMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_core.messages import trim_messages
from langgraph.checkpoint.sqlite import SqliteSaver
import sqlite3

# ════════════════════════════════════════════════════════
# 【L3】输入 Schema ── 图的公开入口，只暴露给调用者的字段
# ════════════════════════════════════════════════════════
class InputState(TypedDict):
    user_message: str   # 调用者只需传这一个字段

# ════════════════════════════════════════════════════════
# 【L3】输出 Schema ── 图的公开出口，过滤掉内部中间字段
# ════════════════════════════════════════════════════════
class OutputState(TypedDict):
    messages: list[AnyMessage]  # 调用者可见的对话历史
    summary:  str               # 调用者可见的当前摘要

# ════════════════════════════════════════════════════════
# 【L1 + L2】OverallState ── 图内部完整状态，包含所有字段
#   · TypedDict（L1）：零安装成本，字典风格访问
#   · Annotated + add_messages（L2）：消息追加 Reducer
# ════════════════════════════════════════════════════════
class ResearchState(TypedDict):
    # 【L2】add_messages Reducer：每轮对话追加，不覆盖
    messages: Annotated[list[AnyMessage], add_messages]

    # 【L5】摘要字段：超过阈值时由 summarize 节点写入
    summary:      str

    # 普通字段：默认覆盖策略（L2：无 Reducer 即覆盖）
    user_message: str

# ════════════════════════════════════════════════════════
# 【L1】Pydantic Schema ── 运行时强校验用户会话配置
#   · @field_validator 确保 session_id 合法
#   · 非法输入在进入图之前就被拦截
# ════════════════════════════════════════════════════════
class UserSession(BaseModel):
    session_id:   str
    max_messages: int = 6   # 触发摘要的消息阈值

    @field_validator("session_id")
    @classmethod
    def validate_session_id(cls, v):
        if not v or len(v) < 4:
            raise ValueError("session_id 至少需要 4 个字符")
        return v.strip()

# ════════════════════════════════════════════════════════
# 【L6】外部持久化配置
#   · SqliteSaver 将图的检查点（Checkpoint）写入 SQLite 文件
#   · 进程重启后，同一 thread_id 的对话历史可完整恢复
#   · 对比 MemorySaver：进程退出即全部丢失
# ════════════════════════════════════════════════════════
conn   = sqlite3.connect("research_assistant.db", check_same_thread=False)
memory = SqliteSaver(conn)   # 替换为 PostgresSaver 即可升级到生产级数据库

# LLM 初始化
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

def call_model(state: ResearchState) -> dict:
    """主对话节点：整合摘要上下文 + Token 控制 + LLM 调用"""
    summary  = state.get("summary", "")
    messages = state["messages"]

    # ── 【L4】trim_messages：按 Token 裁剪，只保留最近的消息 ──
    # strategy="last" 从最新消息开始保留，不满足 token 上限就截断
    # include_system=True 保留 SystemMessage（优先级最高）
    # 这里裁剪的是发给 LLM 的消息，State 中的完整历史不受影响
    trimmed = trim_messages(
        messages,
        max_tokens=1500,
        strategy="last",
        token_counter=llm,
        include_system=True,
        allow_partial=False,
    )

    # ── 【L5】将摘要注入 SystemMessage 作为背景上下文 ──
    # 当摘要存在时，LLM 能"记住"之前压缩掉的所有对话
    # 这解决了：消息被 RemoveMessage 删掉后如何保留语义信息
    system_content = "你是一个专业的智能研究助手，帮助用户深入探讨各类研究问题。"
    if summary:
        system_content += f"\n\n【对话摘要】以下是之前对话的压缩记录：\n{summary}"

    final_messages = [SystemMessage(content=system_content)] + trimmed
    response = llm.invoke(final_messages)

    # ── 【L2】add_messages Reducer 自动追加 response ──
    # 节点只需返回包含新消息的列表，Reducer 负责拼接到历史末尾
    return {"messages": [response]}

def should_summarize(state: ResearchState) -> Literal["summarize", "end"]:
    """条件边：判断消息数量是否触发摘要阈值"""
    # 超过 6 条消息就触发摘要节点，否则直接结束
    if len(state["messages"]) > 6:
        return "summarize"
    return "end"


def summarize_conversation(state: ResearchState) -> dict:
    """摘要节点：压缩对话历史，物理删除旧消息"""
    existing_summary = state.get("summary", "")
    messages         = state["messages"]

    # ── 【L5】构建摘要 Prompt ──
    # 若已有摘要则在其基础上扩展（滚动摘要模式）
    # 若没有则从头生成（首次压缩）
    if existing_summary:
        prompt = (
            f"当前已有摘要：\n{existing_summary}\n\n"
            "请在此基础上，将以下新对话内容也纳入摘要，保持信息完整简洁："
        )
    else:
        prompt = "请简洁总结以下对话的关键信息，保留重要的问答内容："

    summary_messages = messages + [HumanMessage(content=prompt)]
    response = llm.invoke(summary_messages)

    # ── 【L2】RemoveMessage：物理删除旧消息，只保留最近 2 条 ──
    # messages[:-2] 取"最后 2 条之前"的所有消息作为删除目标
    # add_messages Reducer 识别到 RemoveMessage 后，按 id 从列表中移除
    # 这与 trim_messages（仅过滤发给 LLM 的内容）不同：这里是真正修改 State
    delete_messages = [RemoveMessage(id=m.id) for m in messages[:-2]]

    return {
        "summary":  response.content,  # 覆盖旧摘要（默认覆盖策略，L2）
        "messages": delete_messages,   # add_messages Reducer 执行删除（L2）
    }

# ════════════════════════════════════════════════════════
# 【L3】多重 Schema 装配
#   · 第一参数 ResearchState：图内部的完整 Schema
#   · input=InputState：调用者只能传 user_message
#   · output=OutputState：图只返回 messages 和 summary
# ════════════════════════════════════════════════════════
builder = StateGraph(
    ResearchState,
    input=InputState,
    output=OutputState,
)

# 注册节点
builder.add_node("call_model",            call_model)
builder.add_node("summarize_conversation", summarize_conversation)

# 连接边
builder.add_edge(START, "call_model")

# 条件边：call_model 执行后，根据消息数决定走向
builder.add_conditional_edges(
    "call_model",
    should_summarize,
    {"summarize": "summarize_conversation", "end": END}
)
builder.add_edge("summarize_conversation", END)

# ════════════════════════════════════════════════════════
# 【L6】挂载外部检查点，compile 时传入 SqliteSaver
#   · 每个节点执行完后，LangGraph 自动将 State 快照存入数据库
#   · 只需替换 checkpointer=memory 即可切换存储后端
# ════════════════════════════════════════════════════════
graph = builder.compile(checkpointer=memory)

# ── 【L1】Pydantic 校验：确保 session_id 合法再启动 ──
session = UserSession(session_id="researcher_001", max_messages=6)
config  = {"configurable": {"thread_id": session.session_id}}

# ── 第 1 轮对话 ────────────────────────────────────────
result = graph.invoke(
    {"user_message": "深海生物是如何在极端高压环境中生存的？"},
    config=config,
)
# result 只包含 OutputState 的字段：messages 和 summary
print(result["messages"][-1].content)

# ── 第 2 轮：LangGraph 自动从 SQLite 恢复第 1 轮的状态 ──
result = graph.invoke(
    {"user_message": "这与极地生物的耐寒机制有何异同？"},
    config=config,
)

# ── 模拟进程重启后的跨会话恢复 ────────────────────────
# 重新建立数据库连接（模拟重启）
new_conn   = sqlite3.connect("research_assistant.db", check_same_thread=False)
new_memory = SqliteSaver(new_conn)
new_graph  = builder.compile(checkpointer=new_memory)

# 同一个 thread_id → 历史自动恢复，就像从未重启过
result = new_graph.invoke(
    {"user_message": "给我推荐几篇这个领域的经典论文"},
    config=config,  # thread_id 不变，记忆完整继承
)

# ── 多用户隔离示例 ──────────────────────────────────────
user_a_config = {"configurable": {"thread_id": "user_alice"}}
user_b_config = {"configurable": {"thread_id": "user_bob"}}
# alice 和 bob 的对话历史完全独立，互不干扰