Coverage for agentos/checkpoint/base.py: 93%

1""" 

2Checkpointer 抽象基类与数据结构。 

3""" 

4 

5from __future__ import annotations 

6 

7from abc import ABC, abstractmethod 

8from dataclasses import dataclass, field 

9from datetime import UTC, datetime 

10from typing import Any 

11 

12 

13@dataclass 

14class CheckpointMetadata: 

15 """Checkpoint 元信息。""" 

16 

17 thread_id: str # 对话线程 ID 

18 checkpoint_id: str # 唯一 ID 

19 step: int # 步骤序号 

20 parent_checkpoint_id: str | None = None # 父 checkpoint（用于分支/回溯） 

21 created_at: str = field(default_factory=lambda: datetime.now(UTC).isoformat()) 

22 tags: list[str] = field(default_factory=list) # 标签 

23 summary: str = "" # 可选摘要 

24 

25 

26@dataclass 

27class Checkpoint: 

28 """单个 Checkpoint — 完整的运行时状态快照。""" 

29 

30 metadata: CheckpointMetadata # 元信息 

31 messages: list[dict[str, Any]] # 对话消息（序列化后） 

32 state: dict[str, Any] # Agent 运行时状态 

33 tools_result: dict[str, Any] # 工具调用结果 

34 next_node: str = "" # 下一个执行节点 

35 

36 def to_dict(self) -> dict[str, Any]: 

37 return { 

38 "metadata": { 

39 "thread_id": self.metadata.thread_id, 

40 "checkpoint_id": self.metadata.checkpoint_id, 

41 "parent_checkpoint_id": self.metadata.parent_checkpoint_id, 

42 "step": self.metadata.step, 

43 "created_at": self.metadata.created_at, 

44 "tags": self.metadata.tags, 

45 "summary": self.metadata.summary, 

46 }, 

47 "messages": self.messages, 

48 "state": self.state, 

49 "tools_result": self.tools_result, 

50 "next_node": self.next_node, 

51 } 

52 

53 @classmethod 

54 def from_dict(cls, d: dict[str, Any]) -> Checkpoint: 

55 meta = d["metadata"] 

56 return cls( 

57 metadata=CheckpointMetadata( 

58 thread_id=meta["thread_id"], 

59 checkpoint_id=meta["checkpoint_id"], 

60 parent_checkpoint_id=meta.get("parent_checkpoint_id"), 

61 step=meta["step"], 

62 created_at=meta["created_at"], 

63 tags=meta.get("tags", []), 

64 summary=meta.get("summary", ""), 

65 ), 

66 messages=d.get("messages", []), 

67 state=d.get("state", {}), 

68 tools_result=d.get("tools_result", {}), 

69 next_node=d.get("next_node", ""), 

70 ) 

71 

72 

73class CheckpointBackend(ABC): 

74 """Checkpoint 存储后端抽象基类。""" 

75 

76 @abstractmethod 

77 async def put(self, checkpoint: Checkpoint) -> str: 

78 """保存 checkpoint，返回 checkpoint_id。""" 

79 ... 

80 

81 @abstractmethod 

82 async def get(self, checkpoint_id: str) -> Checkpoint | None: 

83 """按 ID 获取 checkpoint。""" 

84 ... 

85 

86 @abstractmethod 

87 async def get_latest(self, thread_id: str) -> Checkpoint | None: 

88 """获取某线程的最新 checkpoint。""" 

89 ... 

90 

91 @abstractmethod 

92 async def list_threads(self, limit: int = 50, offset: int = 0) -> list[CheckpointMetadata]: 

93 """列出所有线程的最新 checkpoint 元信息。""" 

94 ... 

95 

96 @abstractmethod 

97 async def list_checkpoints( 

98 self, thread_id: str, limit: int = 100, offset: int = 0 

99 ) -> list[CheckpointMetadata]: 

100 """列出某线程的所有 checkpoint（支持回溯/时间旅行）。""" 

101 ... 

102 

103 @abstractmethod 

104 async def delete_thread(self, thread_id: str) -> int: 

105 """删除某线程的所有 checkpoint，返回删除数。""" 

106 ... 

107 

108 @abstractmethod 

109 async def delete_before(self, thread_id: str, before_step: int) -> int: 

110 """删除某线程 before_step 之前的所有 checkpoint。""" 

111 ...