深入LangGraph检查点系统：持久化与状态管理源码解析

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

class Checkpoint(TypedDict):
    """给定时间点的状态快照
    
    检查点是LangGraph执行状态的完整记录，包含：
    - 版本信息
    - 唯一标识符
    - 时间戳
    - 通道状态值
    - 通道版本号
    - 节点已见版本
    - 更新通道列表
    """
    
    v: int
    """检查点格式版本，当前为1
    
    用于处理格式升级和向后兼容性
    """
    
    id: str
    """检查点的唯一ID
    
    既唯一又单调递增，可用于从第一个到最后一个排序检查点
    使用uuid6生成，确保时间有序性
    """
    
    ts: str
    """检查点的时间戳，ISO 8601格式
    
    示例: "2024-01-15T10:30:45.123456Z"
    """
    
    channel_values: dict[str, Any]
    """检查点时通道的值
    
    从通道名到反序列化通道快照值的映射
    这是图状态的实际数据内容
    """
    
    channel_versions: ChannelVersions
    """检查点时通道的版本
    
    键是通道名，值是每个通道单调递增的版本字符串
    用于增量更新和变化检测
    """
    
    versions_seen: dict[str, ChannelVersions]
    """从节点ID到从通道名到版本的映射
    
    跟踪每个节点已见的通道版本
    用于确定接下来执行哪些节点
    格式: {node_id: {channel_name: version}}
    """
    
    updated_channels: list[str] | None
    """在此检查点中更新的通道列表
    
    用于优化，快速识别哪些通道发生了变化
    """

# 检查点复制函数
def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
    """深拷贝检查点对象
    
    Args:
        checkpoint: 要复制的检查点
        
    Returns:
        Checkpoint: 检查点的深拷贝
    """
    return Checkpoint(
        v=checkpoint["v"],
        ts=checkpoint["ts"],
        id=checkpoint["id"],
        channel_values=checkpoint["channel_values"].copy(),
        channel_versions=checkpoint["channel_versions"].copy(),
        versions_seen={
            k: v.copy() for k, v in checkpoint["versions_seen"].items()
        },
        pending_sends=checkpoint.get("pending_sends", []).copy(),
        updated_channels=checkpoint.get("updated_channels", None),
    )

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

class CheckpointMetadata(TypedDict, total=False):
    """与检查点关联的元数据"""
    
    source: Literal["input", "loop", "update", "fork"]
    """检查点的来源
    
    - "input": 从invoke/stream/batch的输入创建的检查点
    - "loop": 从pregel循环内部创建的检查点  
    - "update": 从手动状态更新创建的检查点
    - "fork": 作为另一个检查点副本创建的检查点
    """
    
    step: int
    """检查点的步骤编号
    
    -1 表示第一个"input"检查点
    0 表示第一个"loop"检查点
    ... 之后的第n个检查点
    """
    
    parents: dict[str, str]
    """父检查点的ID映射
    
    从检查点命名空间到检查点ID的映射
    用于支持分支和合并操作
    """

# 获取检查点元数据的辅助函数
def get_checkpoint_metadata(
    config: RunnableConfig,
    step: int = -1,
    source: Literal["input", "loop", "update", "fork"] = "input",
) -> CheckpointMetadata:
    """构造检查点元数据
    
    Args:
        config: 运行配置
        step: 步骤编号
        source: 检查点来源
        
    Returns:
        CheckpointMetadata: 构造的元数据
    """
    metadata: CheckpointMetadata = {
        "source": source,
        "step": step,
        "parents": {},
    }
    
    # 从配置中提取父检查点信息
    if "checkpoint_id" in config.get("configurable", {}):
        parent_id = config["configurable"]["checkpoint_id"]
        ns = config.get("configurable", {}).get("checkpoint_ns", "")
        metadata["parents"][ns] = parent_id
    
    return metadata

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

class CheckpointTuple(NamedTuple):
    """包含检查点及其关联数据的元组
    
    这是检查点系统的主要返回类型，包含了
    完整的检查点信息和相关的配置、元数据
    """
    
    config: RunnableConfig
    """与检查点关联的运行配置
    
    包含thread_id、checkpoint_id等标识信息
    """
    
    checkpoint: Checkpoint
    """检查点数据本身"""
    
    metadata: CheckpointMetadata
    """检查点元数据"""
    
    parent_config: RunnableConfig | None = None
    """父检查点的配置（如果有）
    
    用于支持检查点链和历史追踪
    """
    
    pending_writes: list[PendingWrite] | None = None
    """待写入操作列表
    
    当节点执行失败时，成功完成的节点的写入
    会保存为待写入，以便恢复执行时不重复运行
    """

