base_agent/uas/inout/dispatcher.py

# -*- coding: utf-8 -*-
"""
异步分发器
─────────────────────────────────────────────────────────────────
基于 QThreadPool + QRunnable 实现：
  1. 接收 input_received 信号
  2. 在工作线程中执行处理器链（Pipeline）
  3. 处理完成后 emit output_ready 信号
线程间通信完全由 Qt 信号槽队列保证安全
─────────────────────────────────────────────────────────────────
"""

import logging
import traceback
from typing import List, Dict, Callable

from PyQt5.QtCore import QObject, QRunnable, QThreadPool, pyqtSlot, pyqtSignal
from uas.io.base_io import IOData, IOType
from uas.io.signals import io_signals

logger = logging.getLogger(__name__)


# ─────────────────────────────────────────────────────────────────
# Worker：在线程池中执行处理器链
# ─────────────────────────────────────────────────────────────────

class _WorkerSignals(QObject):
    """Worker 内部信号（QRunnable 不能直接定义 pyqtSignal）"""
    finished = pyqtSignal(object)   # IOData
    error    = pyqtSignal(str, str) # data_id, error_msg


class _PipelineWorker(QRunnable):
    """
    在 QThreadPool 工作线程中执行处理器链
    每个 IOData 对应一个独立 Worker 实例
    """

    def __init__(self, data: IOData, processors: List[Callable]):
        super().__init__()
        self.data       = data
        self.processors = processors
        self.signals    = _WorkerSignals()
        self.setAutoDelete(True)

    @pyqtSlot()
    def run(self):
        data = self.data
        try:
            data.mark_processing()
            logger.debug(f"Worker 开始处理: {data}")

            # 顺序执行处理器链
            result = data.raw
            for idx, processor in enumerate(self.processors):
                io_signals.progress_updated.emit(
                    data.data_id,
                    int((idx / len(self.processors)) * 100)
                )
                result = processor(result, data)

            data.mark_success(result)
            io_signals.progress_updated.emit(data.data_id, 100)
            logger.debug(f"Worker 处理完成: {data}")
            self.signals.finished.emit(data)

        except Exception as e:
            msg = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
            data.mark_error(msg)
            logger.error(f"Worker 处理异常: {msg}")
            self.signals.error.emit(data.data_id, msg)


# ─────────────────────────────────────────────────────────────────
# 分发器主类
# ─────────────────────────────────────────────────────────────────

class IODispatcher(QObject):
    """
    IO 异步分发器
    ─────────────────────────────────────────────────────────────
    使用方式：
        dispatcher = IODispatcher()
        dispatcher.register_processor(IOType.TEXT, clean_text)
        dispatcher.register_processor(IOType.TEXT, translate_text)
        dispatcher.start()
    ─────────────────────────────────────────────────────────────
    """

    def __init__(self, max_threads: int = 8, parent=None):
        super().__init__(parent)
        self._pool: QThreadPool = QThreadPool.globalInstance()
        self._pool.setMaxThreadCount(max_threads)
        # { IOType: [processor_fn, ...] }
        self._processors: Dict[str, List[Callable]] = {}
        self._active = False
        logger.info(f"IODispatcher 初始化，最大线程数={max_threads}")

    # ── 处理器注册 ────────────────────────────────────────────────

    def register_processor(
        self,
        io_type:   IOType,
        processor: Callable,
        index:     int = -1,
    ) -> None:
        """
        注册处理器到指定 IO 类型的处理链
        :param io_type:   目标类型（IOType.TEXT 等）
        :param processor: fn(raw, IOData) -> processed_result
        :param index:     插入位置，-1 表示追加到末尾
        """
        key = io_type.value
        if key not in self._processors:
            self._processors[key] = []
        if index == -1:
            self._processors[key].append(processor)
        else:
            self._processors[key].insert(index, processor)
        logger.info(f"注册处理器: {processor.__name__} -> {key}")

    def unregister_processor(self, io_type: IOType, processor: Callable) -> None:
        key = io_type.value
        if key in self._processors:
            self._processors[key] = [
                p for p in self._processors[key] if p is not processor
            ]

    def clear_processors(self, io_type: IOType) -> None:
        self._processors.pop(io_type.value, None)

    # ── 生命周期 ──────────────────────────────────────────────────

    def start(self) -> None:
        """启动分发器，连接输入信号"""
        if self._active:
            return
        io_signals.input_received.connect(self._on_input_received)
        self._active = True
        logger.info("✅ IODispatcher 已启动")

    def stop(self) -> None:
        """停止分发器"""
        if not self._active:
            return
        io_signals.input_received.disconnect(self._on_input_received)
        self._pool.waitForDone(3000)
        self._active = False
        logger.info("🛑 IODispatcher 已停止")

    # ── 核心槽函数 ────────────────────────────────────────────────

    @pyqtSlot(str, str, object)
    def _on_input_received(self, data_id: str, io_type: str, data: object):
        """
        槽：接收原始输入，提交到线程池处理
        Qt 信号槽队列保证此函数在主线程执行，
        实际处理在 QRunnable 工作线程中进行
        """
        processors = self._processors.get(io_type, [])
        if not processors:
            # 无处理器，直接透传
            logger.debug(f"无处理器，透传 data_id={data_id[:8]} type={io_type}")
            data.mark_success()
            io_signals.output_ready.emit(data_id, io_type, data)
            return

        worker = _PipelineWorker(data, processors)
        worker.signals.finished.connect(self._on_worker_finished)
        worker.signals.error.connect(self._on_worker_error)
        self._pool.start(worker)
        logger.debug(f"提交 Worker: data_id={data_id[:8]} type={io_type} "
                     f"processors={len(processors)}")

    @pyqtSlot(object)
    def _on_worker_finished(self, data: object):
        """槽：Worker 完成，发射 output_ready"""
        io_signals.output_ready.emit(
            data.data_id, data.io_type.value, data
        )

    @pyqtSlot(str, str)
    def _on_worker_error(self, data_id: str, error_msg: str):
        logger.error(f"Worker 错误 data_id={data_id[:8]}: {error_msg}")