ruff

src/chat/message_manager/adaptive_stream_manager.py

@@ -17,6 +17,7 @@ logger = get_logger("adaptive_stream_manager")
 
 class StreamPriority(Enum):
     """流优先级"""
+
     LOW = 1
     NORMAL = 2
     HIGH = 3
@@ -26,6 +27,7 @@ class StreamPriority(Enum):
 @dataclass
 class SystemMetrics:
     """系统指标"""
+
     cpu_usage: float = 0.0
     memory_usage: float = 0.0
     active_coroutines: int = 0
@@ -36,6 +38,7 @@ class SystemMetrics:
 @dataclass
 class StreamMetrics:
     """流指标"""
+
     stream_id: str
     priority: StreamPriority
     message_rate: float = 0.0  # 消息速率（消息/分钟）
@@ -56,7 +59,7 @@ class AdaptiveStreamManager:
         metrics_window: float = 60.0,  # 指标窗口时间
         adjustment_interval: float = 30.0,  # 调整间隔
         cpu_threshold_high: float = 0.8,  # CPU高负载阈值
-        cpu_threshold_low: float = 0.3,   # CPU低负载阈值
+        cpu_threshold_low: float = 0.3,  # CPU低负载阈值
         memory_threshold_high: float = 0.85,  # 内存高负载阈值
     ):
         self.base_concurrent_limit = base_concurrent_limit
@@ -139,10 +142,7 @@ class AdaptiveStreamManager:
         logger.info("自适应流管理器已停止")
 
     async def acquire_stream_slot(
-        self,
-        stream_id: str,
-        priority: StreamPriority = StreamPriority.NORMAL,
-        force: bool = False
+        self, stream_id: str, priority: StreamPriority = StreamPriority.NORMAL, force: bool = False
     ) -> bool:
         """
         获取流处理槽位
@@ -165,10 +165,7 @@ class AdaptiveStreamManager:
 
         # 更新流指标
         if stream_id not in self.stream_metrics:
-            self.stream_metrics[stream_id] = StreamMetrics(
-                stream_id=stream_id,
-                priority=priority
-            )
+            self.stream_metrics[stream_id] = StreamMetrics(stream_id=stream_id, priority=priority)
         self.stream_metrics[stream_id].last_activity = current_time
 
         # 检查是否已经活跃
@@ -271,8 +268,10 @@ class AdaptiveStreamManager:
 
         # 如果最近有活跃且响应时间较长，可能需要强制分发
         current_time = time.time()
-        if (current_time - metrics.last_activity < 300 and  # 5分钟内有活动
-            metrics.response_time > 5.0):  # 响应时间超过5秒
+        if (
+            current_time - metrics.last_activity < 300  # 5分钟内有活动
+            and metrics.response_time > 5.0
+        ):  # 响应时间超过5秒
             return True
 
         return False
@@ -324,26 +323,20 @@ class AdaptiveStreamManager:
                 memory_usage=memory_usage,
                 active_coroutines=active_coroutines,
                 event_loop_lag=event_loop_lag,
-                timestamp=time.time()
+                timestamp=time.time(),
             )
 
             self.system_metrics.append(metrics)
 
             # 保持指标窗口大小
             cutoff_time = time.time() - self.metrics_window
-            self.system_metrics = [
-                m for m in self.system_metrics
-                if m.timestamp > cutoff_time
-            ]
+            self.system_metrics = [m for m in self.system_metrics if m.timestamp > cutoff_time]
 
             # 更新统计信息
             self.stats["avg_concurrent_streams"] = (
                 self.stats["avg_concurrent_streams"] * 0.9 + len(self.active_streams) * 0.1
             )
-            self.stats["peak_concurrent_streams"] = max(
-                self.stats["peak_concurrent_streams"],
-                len(self.active_streams)
-            )
+            self.stats["peak_concurrent_streams"] = max(self.stats["peak_concurrent_streams"], len(self.active_streams))
 
         except Exception as e:
             logger.error(f"收集系统指标失败: {e}")
@@ -445,14 +438,16 @@ class AdaptiveStreamManager:
     def get_stats(self) -> dict:
         """获取统计信息"""
         stats = self.stats.copy()
-        stats.update({
-            "current_limit": self.current_limit,
-            "active_streams": len(self.active_streams),
-            "pending_streams": len(self.pending_streams),
-            "is_running": self.is_running,
-            "system_cpu": self.system_metrics[-1].cpu_usage if self.system_metrics else 0,
-            "system_memory": self.system_metrics[-1].memory_usage if self.system_metrics else 0,
-        })
+        stats.update(
+            {
+                "current_limit": self.current_limit,
+                "active_streams": len(self.active_streams),
+                "pending_streams": len(self.pending_streams),
+                "is_running": self.is_running,
+                "system_cpu": self.system_metrics[-1].cpu_usage if self.system_metrics else 0,
+                "system_memory": self.system_metrics[-1].memory_usage if self.system_metrics else 0,
+            }
+        )
 
         # 计算接受率
         if stats["total_requests"] > 0:

src/chat/message_manager/batch_database_writer.py

@@ -20,6 +20,7 @@ logger = get_logger("batch_database_writer")
 @dataclass
 class StreamUpdatePayload:
     """流更新数据结构"""
+
     stream_id: str
     update_data: dict[str, Any]
     priority: int = 0  # 优先级，数字越大优先级越高
@@ -95,12 +96,7 @@ class BatchDatabaseWriter:
 
         logger.info("批量数据库写入器已停止")
 
-    async def schedule_stream_update(
-        self,
-        stream_id: str,
-        update_data: dict[str, Any],
-        priority: int = 0
-    ) -> bool:
+    async def schedule_stream_update(self, stream_id: str, update_data: dict[str, Any], priority: int = 0) -> bool:
         """
         调度流更新
 