# 待写入类型定义
PendingWrite = tuple[str, str, Any]  # (channel, task_id, value)

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189

class BaseCheckpointSaver(Generic[V]):
    """创建图检查点保存器的基类
    
    检查点保存器允许LangGraph智能体在多次交互中持久化状态
    
    Attributes:
        serde (SerializerProtocol): 用于编码/解码检查点的序列化器
        
    Note:
        创建自定义检查点保存器时，考虑实现异步版本以避免阻塞主线程
    """
    
    serde: SerializerProtocol = JsonPlusSerializer()
    """序列化协议实例，默认使用JsonPlus序列化器"""
    
    def __init__(
        self,
        *,
        serde: SerializerProtocol | None = None,
    ) -> None:
        """初始化检查点保存器
        
        Args:
            serde: 序列化器实例，如果为None则使用默认序列化器
        """
        self.serde = maybe_add_typed_methods(serde or self.serde)
    
    @property
    def config_specs(self) -> list:
        """定义检查点保存器的配置选项
        
        Returns:
            list: 配置字段规格列表
        """
        return [
            {
                "name": "thread_id",
                "type": "string",
                "description": "线程唯一标识符",
                "required": True,
            },
            {
                "name": "checkpoint_id", 
                "type": "string",
                "description": "检查点唯一标识符",
                "required": False,
            },
            {
                "name": "checkpoint_ns",
                "type": "string", 
                "description": "检查点命名空间",
                "required": False,
                "default": "",
            },
        ]

    # 核心接口方法
    def get(self, config: RunnableConfig) -> Checkpoint | None:
        """使用给定配置获取检查点
        
        Args:
            config: 指定要检索哪个检查点的配置
            
        Returns:
            Optional[Checkpoint]: 请求的检查点，如果未找到则为None
        """
        if value := self.get_tuple(config):
            return value.checkpoint

    def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
        """使用给定配置获取检查点元组
        
        Args:
            config: 指定要检索哪个检查点的配置
            
        Returns:
            Optional[CheckpointTuple]: 请求的检查点元组，如果未找到则为None
            
        Raises:
            NotImplementedError: 在自定义检查点保存器中实现此方法
        """
        raise NotImplementedError

    def list(
        self,
        config: RunnableConfig | None,
        *,
        filter: dict[str, Any] | None = None,
        before: RunnableConfig | None = None,
        limit: int | None = None,
    ) -> Iterator[CheckpointTuple]:
        """列出与给定条件匹配的检查点
        
        Args:
            config: 基础配置，包含thread_id等
            filter: 过滤条件字典
            before: 返回此配置之前的检查点
            limit: 返回的最大检查点数量
            
        Yields:
            CheckpointTuple: 匹配的检查点元组
            
        Raises:
            NotImplementedError: 在自定义检查点保存器中实现此方法
        """
        raise NotImplementedError

    def put(
        self,
        config: RunnableConfig,
        checkpoint: Checkpoint,
        metadata: CheckpointMetadata,
        new_versions: ChannelVersions,
    ) -> RunnableConfig:
        """存储检查点及其配置和元数据
        
        Args:
            config: 与检查点关联的配置
            checkpoint: 要存储的检查点
            metadata: 检查点元数据
            new_versions: 新的通道版本
            
        Returns:
            RunnableConfig: 更新的配置（可能包含新的checkpoint_id）
            
        Raises:
            NotImplementedError: 在自定义检查点保存器中实现此方法
        """
        raise NotImplementedError

    def put_writes(
        self,
        config: RunnableConfig,
        writes: Sequence[tuple[str, Any]],
        task_id: str,
    ) -> None:
        """存储与检查点关联的中间写入
        
        Args:
            config: 与检查点关联的配置
            writes: 要存储的写入操作序列
            task_id: 执行写入的任务ID
            
        Note:
            中间写入用于在节点执行失败时保存成功节点的结果，
            以便恢复时不需要重新执行
        """
        raise NotImplementedError

    # 异步接口方法（可选实现）
    async def aget(self, config: RunnableConfig) -> Checkpoint | None:
        """get方法的异步版本"""
        if value := await self.aget_tuple(config):
            return value.checkpoint

    async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
        """get_tuple方法的异步版本"""
        raise NotImplementedError

    async def alist(
        self,
        config: RunnableConfig | None,
        *,
        filter: dict[str, Any] | None = None,
        before: RunnableConfig | None = None,
        limit: int | None = None,
    ) -> AsyncIterator[CheckpointTuple]:
        """list方法的异步版本"""
        raise NotImplementedError
        yield  # 使其成为异步生成器

    async def aput(
        self,
        config: RunnableConfig,
        checkpoint: Checkpoint,
        metadata: CheckpointMetadata,
        new_versions: ChannelVersions,
    ) -> RunnableConfig:
        """put方法的异步版本"""
        raise NotImplementedError

    async def aput_writes(
        self,
        config: RunnableConfig,
        writes: Sequence[tuple[str, Any]],
        task_id: str,
    ) -> None:
        """put_writes方法的异步版本"""
        raise NotImplementedError

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

