Coverage for aipyapp/aipy/event_recorder.py: 29%

1#!/usr/bin/env python 

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

3 

4import time 

5import json 

6from typing import List, Dict, Any, Optional 

7from datetime import datetime 

8from loguru import logger 

9 

10from .event_serializer import EventSerializer 

11from ..interface import Trackable 

12 

13class EventRecorder(Trackable): 

14 """事件记录器 - 记录任务执行过程中的所有重要事件""" 

15 

16 def __init__(self, enabled: bool = True): 

17 self.enabled = enabled 

18 self.events: List[Dict[str, Any]] = [] 

19 self.start_time: Optional[float] = None 

20 self.log = logger.bind(src='event_recorder') 

21 

22 def start_recording(self): 

23 """开始记录""" 

24 self.start_time = time.time() 

25 self.events.clear() 

26 self.record_event('recording_start', {'timestamp': self.start_time}) 

27 self.log.info('Started event recording') 

28 

29 def stop_recording(self): 

30 """停止记录""" 

31 if self.start_time: 

32 self.record_event('recording_end', {'timestamp': time.time()}) 

33 self.log.info(f'Stopped event recording, total events: {len(self.events)}') 

34 

35 def record_event(self, event_type: str, data: Dict[str, Any], timestamp: Optional[float] = None): 

36 """记录事件 

37  

38 Args: 

39 event_type: 事件类型 

40 data: 事件数据 

41 timestamp: 时间戳，为None时使用当前时间 

42 """ 

43 if not self.enabled: 

44 return 

45 

46 if timestamp is None: 

47 timestamp = time.time() 

48 

49 # 计算相对时间 

50 relative_time = timestamp - self.start_time if self.start_time else 0 

51 

52 # 序列化对象参数 

53 

54 serialized_data = EventSerializer.serialize_event_data(data.copy() if isinstance(data, dict) else data) 

55 

56 event = { 

57 'type': event_type, 

58 'data': serialized_data, 

59 'timestamp': timestamp, 

60 'relative_time': relative_time, 

61 'datetime': datetime.fromtimestamp(timestamp).isoformat() 

62 } 

63 

64 self.events.append(event) 

65 

66 # 记录调试信息（简化版本，避免logger level检查） 

67 if len(self.events) % 100 == 0: # 每100个事件记录一次调试信息 

68 self.log.debug(f'Recorded {len(self.events)} events, latest: {event_type} at {relative_time:.3f}s') 

69 

70 

71 def get_events(self) -> List[Dict[str, Any]]: 

72 """获取所有事件""" 

73 return self.events.copy() 

74 

75 def get_events_for_replay(self) -> List[Dict[str, Any]]: 

76 """获取用于重放的事件（反序列化对象）""" 

77 return EventSerializer.deserialize_events(self.events) 

78 

79 def get_events_by_type(self, event_type: str) -> List[Dict[str, Any]]: 

80 """获取指定类型的事件""" 

81 return [event for event in self.events if event['type'] == event_type] 

82 

83 def get_events_in_range(self, start_time: float, end_time: float) -> List[Dict[str, Any]]: 

84 """获取指定时间范围内的事件""" 

85 return [event for event in self.events 

86 if start_time <= event['relative_time'] <= end_time] 

87 

88 def clear_events(self): 

89 """清空所有事件""" 

90 self.events.clear() 

91 self.start_time = None 

92 self.log.info('Cleared all events') 

93 

94 # Trackable接口实现 

95 def get_checkpoint(self) -> int: 

96 """获取当前检查点状态 - 返回事件数量""" 

97 return len(self.events) 

98 

99 def restore_to_checkpoint(self, checkpoint: Optional[int]): 

100 """恢复到指定检查点""" 

101 if checkpoint is None: 

102 # 恢复到初始状态 

103 self.clear_events() 

104 else: 

105 # 恢复到指定事件数量 

106 if checkpoint < len(self.events): 

107 deleted_count = len(self.events) - checkpoint 

108 self.events = self.events[:checkpoint] 

109 self.log.info(f'Restored to checkpoint {checkpoint}, deleted {deleted_count} events') 

110 

111 def get_state(self) -> Dict[str, Any]: 

112 """获取需要持久化的状态数据""" 

113 return { 

114 'enabled': self.enabled, 

115 'start_time': self.start_time, 

116 'events': self.events 

117 } 

118 

119 def restore_state(self, state_data: Dict[str, Any]): 

120 """从状态数据恢复事件记录器""" 

121 self.enabled = state_data.get('enabled', True) 

122 self.start_time = state_data.get('start_time') 

123 self.events = state_data.get('events', []) 

124 self.log.info(f'Restored event recorder with {len(self.events)} events') 

125 

126 def export_to_file(self, filepath: str): 

127 """导出事件到文件""" 

128 try: 

129 with open(filepath, 'w', encoding='utf-8') as f: 

130 json.dump({ 

131 'metadata': { 

132 'start_time': self.start_time, 

133 'total_events': len(self.events), 

134 'duration': self.events[-1]['relative_time'] if self.events else 0 

135 }, 

136 'events': self.events 

137 }, f, indent=2, ensure_ascii=False) 

138 self.log.info(f'Exported {len(self.events)} events to {filepath}') 

139 except Exception as e: 

140 self.log.error(f'Failed to export events: {e}') 

141 raise 

142 

143 def import_from_file(self, filepath: str): 

144 """从文件导入事件""" 

145 try: 

146 with open(filepath, 'r', encoding='utf-8') as f: 

147 data = json.load(f) 

148 

149 self.start_time = data.get('metadata', {}).get('start_time') 

150 self.events = data.get('events', []) 

151 self.log.info(f'Imported {len(self.events)} events from {filepath}') 

152 except Exception as e: 

153 self.log.error(f'Failed to import events: {e}') 

154 raise 

155 

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

157 """获取事件摘要统计""" 

158 if not self.events: 

159 return {'total_events': 0, 'duration': 0, 'event_types': {}} 

160 

161 # 统计事件类型 

162 event_types = {} 

163 for event in self.events: 

164 event_type = event['type'] 

165 event_types[event_type] = event_types.get(event_type, 0) + 1 

166 

167 return { 

168 'total_events': len(self.events), 

169 'duration': self.events[-1]['relative_time'] if self.events else 0, 

170 'start_time': self.start_time, 

171 'event_types': event_types 

172 } 

173 

174 def __len__(self): 

175 """返回事件数量""" 

176 return len(self.events) 

177 

178 def __bool__(self): 

179 """检查是否有事件""" 

180 return len(self.events) > 0