@@ -119,11 +115,7 @@ class BatchDatabaseWriter:
                 return True
 
             # 创建更新载荷
-            payload = StreamUpdatePayload(
-                stream_id=stream_id,
-                update_data=update_data,
-                priority=priority
-            )
+            payload = StreamUpdatePayload(stream_id=stream_id, update_data=update_data, priority=priority)
 
             # 非阻塞方式加入队列
             try:
@@ -178,10 +170,7 @@ class BatchDatabaseWriter:
                 if remaining_time == 0:
                     break
 
-                payload = await asyncio.wait_for(
-                    self.write_queue.get(),
-                    timeout=remaining_time
-                )
+                payload = await asyncio.wait_for(self.write_queue.get(), timeout=remaining_time)
                 batch.append(payload)
 
             except asyncio.TimeoutError:
@@ -203,7 +192,10 @@ class BatchDatabaseWriter:
             # 合并同一流ID的更新（保留最新的）
             merged_updates = {}
             for payload in batch:
-                if payload.stream_id not in merged_updates or payload.timestamp > merged_updates[payload.stream_id].timestamp:
+                if (
+                    payload.stream_id not in merged_updates
+                    or payload.timestamp > merged_updates[payload.stream_id].timestamp
+                ):
                     merged_updates[payload.stream_id] = payload
 
             # 批量写入
@@ -211,9 +203,7 @@ class BatchDatabaseWriter:
 
             # 更新统计
             self.stats["batch_writes"] += 1
-            self.stats["avg_batch_size"] = (
-                self.stats["avg_batch_size"] * 0.9 + len(batch) * 0.1
-            )  # 滑动平均
+            self.stats["avg_batch_size"] = self.stats["avg_batch_size"] * 0.9 + len(batch) * 0.1  # 滑动平均
             self.stats["last_flush_time"] = start_time
 
             logger.debug(f"批量写入完成: {len(batch)} 个更新，耗时 {time.time() - start_time:.3f}s")
@@ -238,31 +228,22 @@ class BatchDatabaseWriter:
                 # 根据数据库类型选择不同的插入/更新策略
                 if global_config.database.database_type == "sqlite":
                     from sqlalchemy.dialects.sqlite import insert as sqlite_insert
-                    stmt = sqlite_insert(ChatStreams).values(
-                        stream_id=stream_id, **update_data
-                    )
-                    stmt = stmt.on_conflict_do_update(
-                        index_elements=["stream_id"],
-                        set_=update_data
-                    )
+
+                    stmt = sqlite_insert(ChatStreams).values(stream_id=stream_id, **update_data)
+                    stmt = stmt.on_conflict_do_update(index_elements=["stream_id"], set_=update_data)
                 elif global_config.database.database_type == "mysql":
                     from sqlalchemy.dialects.mysql import insert as mysql_insert