def get_next_version(current: V | None, channel: BaseChannel) -> V:
    """获取通道的下一个版本号
    
    Args:
        current: 当前版本号
        channel: 通道实例
        
    Returns:
        V: 下一个版本号
    """
    if current is None:
        return 1
    elif isinstance(current, int):
        return current + 1
    elif isinstance(current, float):
        return current + 1.0
    elif isinstance(current, str):
        try:
            # 尝试解析为整数并递增
            return str(int(current) + 1)
        except ValueError:
            # 如果不能解析，生成新的UUID
            from langgraph.checkpoint.base.id import uuid6
            return str(uuid6())
    else:
        raise ValueError(f"Unsupported version type: {type(current)}")

def get_checkpoint_id(config: RunnableConfig, checkpoint: Checkpoint) -> str:
    """从配置或检查点获取检查点ID
    
    Args:
        config: 运行配置
        checkpoint: 检查点数据
        
    Returns:
        str: 检查点ID
    """
    # 首先尝试从配置获取
    if checkpoint_id := config.get("configurable", {}).get("checkpoint_id"):
        return checkpoint_id
    
    # 否则从检查点本身获取
    return checkpoint["id"]

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218

class JsonPlusSerializer(SerializerProtocol):
    """使用ormsgpack的序列化器，具有扩展JSON序列化的回退功能
    
    该序列化器支持：
    - MessagePack高效二进制编码
    - LangChain对象序列化
    - Pydantic模型序列化  
    - Python标准库类型
    - 自定义类型扩展
    """
    
    def __init__(
        self,
        *,
        pickle_fallback: bool = False,
        __unpack_ext_hook__: Callable[[int, bytes], Any] | None = None,
    ) -> None:
        """初始化序列化器
        
        Args:
            pickle_fallback: 是否使用pickle作为最后的回退方案
            __unpack_ext_hook__: MessagePack扩展类型解包钩子
        """
        self.pickle_fallback = pickle_fallback
        self._unpack_ext_hook = (
            __unpack_ext_hook__
            if __unpack_ext_hook__ is not None
            else _msgpack_ext_hook
        )
    
    def dumps(self, obj: Any) -> bytes:
        """将对象序列化为字节
        
        Args:
            obj: 要序列化的对象
            
        Returns:
            bytes: 序列化后的字节数据
        """
        try:
            # 首先尝试使用MessagePack
            return ormsgpack.packb(
                obj, 
                default=self._default,
                option=ormsgpack.OPT_NON_STR_KEYS | ormsgpack.OPT_SERIALIZE_DATACLASS
            )
        except (TypeError, ValueError) as e:
            if self.pickle_fallback:
                # 使用pickle作为回退
                return pickle.dumps(obj)
            else:
                raise ValueError(f"Failed to serialize object: {e}") from e
    
    def loads(self, data: bytes) -> Any:
        """从字节反序列化对象
        
        Args:
            data: 要反序列化的字节数据
            
        Returns:
            Any: 反序列化后的对象
        """
        if not data:
            return None
            
        try:
            # 首先尝试MessagePack解包
            return ormsgpack.unpackb(
                data, 
                ext_hook=self._unpack_ext_hook,
                str_hook=self._str_hook,
            )
        except (ormsgpack.MsgpackDecodeError, ValueError):
            if self.pickle_fallback:
                # 尝试pickle解包
                return pickle.loads(data)
            else:
                raise

    def _default(self, obj: Any) -> str | dict[str, Any]:
        """处理不能直接序列化的对象
        
        Args:
            obj: 要处理的对象
            
        Returns:
            str | dict: 可序列化的表示
        """
        # LangChain Serializable对象
        if isinstance(obj, Serializable):
            return cast(dict[str, Any], obj.to_json())
        
        # Pydantic模型（v2）
        elif hasattr(obj, "model_dump") and callable(obj.model_dump):
            return self._encode_constructor_args(
                obj.__class__, 
                method=(None, "model_construct"), 
                kwargs=obj.model_dump()
            )
        
        # Pydantic模型（v1）
        elif hasattr(obj, "dict") and callable(obj.dict):
            return self._encode_constructor_args(
                obj.__class__, 
                method=(None, "construct"), 
                kwargs=obj.dict()
            )
        
        # NamedTuple
        elif hasattr(obj, "_asdict") and callable(obj._asdict):
            return self._encode_constructor_args(
                obj.__class__, 
                kwargs=obj._asdict()
            )
        
        # 路径对象
        elif isinstance(obj, pathlib.Path):
            return self._encode_constructor_args(
                pathlib.Path, 
                args=obj.parts
            )
        
        # 正则表达式
        elif isinstance(obj, re.Pattern):
            return self._encode_constructor_args(
                re.compile, 
                args=(obj.pattern, obj.flags)
            )
        
        # UUID对象
        elif isinstance(obj, UUID):
            return self._encode_constructor_args(
                UUID, 
                args=(obj.hex,)
            )
        
        # Decimal对象
        elif isinstance(obj, decimal.Decimal):
            return self._encode_constructor_args(
                decimal.Decimal, 
                args=(str(obj),)
            )
        
        # 集合类型
        elif isinstance(obj, (set, frozenset, deque)):
            return self._encode_constructor_args(
                obj.__class__, 
                args=(list(obj),)
            )
        
        # 日期时间类型
        elif isinstance(obj, (datetime, date, time)):
            return self._encode_datetime(obj)
        
        # IP地址类型
        elif isinstance(obj, (IPv4Address, IPv6Address, IPv4Network, IPv6Network)):
            return self._encode_constructor_args(
                obj.__class__, 
                args=(str(obj),)
            )
        
        # 枚举类型
        elif isinstance(obj, Enum):
            return self._encode_constructor_args(
                obj.__class__, 
                args=(obj.value,)
            )
        
        # 数据类
        elif dataclasses.is_dataclass(obj):
            return self._encode_constructor_args(
                obj.__class__, 
                kwargs=dataclasses.asdict(obj)
            )
        
        # 如果启用了pickle回退
        elif self.pickle_fallback:
            # 使用pickle编码为MessagePack扩展类型
            return ormsgpack.packb(
                pickle.dumps(obj), 
                option=ormsgpack.OPT_NON_STR_KEYS
            )
        
        else:
            raise TypeError(f"Object of type {type(obj)} is not serializable")

    def _encode_constructor_args(
        self,
        constructor: Callable | type[Any],
        *,
        method: None | str | Sequence[None | str] = None,
        args: Sequence[Any] | None = None,
        kwargs: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """编码构造函数参数为可序列化格式
        
        这个格式与LangChain的序列化格式兼容，
        支持在反序列化时重建对象
        """
        out = {
            "lc": 2,  # LangChain序列化版本
            "type": "constructor",
            "id": (*constructor.__module__.split("."), constructor.__name__),
        }
        if method is not None:
            out["method"] = method
        if args is not None:
            out["args"] = args
        if kwargs is not None:
            out["kwargs"] = kwargs
        return out

    def _str_hook(self, obj: str) -> Any:
        """字符串反序列化钩子，处理特殊字符串格式"""
        # 处理LangChain序列化标记
        if obj.startswith("lc://"):
            return LC_REVIVER(json.loads(obj[5:]))
        return obj

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182

class OptimizedJsonPlusSerializer(JsonPlusSerializer):
    """优化的JsonPlus序列化器：针对大规模数据的性能优化"""
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.compression_threshold = 1024  # 1KB压缩阈值
        self.lru_cache = {}
        self.cache_size = 10000
        self.type_registry = {}  # 类型注册表
        
    def dumps_optimized(self, obj: Any, compress: bool = True) -> bytes:
        """优化的序列化方法"""
        # 1. 类型检查和快速路径
        if self._is_simple_type(obj):
            return self._serialize_simple(obj)
        
        # 2. 检查缓存
        obj_hash = self._compute_hash(obj)
        if obj_hash in self.lru_cache:
            return self.lru_cache[obj_hash]
        
        # 3. 常规序列化
        serialized = super().dumps(obj)
        
        # 4. 压缩处理
        if compress and len(serialized) > self.compression_threshold:
            import zlib
            compressed = zlib.compress(serialized, level=6)
            if len(compressed) < len(serialized) * 0.8:  # 压缩率阈值
                serialized = b'\x01' + compressed  # 添加压缩标记
        
        # 5. 更新缓存
        self._update_cache(obj_hash, serialized)
        
        return serialized
    
    def loads_optimized(self, data: bytes) -> Any:
        """优化的反序列化方法"""
        if not data:
            return None
        
        # 检查压缩标记
        if data[0:1] == b'\x01':
            import zlib
            data = zlib.decompress(data[1:])
        
        # 反序列化
        return super().loads(data)
    
    def _is_simple_type(self, obj: Any) -> bool:
        """检查是否为简单类型（可快速序列化）"""
        return isinstance(obj, (int, float, str, bool, type(None)))
    
    def _serialize_simple(self, obj: Any) -> bytes:
        """简单类型的快速序列化"""
        return json.dumps(obj).encode('utf-8')
    
    def _compute_hash(self, obj: Any) -> str:
        """计算对象哈希用于缓存"""
        try:
            return hashlib.sha256(
                json.dumps(obj, sort_keys=True, default=str).encode()
            ).hexdigest()[:16]  # 取前16位作为缓存键
        except:
            return str(id(obj))
    
    def _update_cache(self, key: str, value: bytes):
        """更新LRU缓存"""
        if len(self.lru_cache) >= self.cache_size:
            # 移除最老的缓存项
            oldest_key = next(iter(self.lru_cache))
            del self.lru_cache[oldest_key]
        
        self.lru_cache[key] = value

# 智能类型注册机制
class TypeRegistryManager:
    """类型注册管理器：智能识别和处理自定义类型"""
    
    def __init__(self):
        self.type_handlers = {}
        self.auto_discovery = True
        self.reflection_cache = {}
    
    def register_type_handler(
        self, 
        type_class: type, 
        serializer: Callable, 
        deserializer: Callable
    ):
        """注册自定义类型处理器"""
        self.type_handlers[type_class] = {
            "serialize": serializer,
            "deserialize": deserializer,
            "registered_at": time.time(),
        }
    
    def auto_register_pydantic_models(self, module_name: str):
        """自动注册模块中的Pydantic模型"""
        import importlib
        import inspect
        from pydantic import BaseModel
        
        try:
            module = importlib.import_module(module_name)
            for name, obj in inspect.getmembers(module):
                if (inspect.isclass(obj) and 
                    issubclass(obj, BaseModel) and 
                    obj != BaseModel):
                    
                    self.register_type_handler(
                        obj,
                        lambda instance: instance.model_dump_json(),
                        lambda data: obj.model_validate_json(data)
                    )
        except ImportError:
            pass

# 增量序列化机制
class IncrementalSerializer:
    """增量序列化器：支持大型状态对象的增量序列化"""
    
    def __init__(self, base_serializer: SerializerProtocol):
        self.base_serializer = base_serializer
        self.state_snapshots = {}
        self.change_tracker = {}
    
    def serialize_incremental(
        self, 
        obj: Any, 
        obj_id: str, 
        force_full: bool = False
    ) -> tuple[bytes, bool]:
        """增量序列化
        
        Returns:
            tuple: (序列化数据, 是否为增量数据)
        """
        if force_full or obj_id not in self.state_snapshots:
            # 全量序列化
            serialized = self.base_serializer.dumps(obj)
            self.state_snapshots[obj_id] = obj
            return serialized, False
        
        # 计算增量
        previous_obj = self.state_snapshots[obj_id]
        changes = self._compute_changes(previous_obj, obj)
        
        if len(changes) > len(str(obj)) * 0.5:  # 变化太大，使用全量
            serialized = self.base_serializer.dumps(obj)
            self.state_snapshots[obj_id] = obj
            return serialized, False
        else:
            # 增量序列化
            delta_data = {
                "type": "incremental",
                "base_id": obj_id,
                "changes": changes,
                "timestamp": time.time(),
            }
            self.state_snapshots[obj_id] = obj
            return self.base_serializer.dumps(delta_data), True
    
    def deserialize_incremental(
        self, 
        data: bytes, 
        obj_id: str
    ) -> Any:
        """增量反序列化"""
        deserialized = self.base_serializer.loads(data)
        
        if isinstance(deserialized, dict) and deserialized.get("type") == "incremental":
            # 增量数据
            base_obj = self.state_snapshots.get(obj_id)
            if base_obj is None:
                raise ValueError(f"Base object for {obj_id} not found")
            
            return self._apply_changes(base_obj, deserialized["changes"])
        else:
            # 全量数据
            self.state_snapshots[obj_id] = deserialized
            return deserialized

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

def _msgpack_ext_hook(code: int, data: bytes) -> Any:
    """MessagePack扩展类型钩子
    
    Args:
        code: 扩展类型代码
        data: 扩展数据
        
    Returns:
        Any: 反序列化后的对象
    """
    if code == 1:  # datetime
        return datetime.fromisoformat(data.decode())
    elif code == 2:  # date
        return date.fromisoformat(data.decode())
    elif code == 3:  # time
        return time.fromisoformat(data.decode())
    elif code == 4:  # timedelta
        return timedelta(seconds=float(data.decode()))
    elif code == 5:  # timezone
        return ZoneInfo(data.decode())
    elif code == 6:  # decimal
        return decimal.Decimal(data.decode())
    elif code == 7:  # uuid
        return UUID(data.decode())
    elif code == 8:  # pickle fallback
        return pickle.loads(data)
    else:
        raise ValueError(f"Unknown extension type: {code}")

def _encode_datetime(obj: datetime | date | time) -> dict[str, Any]:
    """编码日期时间对象
    
    Args:
        obj: 日期时间对象
        
    Returns:
        dict: 编码后的字典
    """
    if isinstance(obj, datetime):
        return {
            "__msgpack_ext__": {
                "code": 1,
                "data": obj.isoformat().encode(),
            }
        }
    elif isinstance(obj, date):
        return {
            "__msgpack_ext__": {
                "code": 2,
                "data": obj.isoformat().encode(),
            }
        }
    elif isinstance(obj, time):
        return {
            "__msgpack_ext__": {
                "code": 3,
                "data": obj.isoformat().encode(),
            }
        }
    else:
        raise TypeError(f"Unsupported datetime type: {type(obj)}")

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313

class InMemorySaver(
    BaseCheckpointSaver[str], 
    AbstractContextManager, 
    AbstractAsyncContextManager
):
    """内存中的检查点保存器
    
    此检查点保存器使用defaultdict在内存中存储检查点
    
    Note:
        只能用于调试或测试目的
        对于生产用例，我们推荐安装langgraph-checkpoint-postgres
        并使用PostgresSaver / AsyncPostgresSaver
        
        如果您使用LangGraph平台，无需指定检查点保存器
        将自动使用正确的托管检查点保存器
    """
    
    def __init__(
        self,
        *,
        serde: SerializerProtocol | None = None,
    ) -> None:
        """初始化内存保存器
        
        Args:
            serde: 序列化器，默认为None使用JsonPlusSerializer
        """
        super().__init__(serde=serde)
        
        # 存储结构：thread_id -> checkpoint_ns -> checkpoint_id -> (checkpoint, metadata)
        self.storage: defaultdict[
            str, defaultdict[str, dict[str, tuple[Checkpoint, CheckpointMetadata]]]
        ] = defaultdict(lambda: defaultdict(dict))
        
        # 待写入存储：thread_id -> checkpoint_ns -> checkpoint_id -> task_id -> writes
        self.writes: defaultdict[
            str, defaultdict[str, defaultdict[str, defaultdict[str, list[tuple[str, Any]]]]]
        ] = defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(list))))
    
    def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
        """获取检查点元组
        
        Args:
            config: 运行配置，必须包含thread_id
            
        Returns:
            CheckpointTuple | None: 检查点元组或None
        """
        # 提取配置参数
        thread_id = config["configurable"]["thread_id"]
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
        checkpoint_id = config["configurable"].get("checkpoint_id")
        
        # 获取线程的存储
        thread_checkpoints = self.storage.get(thread_id, {}).get(checkpoint_ns, {})
        
        if checkpoint_id is not None:
            # 获取特定检查点
            if checkpoint_id in thread_checkpoints:
                checkpoint, metadata = thread_checkpoints[checkpoint_id]
                return CheckpointTuple(
                    config=config,
                    checkpoint=checkpoint,
                    metadata=metadata,
                    parent_config=self._get_parent_config(config, metadata),
                    pending_writes=self._get_pending_writes(config),
                )
        else:
            # 获取最新检查点
            if thread_checkpoints:
                # 按时间戳排序获取最新的
                latest_id = max(
                    thread_checkpoints.keys(),
                    key=lambda x: thread_checkpoints[x][0]["ts"]
                )
                checkpoint, metadata = thread_checkpoints[latest_id]
                updated_config = {
                    **config,
                    "configurable": {
                        **config["configurable"],
                        "checkpoint_id": latest_id,
                    }
                }
                return CheckpointTuple(
                    config=updated_config,
                    checkpoint=checkpoint,
                    metadata=metadata,
                    parent_config=self._get_parent_config(updated_config, metadata),
                    pending_writes=self._get_pending_writes(updated_config),
                )
        
        return None
    
    def list(
        self,
        config: RunnableConfig | None,
        *,
        filter: dict[str, Any] | None = None,
        before: RunnableConfig | None = None,
        limit: int | None = None,
    ) -> Iterator[CheckpointTuple]:
        """列出检查点
        
        Args:
            config: 基础配置，包含thread_id
            filter: 过滤条件
            before: 返回此配置之前的检查点
            limit: 最大返回数量
            
        Yields:
            CheckpointTuple: 匹配的检查点元组
        """
        if config is None:
            return
            
        thread_id = config["configurable"]["thread_id"]
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
        
        # 获取线程的所有检查点
        thread_checkpoints = self.storage.get(thread_id, {}).get(checkpoint_ns, {})
        
        if not thread_checkpoints:
            return
        
        # 构建检查点列表
        checkpoints = []
        for checkpoint_id, (checkpoint, metadata) in thread_checkpoints.items():
            # 应用过滤器
            if filter:
                if not self._matches_filter(metadata, filter):
                    continue
            
            # 应用before过滤
            if before:
                before_ts = before.get("configurable", {}).get("checkpoint_ts")
                if before_ts and checkpoint["ts"] >= before_ts:
                    continue
            
            checkpoint_config = {
                **config,
                "configurable": {
                    **config["configurable"],
                    "checkpoint_id": checkpoint_id,
                }
            }
            
            checkpoints.append(CheckpointTuple(
                config=checkpoint_config,
                checkpoint=checkpoint,
                metadata=metadata,
                parent_config=self._get_parent_config(checkpoint_config, metadata),
                pending_writes=self._get_pending_writes(checkpoint_config),
            ))
        
        # 按时间戳排序（最新的在前）
        checkpoints.sort(key=lambda x: x.checkpoint["ts"], reverse=True)
        
        # 应用限制
        if limit is not None:
            checkpoints = checkpoints[:limit]
        
        yield from checkpoints
    
    def put(
        self,
        config: RunnableConfig,
        checkpoint: Checkpoint,
        metadata: CheckpointMetadata,
        new_versions: ChannelVersions,
    ) -> RunnableConfig:
        """存储检查点
        
        Args:
            config: 运行配置
            checkpoint: 检查点数据
            metadata: 检查点元数据
            new_versions: 新的通道版本
            
        Returns:
            RunnableConfig: 更新的配置
        """
        # 提取配置参数
        thread_id = config["configurable"]["thread_id"]
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
        
        # 生成检查点ID（如果没有提供）
        checkpoint_id = checkpoint.get("id")
        if not checkpoint_id:
            from langgraph.checkpoint.base.id import uuid6
            checkpoint_id = str(uuid6())
            checkpoint["id"] = checkpoint_id
        
        # 存储检查点
        self.storage[thread_id][checkpoint_ns][checkpoint_id] = (
            checkpoint, 
            metadata
        )
        
        # 返回更新的配置
        return {
            **config,
            "configurable": {
                **config["configurable"],
                "checkpoint_id": checkpoint_id,
            }
        }
    
    def put_writes(
        self,
        config: RunnableConfig,
        writes: Sequence[tuple[str, Any]],
        task_id: str,
    ) -> None:
        """存储待写入操作
        
        Args:
            config: 运行配置
            writes: 写入操作序列
            task_id: 任务ID
        """
        # 提取配置参数
        thread_id = config["configurable"]["thread_id"]
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
        checkpoint_id = config["configurable"]["checkpoint_id"]
        
        # 存储写入操作
        self.writes[thread_id][checkpoint_ns][checkpoint_id][task_id].extend(
            writes
        )
    
    def _get_pending_writes(self, config: RunnableConfig) -> list[PendingWrite]:
        """获取待写入操作
        
        Args:
            config: 运行配置
            
        Returns:
            list[PendingWrite]: 待写入操作列表
        """
        thread_id = config["configurable"]["thread_id"]
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
        checkpoint_id = config["configurable"].get("checkpoint_id")
        
        if not checkpoint_id:
            return []
        
        pending = []
        checkpoint_writes = self.writes.get(thread_id, {}).get(checkpoint_ns, {}).get(checkpoint_id, {})
        
        for task_id, writes in checkpoint_writes.items():
            for channel, value in writes:
                pending.append((channel, task_id, value))
        
        return pending
    
    def _get_parent_config(
        self, 
        config: RunnableConfig, 
        metadata: CheckpointMetadata
    ) -> RunnableConfig | None:
        """获取父检查点配置
        
        Args:
            config: 当前配置
            metadata: 检查点元数据
            
        Returns:
            RunnableConfig | None: 父检查点配置或None
        """
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
        parent_id = metadata.get("parents", {}).get(checkpoint_ns)
        
        if parent_id:
            return {
                **config,
                "configurable": {
                    **config["configurable"],
                    "checkpoint_id": parent_id,
                }
            }
        
        return None
    
    def _matches_filter(self, metadata: CheckpointMetadata, filter: dict[str, Any]) -> bool:
        """检查元数据是否匹配过滤器
        
        Args:
            metadata: 检查点元数据
            filter: 过滤条件
            
        Returns:
            bool: 是否匹配
        """
        for key, value in filter.items():
            if key not in metadata:
                return False
            if metadata[key] != value:
                return False
        return True
    
    # 上下文管理器方法
    def __enter__(self) -> InMemorySaver:
        return self
    
    def __exit__(self, *args) -> None:
        pass
    
    async def __aenter__(self) -> InMemorySaver:
        return self
    
    async def __aexit__(self, *args) -> None:
        pass

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

