Coverage for aipyapp/aipy/task_state.py: 22%

1#!/usr/bin/env python

2# -*- coding: utf-8 -*-

3

4 import json

5 import time

6 from pathlib import Path

7 from typing import Dict, Any, Optional, Union, TYPE_CHECKING

8 from collections import OrderedDict

9

10 from loguru import logger

11

12 if TYPE_CHECKING:

13 from .task import Task

14

15# 任务版本常量

16 TASK_VERSION = 20250806

17

18 class TaskStateError(Exception):

19 """任务状态异常"""

20 pass

21

22 class TaskState:

23 """任务状态管理器 - 封装任务状态的序列化、反序列化和文件操作"""

24

25 def __init__(self, task: Optional['Task'] = None):

26 self.log = logger.bind(src='task_state')

27

28 # 任务基本信息

29 self.version: int = TASK_VERSION

30 self.task_id: Optional[str] = None

31 self.instruction: Optional[str] = None

32 self.start_time: Optional[float] = None

33 self.done_time: Optional[float] = None

34

35 # 组件状态

36 self._component_states: Dict[str, Any] = {}

37

38 if task:

39 self.from_task(task)

40

41 @classmethod

42 def from_file(cls, path: Union[str, Path]) -> 'TaskState':

43 """从文件创建 TaskState 对象"""

44 instance = cls()

45 instance.load_from_file(path)

46 return instance

47

48 @classmethod

49 def from_dict(cls, data: Dict[str, Any]) -> 'TaskState':

50 """从字典创建 TaskState 对象"""

51 instance = cls()

52 instance._load_from_dict(data)

53 return instance

54

55 def from_task(self, task: 'Task') -> None:

56 """从 Task 对象提取状态"""

57 self.task_id = task.task_id

58 self.instruction = task.instruction

59 self.start_time = task.start_time

60 self.done_time = task.done_time

61

62 # 提取各组件状态

63 self._component_states['steps'] = task.step_manager.get_state()

64 self._component_states['context_manager'] = task.context_manager.get_state()

65 self._component_states['runner'] = task.runner.get_state()

66 self._component_states['blocks'] = task.code_blocks.get_state()

67

68 # 提取事件记录（如果存在）

69 if task.event_recorder:

70 self._component_states['events'] = task.event_recorder.get_events()

71

72 self.log.debug('Extracted state from task', task_id=self.task_id)

73

74 def restore_to_task(self, task: 'Task') -> None:

75 """恢复状态到 Task 对象"""

76 # 验证版本兼容性

77 if self.version != TASK_VERSION:

78 raise TaskStateError(f'Task version mismatch: expected {TASK_VERSION}, got {self.version}')

79

80 # 恢复任务基本信息

81 task.task_id = self.task_id

82 task.instruction = self.instruction

83 task.start_time = self.start_time

84 task.done_time = self.done_time

85

86 # 恢复各组件状态

87 if 'steps' in self._component_states:

88 task.step_manager.restore_state(self._component_states['steps'])

89

90 if 'context_manager' in self._component_states:

91 task.context_manager.restore_state(self._component_states['context_manager'])

92

93 if 'runner' in self._component_states:

94 task.runner.restore_state(self._component_states['runner'])

95

96 if 'blocks' in self._component_states:

97 task.code_blocks.restore_state(self._component_states['blocks'])

98

99 # 恢复事件记录器状态（如果存在）

100 if 'events' in self._component_states and task.event_recorder:

101 events_data = self._component_states['events']

102 task.event_recorder.restore_state({'events': events_data, 'enabled': True})

103

104 self.log.info('Restored state to task', task_id=self.task_id)

105

106 def to_dict(self) -> Dict[str, Any]:

107 """序列化为字典"""

108 # 使用 OrderedDict 保持字段顺序

109 data = OrderedDict()

110 data['version'] = self.version

111 data['task_id'] = self.task_id

112 data['instruction'] = self.instruction

113 data['start_time'] = int(self.start_time) if self.start_time else None

114 data['done_time'] = int(self.done_time) if self.done_time else None

115

116 # 添加组件状态

117 for key, state in self._component_states.items():

118 data[key] = state

119

120 return data

121

122 def _load_from_dict(self, data: Dict[str, Any]) -> None:

123 """从字典加载状态"""

124 # 加载基本字段

125 basic_fields = {'version', 'task_id', 'instruction', 'start_time', 'done_time'}

126 self.version = data.get('version', TASK_VERSION)

127 self.task_id = data.get('task_id')

128 self.instruction = data.get('instruction')

129 self.start_time = data.get('start_time')

130 self.done_time = data.get('done_time')

131

132 # 加载所有非基本字段作为组件状态

133 for key, value in data.items():

134 if key not in basic_fields:

135 self._component_states[key] = value

136

137 def save_to_file(self, path: Union[str, Path]) -> None:

138 """保存到文件"""

139 path = Path(path)

140

141 # 确保目录存在

142 path.parent.mkdir(parents=True, exist_ok=True)

143

144 # 更新完成时间

145 if not self.done_time:

146 self.done_time = time.time()

147

148 try:

149 with open(path, 'w', encoding='utf-8') as f:

150 json.dump(self.to_dict(), f, ensure_ascii=False, indent=4, default=str)

151 self.log.info('Saved task state to file', path=str(path))

152 except Exception as e:

153 self.log.exception('Failed to save task state', path=str(path))

154 raise TaskStateError(f'Failed to save task state: {e}') from e

155

156 def load_from_file(self, path: Union[str, Path]) -> None:

157 """从文件加载状态"""

158 path = Path(path)

159 self.validate_file(path)

160

161 try:

162 with open(path, 'r', encoding='utf-8') as f:

163 data = json.load(f)

164 self._load_from_dict(data)

165 self.log.info('Loaded task state from file', path=str(path), task_id=self.task_id)

166 except json.JSONDecodeError as e:

167 raise TaskStateError(f'Invalid JSON file: {e}') from e

168 except Exception as e:

169 self.log.exception('Failed to load task state', path=str(path))

170 raise TaskStateError(f'Failed to load task state: {e}') from e

171

172 def validate_file(self, path: Union[str, Path]) -> None:

173 """验证文件格式和存在性"""

174 path = Path(path)

175

176 if not path.exists():

177 raise FileNotFoundError(f"Task file not found: {path}")

178

179 if not path.name.endswith('.json'):

180 raise ValueError("Task file must be a .json file")

181

182 if not path.is_file():

183 raise ValueError(f"Path is not a file: {path}")

184

185 def get_component_state(self, name: str) -> Any:

186 """获取组件状态"""

187 return self._component_states.get(name)

188

189 def set_component_state(self, name: str, state: Any) -> None:

190 """设置组件状态"""

191 self._component_states[name] = state

192

193 def has_component_state(self, name: str) -> bool:

194 """检查是否有指定组件的状态"""

195 return name in self._component_states

196

197 def get_summary(self) -> Dict[str, Any]:

198 """获取状态摘要信息"""

199 return {

200 'version': self.version,

201 'task_id': self.task_id,

202 'instruction': self.instruction[:50] + '...' if self.instruction and len(self.instruction) > 50 else self.instruction,

203 'start_time': self.start_time,

204 'done_time': self.done_time,

205 'components': list(self._component_states.keys())

206 }

207

208 def __repr__(self):

209 return f"<TaskState task_id={self.task_id}, version={self.version}, components={list(self._component_states.keys())}>"