-                    stmt = mysql_insert(ChatStreams).values(
-                        stream_id=stream_id, **update_data
-                    )
+
+                    stmt = mysql_insert(ChatStreams).values(stream_id=stream_id, **update_data)
                     stmt = stmt.on_duplicate_key_update(
                         **{key: value for key, value in update_data.items() if key != "stream_id"}
                     )
                 else:
                     # 默认使用SQLite语法
                     from sqlalchemy.dialects.sqlite import insert as sqlite_insert
-                    stmt = sqlite_insert(ChatStreams).values(
-                        stream_id=stream_id, **update_data
-                    )
-                    stmt = stmt.on_conflict_do_update(
-                        index_elements=["stream_id"],
-                        set_=update_data
-                    )
+
+                    stmt = sqlite_insert(ChatStreams).values(stream_id=stream_id, **update_data)
+                    stmt = stmt.on_conflict_do_update(index_elements=["stream_id"], set_=update_data)
 
                 await session.execute(stmt)
 
@@ -273,30 +254,21 @@ class BatchDatabaseWriter:
         async with get_db_session() as session:
             if global_config.database.database_type == "sqlite":
                 from sqlalchemy.dialects.sqlite import insert as sqlite_insert
-                stmt = sqlite_insert(ChatStreams).values(
-                    stream_id=stream_id, **update_data
-                )
-                stmt = stmt.on_conflict_do_update(
-                    index_elements=["stream_id"],
-                    set_=update_data
-                )
+
+                stmt = sqlite_insert(ChatStreams).values(stream_id=stream_id, **update_data)
+                stmt = stmt.on_conflict_do_update(index_elements=["stream_id"], set_=update_data)
             elif global_config.database.database_type == "mysql":
                 from sqlalchemy.dialects.mysql import insert as mysql_insert
-                stmt = mysql_insert(ChatStreams).values(
-                    stream_id=stream_id, **update_data
-                )
+
+                stmt = mysql_insert(ChatStreams).values(stream_id=stream_id, **update_data)
                 stmt = stmt.on_duplicate_key_update(
                     **{key: value for key, value in update_data.items() if key != "stream_id"}
                 )
             else:
                 from sqlalchemy.dialects.sqlite import insert as sqlite_insert
-                stmt = sqlite_insert(ChatStreams).values(
-                    stream_id=stream_id, **update_data
-                )
-                stmt = stmt.on_conflict_do_update(
-                    index_elements=["stream_id"],
-                    set_=update_data
-                )
+
+                stmt = sqlite_insert(ChatStreams).values(stream_id=stream_id, **update_data)
+                stmt = stmt.on_conflict_do_update(index_elements=["stream_id"], set_=update_data)
 
             await session.execute(stmt)
             await session.commit()

src/chat/message_manager/context_manager.py

@@ -273,8 +273,10 @@ class SingleStreamContextManager:
                     message.should_reply = result.should_reply
                     message.should_act = result.should_act
 
-                    logger.debug(f"消息 {message.message_id} 兴趣值已更新: {result.interest_value:.3f}, "
-                                f"should_reply: {result.should_reply}, should_act: {result.should_act}")
+                    logger.debug(
+                        f"消息 {message.message_id} 兴趣值已更新: {result.interest_value:.3f}, "
+                        f"should_reply: {result.should_reply}, should_act: {result.should_act}"
+                    )
                     return result.interest_value
                 else:
                     logger.warning(f"消息 {message.message_id} 兴趣值计算失败: {result.error_message}")

src/chat/message_manager/distribution_manager.py

@@ -79,7 +79,7 @@ class StreamLoopManager:
                 logger.info(f"正在取消 {len(cancel_tasks)} 个流循环任务...")
                 await asyncio.gather(
                     *[self._wait_for_task_cancel(stream_id, task) for stream_id, task in cancel_tasks],
-                    return_exceptions=True
+                    return_exceptions=True,
                 )
 
             # 取消所有活跃的 chatter 处理任务
@@ -115,6 +115,7 @@ class StreamLoopManager:
         # 使用自适应流管理器获取槽位
         try:
             from src.chat.message_manager.adaptive_stream_manager import get_adaptive_stream_manager
+
             adaptive_manager = get_adaptive_stream_manager()
 
             if adaptive_manager.is_running:
@@ -123,9 +124,7 @@ class StreamLoopManager:
 
                 # 获取处理槽位
                 slot_acquired = await adaptive_manager.acquire_stream_slot(
-                    stream_id=stream_id,
-                    priority=priority,
-                    force=force
+                    stream_id=stream_id, priority=priority, force=force
                 )
 
                 if slot_acquired:
