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

1#!/usr/bin/env python

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

3

4 import time

5 import json

6 from typing import List, Dict, Any, Optional

7 from datetime import datetime

8 from loguru import logger

9

10 from .event_serializer import EventSerializer

11 from ..interface import Trackable

12

13 class 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