-- 检查点主表
CREATE TABLE IF NOT EXISTS checkpoints (
    thread_id TEXT NOT NULL,
    checkpoint_ns TEXT NOT NULL DEFAULT '',
    checkpoint_id TEXT NOT NULL,
    parent_checkpoint_id TEXT,
    type TEXT,
    checkpoint JSONB NOT NULL,
    metadata JSONB NOT NULL DEFAULT '{}',
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
);

-- 索引优化
CREATE INDEX IF NOT EXISTS checkpoints_thread_id_idx 
ON checkpoints (thread_id);

CREATE INDEX IF NOT EXISTS checkpoints_created_at_idx 
ON checkpoints (created_at);

CREATE INDEX IF NOT EXISTS checkpoints_parent_id_idx 
ON checkpoints (parent_checkpoint_id);

-- 检查点写入表
CREATE TABLE IF NOT EXISTS checkpoint_writes (
    thread_id TEXT NOT NULL,
    checkpoint_ns TEXT NOT NULL DEFAULT '',
    checkpoint_id TEXT NOT NULL,
    task_id TEXT NOT NULL,
    idx INTEGER NOT NULL,
    channel TEXT NOT NULL,
    type TEXT,
    value JSONB,
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
);

-- 写入表索引
CREATE INDEX IF NOT EXISTS checkpoint_writes_lookup_idx 
ON checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id);

-- 迁移版本表
CREATE TABLE IF NOT EXISTS checkpoint_migrations (
    v INTEGER PRIMARY KEY
);