⚡ 短期记忆深度

class BaseCheckpointSaver(Generic[V]):
    """checkpointer 基类。让 agent 在多次交互间持久化状态。

    thread_id 是存取 checkpoint 的主键。没有它，
    checkpointer 无法保存状态、从中断恢复、或启用时间旅行。
    """

    serde: SerializerProtocol = JsonPlusSerializer()  # 序列化器

    # ====== 四个核心方法 ======

    def get_tuple(self, config) -> CheckpointTuple | None:
        """① 取【最新】的 checkpoint（含元数据、父指针）。
        调用时机：每次 invoke/stream 开始时，加载上次状态。"""

    def put(self, config, checkpoint, metadata, new_versions) -> RunnableConfig:
        """② 存一个完整的 checkpoint（状态快照）。
        调用时机：每个 super-step 结束（LG5 的 Update 阶段）。
        返回新 checkpoint 的 config（含 checkpoint_id）。"""

    def put_writes(self, config, writes, task_id) -> None:
        """③ 存【单个任务的写入】（增量，未提交的状态）。
        调用时机：任务执行中实时写入，支持中断恢复（部分完成可续）。
        这是实现"断点续跑"的关键——不必等整个 step 完成。"""

    def list(self, config, *, filter=None, before=None, limit=None):
        """④ 列出某 thread 的【历史 checkpoint 序列】。
        调用时机：时间旅行调试、回溯分析。
        支持 filter（按元数据过滤）、before（某点之前）、limit（数量）。"""

    # get 是 get_tuple 的便捷封装（只取 .checkpoint）
    def get(self, config) -> Checkpoint | None:
        if value := self.get_tuple(config):
            return value.checkpoint

class CheckpointTuple(NamedTuple):
    config: RunnableConfig          # 本 checkpoint 的 config（含 checkpoint_id）
    checkpoint: Checkpoint          # 状态快照本体
    metadata: CheckpointMetadata     # 元数据（step/source/writes/parents）
    parent_config: RunnableConfig | None   # ★ 父 checkpoint 指针（构成历史链）
    pending_writes: list             # 待写入（未完成的任务）

class Checkpoint(TypedDict):
    v: int                           # 版本号
    id: str                          # checkpoint 唯一 ID
    ts: str                          # 时间戳
    channel_values: dict             # ★ 各 channel 的当前值（就是状态）
    channel_versions: dict           # 各 channel 的版本（追踪变化）
    versions_seen: dict              # 各节点已看到的版本（决定是否要执行）
    pending_sends: list              # 待处理的 Send

class CheckpointMetadata(TypedDict):
    source: Literal["input", "loop", "update"]   # 来源：输入/循环/手动更新
    step: int                          # 第几步
    writes: dict                       # 本步的写入
    parents: list                      # 父 checkpoint IDs（分支用）

# pip install langgraph
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict

class State(TypedDict):
    count: int

def step(state):
    return {"count": state["count"] + 1}

g = StateGraph(State)
g.add_node("step", step)
g.add_edge(START, "step")
g.add_edge("step", "step")  # 自循环（演示用，实际要加退出条件）
app = g.compile(checkpointer=InMemorySaver())

config = {"configurable": {"thread_id": "t1"}}

# 跑几步
for _ in range(3):
    app.invoke({"count": 0} if _ == 0 else None, config)

# ★ 时间旅行：列出所有历史状态
print("=== 历史状态 ===")
states = list(app.get_state_history(config))
for s in reversed(states):
    print(f"step {s.metadata.get('step')}: count={s.values.get('count')}")

# ★ 回溯到第 1 步，从那里【分叉】重新跑（不影响原历史）
print("\n=== 从 step 1 分叉重跑 ===")
step1_config = states[-2].config   # 第1步的 config
result = app.invoke(None, step1_config)   # 从该 checkpoint 继续
print("分叉后的 count:", result["count"])
# 原历史还在，这是【新分支】——状态树！

from langchain_core.messages import trim_messages
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage

messages = [
    SystemMessage(content="你是助手"),
    HumanMessage(content="问题1"), AIMessage(content="答1"),
    HumanMessage(content="问题2"), AIMessage(content="答2"),
    HumanMessage(content="问题3"), AIMessage(content="答3"),
]

# ① 按 token 数裁剪（精确控量）
trimmed = trim_messages(
    messages,
    max_tokens=100,
    strategy="last",         # 保留最近的
    token_counter=len,       # 简化：用字符数当 token
    include_system=True,     # 始终保留 system
    start_on="human",        # 从 human 消息开始（不截断半句对话）
)

# ② 保留首尾（system + 最近 N 条）
trimmed2 = trim_messages(
    messages,
    max_tokens=100,
    strategy="last",
    include_system=True,
    allow_partial=False,     # 不截断单条消息
)

# ③ 用 LangGraph 的 pre_model_hook 在每次调 LLM 前自动裁剪
def auto_trim(state):
    return {"messages": trim_messages(
        state["messages"], max_tokens=4000, strategy="last",
        include_system=True, token_counter=model,
    )}

agent = create_react_agent(model, tools, pre_model_hook=auto_trim)
# 每次 LLM 调用前自动裁剪，无需手动管

# 假设用户在对话中途想"重来"，从某轮分叉
config_original = {"configurable": {"thread_id": "chat-1"}}

# 跑了5轮后，用户说"从第3轮重新开始，换个方向"
history = list(app.get_state_history(config_original))
round3_state = history[-3]   # 第3轮的状态（倒数第3个）

# 方法①：用原 thread 从 round3 继续（会覆盖后续）
app.update_state(round3_state.config, {"messages": [HumanMessage("新方向")]})
app.invoke(None, round3_state.config)

# 方法②：创建新 thread，复制 round3 状态（保留原历史）
config_fork = {"configurable": {"thread_id": "chat-1-fork"}}
# 把 round3 的状态写入新 thread，从此分叉独立发展
# 原 chat-1 的5轮历史完整保留，chat-1-fork 从第3轮另起

from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointTuple

class RedisCheckpointer(BaseCheckpointSaver):
    def __init__(self, redis_client):
        super().__init__()
        self.redis = redis_client

    def get_tuple(self, config):
        thread_id = config["configurable"]["thread_id"]
        data = self.redis.hgetall(f"checkpoint:{thread_id}:latest")
        if not data:
            return None
        # 反序列化返回 CheckpointTuple
        return self.serde.loads_typed(data)

    def put(self, config, checkpoint, metadata, new_versions):
        thread_id = config["configurable"]["thread_id"]
        # 序列化存储，同时记 parent 链
        key = f"checkpoint:{thread_id}:{checkpoint['id']}"
        self.redis.hset(key, mapping=self.serde.dumps_typed((checkpoint, metadata)))
        self.redis.set(f"checkpoint:{thread_id}:latest", checkpoint["id"])
        # ... put_writes 和 list 类似