@@ -140,10 +139,7 @@ class StreamLoopManager:
 
         # 创建流循环任务
         try:
-            loop_task = asyncio.create_task(
-                self._stream_loop_worker(stream_id),
-                name=f"stream_loop_{stream_id}"
-            )
+            loop_task = asyncio.create_task(self._stream_loop_worker(stream_id), name=f"stream_loop_{stream_id}")
             self.stream_loops[stream_id] = loop_task
             # 更新统计信息
             self.stats["active_streams"] += 1
@@ -156,6 +152,7 @@ class StreamLoopManager:
             logger.error(f"启动流循环任务失败 {stream_id}: {e}")
             # 释放槽位
             from src.chat.message_manager.adaptive_stream_manager import get_adaptive_stream_manager
+
             adaptive_manager = get_adaptive_stream_manager()
             adaptive_manager.release_stream_slot(stream_id)
 
@@ -179,8 +176,8 @@ class StreamLoopManager:
 
         except Exception:
             from src.chat.message_manager.adaptive_stream_manager import StreamPriority
-            return StreamPriority.NORMAL
 
+            return StreamPriority.NORMAL
 
     async def stop_stream_loop(self, stream_id: str) -> bool:
         """停止指定流的循环任务
@@ -244,11 +241,12 @@ class StreamLoopManager:
                     # 3. 更新自适应管理器指标
                     try:
                         from src.chat.message_manager.adaptive_stream_manager import get_adaptive_stream_manager
+
                         adaptive_manager = get_adaptive_stream_manager()
                         adaptive_manager.update_stream_metrics(
                             stream_id,
                             message_rate=unread_count / 5.0 if unread_count > 0 else 0.0,  # 简化计算
-                            last_activity=time.time()
+                            last_activity=time.time(),
                         )
                     except Exception as e:
                         logger.debug(f"更新流指标失败: {e}")
@@ -300,6 +298,7 @@ class StreamLoopManager:
             # 释放自适应管理器的槽位
             try:
                 from src.chat.message_manager.adaptive_stream_manager import get_adaptive_stream_manager
+
                 adaptive_manager = get_adaptive_stream_manager()
                 adaptive_manager.release_stream_slot(stream_id)
                 logger.debug(f"释放自适应流处理槽位: {stream_id}")
@@ -553,12 +552,12 @@ class StreamLoopManager:
                     existing_task.cancel()
                     # 创建异步任务来等待取消完成，并添加异常处理
                     cancel_task = asyncio.create_task(
-                        self._wait_for_task_cancel(stream_id, existing_task),
-                        name=f"cancel_existing_loop_{stream_id}"
+                        self._wait_for_task_cancel(stream_id, existing_task), name=f"cancel_existing_loop_{stream_id}"
                     )
                     # 为取消任务添加异常处理，避免孤儿任务
                     cancel_task.add_done_callback(
-                        lambda task: logger.debug(f"取消任务完成: {stream_id}") if not task.exception()
+                        lambda task: logger.debug(f"取消任务完成: {stream_id}")
+                        if not task.exception()
                         else logger.error(f"取消任务异常: {stream_id} - {task.exception()}")
                     )
                 # 从字典中移除
@@ -582,10 +581,7 @@ class StreamLoopManager:
             logger.info(f"流 {stream_id} 当前未读消息数: {unread_count}")
 
             # 创建新的流循环任务
-            new_task = asyncio.create_task(
-                self._stream_loop(stream_id),
-                name=f"force_stream_loop_{stream_id}"
-            )
+            new_task = asyncio.create_task(self._stream_loop(stream_id), name=f"force_stream_loop_{stream_id}")
             self.stream_loops[stream_id] = new_task
             self.stats["total_loops"] += 1
 

src/chat/message_manager/message_manager.py

@@ -59,6 +59,7 @@ class MessageManager:
         # 启动批量数据库写入器
         try:
             from src.chat.message_manager.batch_database_writer import init_batch_writer
+
             await init_batch_writer()
         except Exception as e:
             logger.error(f"启动批量数据库写入器失败: {e}")
@@ -66,6 +67,7 @@ class MessageManager:
         # 启动流缓存管理器
         try:
             from src.chat.message_manager.stream_cache_manager import init_stream_cache_manager
+
             await init_stream_cache_manager()
         except Exception as e:
             logger.error(f"启动流缓存管理器失败: {e}")
@@ -73,6 +75,7 @@ class MessageManager:
         # 启动自适应流管理器
         try:
             from src.chat.message_manager.adaptive_stream_manager import init_adaptive_stream_manager
+
             await init_adaptive_stream_manager()
             logger.info("🎯 自适应流管理器已启动")
         except Exception as e:
@@ -97,6 +100,7 @@ class MessageManager:
         # 停止批量数据库写入器
         try:
             from src.chat.message_manager.batch_database_writer import shutdown_batch_writer
+
             await shutdown_batch_writer()
             logger.info("📦 批量数据库写入器已停止")
         except Exception as e:
@@ -105,6 +109,7 @@ class MessageManager:
         # 停止流缓存管理器
         try:
             from src.chat.message_manager.stream_cache_manager import shutdown_stream_cache_manager
+
             await shutdown_stream_cache_manager()
             logger.info("🗄️ 流缓存管理器已停止")
         except Exception as e:
@@ -113,6 +118,7 @@ class MessageManager:
         # 停止自适应流管理器
         try:
             from src.chat.message_manager.adaptive_stream_manager import shutdown_adaptive_stream_manager
+
             await shutdown_adaptive_stream_manager()
             logger.info("🎯 自适应流管理器已停止")
         except Exception as e:

src/chat/message_manager/stream_cache_manager.py

@@ -19,6 +19,7 @@ logger = get_logger("stream_cache_manager")
 @dataclass
 class StreamCacheStats:
     """缓存统计信息"""
+
     hot_cache_size: int = 0
     warm_storage_size: int = 0
     cold_storage_size: int = 0
@@ -38,9 +39,9 @@ class TieredStreamCache:
         max_warm_size: int = 500,
         max_cold_size: int = 2000,
         cleanup_interval: float = 300.0,  # 5分钟清理一次
-        hot_timeout: float = 1800.0,      # 30分钟未访问降级到warm
-        warm_timeout: float = 7200.0,     # 2小时未访问降级到cold
-        cold_timeout: float = 86400.0,    # 24小时未访问删除
+        hot_timeout: float = 1800.0,  # 30分钟未访问降级到warm
+        warm_timeout: float = 7200.0,  # 2小时未访问降级到cold
+        cold_timeout: float = 86400.0,  # 24小时未访问删除
     ):
         self.max_hot_size = max_hot_size
         self.max_warm_size = max_warm_size
@@ -52,8 +53,8 @@ class TieredStreamCache:
 
         # 三层缓存存储
         self.hot_cache: OrderedDict[str, OptimizedChatStream] = OrderedDict()  # 热数据（LRU）
-        self.warm_storage: dict[str, tuple[OptimizedChatStream, float]] = {}   # 温数据（最后访问时间）
-        self.cold_storage: dict[str, tuple[OptimizedChatStream, float]] = {}   # 冷数据（最后访问时间）
+        self.warm_storage: dict[str, tuple[OptimizedChatStream, float]] = {}  # 温数据（最后访问时间）
+        self.cold_storage: dict[str, tuple[OptimizedChatStream, float]] = {}  # 冷数据（最后访问时间）
 
         # 统计信息
         self.stats = StreamCacheStats()
@@ -134,11 +135,7 @@ class TieredStreamCache:
         # 4. 缓存未命中，创建新流
         self.stats.cache_misses += 1
         stream = create_optimized_chat_stream(
-            stream_id=stream_id,
-            platform=platform,
-            user_info=user_info,
-            group_info=group_info,
-            data=data
+            stream_id=stream_id, platform=platform, user_info=user_info, group_info=group_info, data=data
         )
         logger.debug(f"缓存未命中，创建新流: {stream_id}")
 
@@ -294,9 +291,9 @@ class TieredStreamCache:
 
         # 估算内存使用（粗略估计）
         self.stats.total_memory_usage = (
-            len(self.hot_cache) * 1024 +      # 每个热流约1KB
-            len(self.warm_storage) * 512 +    # 每个温流约512B
-            len(self.cold_storage) * 256      # 每个冷流约256B
+            len(self.hot_cache) * 1024  # 每个热流约1KB
+            + len(self.warm_storage) * 512  # 每个温流约512B
+            + len(self.cold_storage) * 256  # 每个冷流约256B
         )
 
         if sum(cleanup_stats.values()) > 0: