136 lines
5.4 KiB
Python
136 lines
5.4 KiB
Python
import logging
|
||
import time
|
||
from typing import Dict, Any, Tuple, Optional
|
||
from datetime import datetime
|
||
|
||
from plugin_common.plugin_interface import PluginInterface
|
||
from event_system.event_manager import EventManager
|
||
# 修正导入，使用与装饰器相同的事件类型
|
||
from event_system.events.plugin_events import PluginCallStartEvent, PluginCallEndEvent, PluginCallErrorEvent
|
||
from db.stats_db import StatsDBOperator
|
||
from db.db_manager import DBConnectionManager
|
||
from job_decorators import register_job_decorator
|
||
|
||
from .decorators import plugin_stats_decorator
|
||
|
||
|
||
class StatsCollectorPlugin(PluginInterface):
|
||
"""统计收集插件"""
|
||
|
||
def __init__(self):
|
||
self.name = "统计收集器"
|
||
self.version = "1.0.0"
|
||
self.description = "收集插件调用统计数据"
|
||
self.author = "Trae AI"
|
||
self.logger = logging.getLogger("StatsCollector")
|
||
|
||
# 默认配置
|
||
self.config = {
|
||
"enable": True,
|
||
"record_all_plugins": True, # 是否记录所有插件的调用
|
||
"excluded_plugins": [], # 排除的插件列表
|
||
}
|
||
|
||
self.event_manager = EventManager.get_instance()
|
||
# 修正获取数据库连接管理器的方式
|
||
self.db_manager = DBConnectionManager.get_instance()
|
||
self.stats_db = StatsDBOperator(self.db_manager)
|
||
|
||
# 用于临时存储插件调用开始时间的字典
|
||
self.plugin_call_start_times = {}
|
||
|
||
def initialize(self, config: Dict[str, Any]) -> bool:
|
||
"""初始化插件"""
|
||
if config:
|
||
self.config.update(config)
|
||
|
||
if not self.config["enable"]:
|
||
self.logger.info("统计收集插件已禁用")
|
||
return False
|
||
|
||
# 注册事件处理器
|
||
self.event_manager.register(PluginCallStartEvent, self.handle_plugin_call_start)
|
||
self.event_manager.register(PluginCallEndEvent, self.handle_plugin_call_end)
|
||
self.event_manager.register(PluginCallErrorEvent, self.handle_plugin_error)
|
||
|
||
self.logger.info("统计收集插件已初始化")
|
||
return True
|
||
|
||
def handle_plugin_call_start(self, event: PluginCallStartEvent) -> None:
|
||
"""处理插件调用开始事件"""
|
||
# 检查是否需要记录该插件
|
||
if not self._should_record_plugin(event.plugin_name):
|
||
return
|
||
|
||
# 记录开始时间和相关信息
|
||
# 注意：plugin_events.py 中的事件结构与 stats_events.py 不同
|
||
self.logger.debug(f"记录插件调用开始: {event.plugin_name} - {event.command}")
|
||
|
||
def handle_plugin_call_end(self, event: PluginCallEndEvent) -> None:
|
||
"""处理插件调用结束事件"""
|
||
# 检查是否需要记录该插件
|
||
if not self._should_record_plugin(event.plugin_name):
|
||
return
|
||
|
||
# 记录统计数据
|
||
try:
|
||
self.stats_db.record_plugin_call(
|
||
plugin_name=event.plugin_name,
|
||
command=event.command,
|
||
user_id=event.user_id,
|
||
group_id=event.group_id,
|
||
success=event.process_result, # 注意字段名不同
|
||
process_time_ms=event.process_time # 注意字段名不同
|
||
)
|
||
self.logger.debug(f"记录插件调用结束: {event.plugin_name} - {event.command} - 成功: {event.process_result} - 处理时间: {event.process_time}ms")
|
||
except Exception as e:
|
||
self.logger.error(f"记录插件调用统计数据出错: {e}")
|
||
|
||
def handle_plugin_error(self, event: PluginCallErrorEvent) -> None:
|
||
"""处理插件调用错误事件"""
|
||
# 检查是否需要记录该插件
|
||
if not self._should_record_plugin(event.plugin_name):
|
||
return
|
||
|
||
# 记录错误信息
|
||
try:
|
||
self.stats_db.record_error(
|
||
plugin_name=event.plugin_name,
|
||
command=event.command,
|
||
user_id=event.user_id,
|
||
group_id=event.group_id,
|
||
error_message=event.error_message,
|
||
stack_trace=event.stack_trace
|
||
)
|
||
self.logger.debug(f"记录插件调用错误: {event.plugin_name} - {event.command} - {event.error_message}")
|
||
except Exception as e:
|
||
self.logger.error(f"记录插件错误信息出错: {e}")
|
||
|
||
def _should_record_plugin(self, plugin_name: str) -> bool:
|
||
"""检查是否应该记录该插件的调用"""
|
||
if not self.config["record_all_plugins"]:
|
||
return False
|
||
|
||
if plugin_name in self.config["excluded_plugins"]:
|
||
return False
|
||
|
||
return True
|
||
|
||
def match_command(self, content: str) -> bool:
|
||
"""匹配命令"""
|
||
# 该插件不处理用户消息
|
||
return False
|
||
|
||
def process_message(self, message: Dict[str, Any]) -> Tuple[bool, str]:
|
||
"""处理消息"""
|
||
# 该插件不处理用户消息
|
||
return False, ""
|
||
|
||
def shutdown(self) -> None:
|
||
"""关闭插件"""
|
||
# 取消注册事件处理器
|
||
self.event_manager.unregister(PluginCallStartEvent, self.handle_plugin_call_start)
|
||
self.event_manager.unregister(PluginCallEndEvent, self.handle_plugin_call_end)
|
||
self.event_manager.unregister(PluginCallErrorEvent, self.handle_plugin_error)
|
||
|
||
self.logger.info("统计收集插件已关闭") |