ForcePilot/backend/package/yuxi/channel/message/bridge.py

import asyncio
import contextlib
import json
import logging
import time
from collections.abc import AsyncIterator, Callable

from yuxi.channel.message.models import (
    ChunkType,
    DispatchResult,
    ReplyStage,
    StreamCancellationToken,
    StreamCancelledError,
    StreamingChunk,
    StreamResult,
    StreamStatus,
    StreamTimeoutError,
    UnifiedMessage,
)
from yuxi.channel.protocols import AgentPromptContext, AgentPromptProtocol, AgentToolProtocol
from yuxi.channel.plugins.registry import ChannelPluginRegistry
from yuxi.channel.security.external_content import sanitize_external_content

logger = logging.getLogger(__name__)

_STREAMING_TEXT_STATUSES = frozenset({"init", "loading", "agent_state", "streaming"})
_STREAMING_REASONING_STATUSES = frozenset({"thinking"})
_TOOL_STATUSES = frozenset({"tool_call", "tool_result"})
_TERMINAL_STATUSES = frozenset({"finished", "error", "interrupted", "warning"})
_INTERACTIVE_STATUSES = frozenset({"ask_user_question_required"})

_DEFAULT_TIMEOUT_SECONDS = 300.0


class AgentBridge:
    def __init__(
        self,
        stream_fn: Callable[..., AsyncIterator[bytes]],
        channel_manager=None,
    ):
        self._stream_fn = stream_fn
        self._channel_manager = channel_manager

    def collect_channel_tools(self, channel_type: str) -> list[dict]:
        plugin = ChannelPluginRegistry.get(channel_type)
        if plugin is None or not isinstance(plugin, AgentToolProtocol):
            return []
        tools = plugin.get_agent_tools()
        return [t.to_openai_schema() for t in tools]

    async def execute_channel_tool(
        self,
        channel_type: str,
        tool_name: str,
        params: dict,
        context: dict,
    ) -> dict:
        plugin = ChannelPluginRegistry.get(channel_type)
        if plugin is None or not isinstance(plugin, AgentToolProtocol):
            return {"success": False, "error": f"Channel {channel_type} does not support agent tools"}

        try:
            return await plugin.execute_agent_tool(tool_name, params, context)
        except Exception as e:
            logger.exception("Channel tool execution failed: %s.%s", channel_type, tool_name)
            return {"success": False, "error": str(e)}

    # ── 构建 meta ──────────────────────────────────────────

    def _build_meta(self, msg: UnifiedMessage) -> tuple[dict, str]:
        meta = {
            "source": "channel",
            "channel_type": msg.channel_type,
            "account_id": msg.account_id,
            "peer_id": msg.sender.id,
        }

        plugin = ChannelPluginRegistry.get(msg.channel_type)

        safe_content, _detection = sanitize_external_content(msg.content)
        query = safe_content

        if isinstance(plugin, AgentPromptProtocol):
            context = AgentPromptContext(
                channel_type=msg.channel_type,
                account_id=msg.account_id,
                peer_id=msg.sender.id,
                peer_name=msg.sender.display_name,
                group_id=msg.group.id if msg.group else None,
                group_name=msg.group.name if msg.group else None,
                thread_id=msg.group.thread_id if msg.group else None,
            )
            system_prompt = plugin.build_system_prompt(context)
            if system_prompt:
                meta["channel_system_prompt"] = system_prompt
            context_note = plugin.build_context_note(context)
            if context_note:
                query = f"{context_note}\n{safe_content}"

        channel_tools = self.collect_channel_tools(msg.channel_type)
        if channel_tools:
            meta["channel_tools"] = channel_tools

        return meta, query

    # ── invoke (原始字节流) ─────────────────────────────────

    async def invoke(
        self,
        msg: UnifiedMessage,
        dispatch_result: DispatchResult,
        agent_config_id: int,
        current_user,
        db,
    ) -> AsyncIterator[bytes]:
        if self._channel_manager:
            self._channel_manager.inc_active_runs(msg.channel_type, msg.account_id)

        try:
            meta, query = self._build_meta(msg)

            async for chunk in self._stream_fn(
                query=query,
                agent_config_id=agent_config_id,
                thread_id=dispatch_result.thread_id,
                meta=meta,
                image_content=msg.image_base64,
                current_user=current_user,
                db=db,
            ):
                yield chunk
        finally:
            if self._channel_manager:
                self._channel_manager.dec_active_runs(msg.channel_type, msg.account_id)

    # ── invoke_stream (结构化流) ────────────────────────────

    async def invoke_stream(
        self,
        msg: UnifiedMessage,
        dispatch_result: DispatchResult,
        agent_config_id: int,
        current_user,
        db,
        *,
        cancel_token: StreamCancellationToken | None = None,
        on_chunk: Callable[[StreamingChunk], None] | None = None,
        timeout: float = _DEFAULT_TIMEOUT_SECONDS,
    ) -> AsyncIterator[StreamingChunk]:
        if self._channel_manager:
            self._channel_manager.inc_active_runs(msg.channel_type, msg.account_id)

        try:
            meta, query = self._build_meta(msg)

            async def _stream_all() -> AsyncIterator[bytes]:
                async for raw in self._stream_fn(
                    query=query,
                    agent_config_id=agent_config_id,
                    thread_id=dispatch_result.thread_id,
                    meta=meta,
                    image_content=msg.image_base64,
                    current_user=current_user,
                    db=db,
                ):
                    yield raw

            stream_iter = _stream_all()

            async with contextlib.aclosing(stream_iter):
                while True:
                    if cancel_token is not None and cancel_token.is_cancelled:
                        raise StreamCancelledError("stream cancelled by cancel token")

                    try:
                        raw_chunk = await asyncio.wait_for(
                            stream_iter.__anext__(),
                            timeout=timeout,
                        )
                    except TimeoutError:
                        raise StreamTimeoutError(
                            f"No chunk received within {timeout}s"
                        )
                    except StopAsyncIteration:
                        break

                    parsed = self._parse_chunk(raw_chunk)
                    if parsed is None:
                        continue

                    if on_chunk is not None:
                        on_chunk(parsed)

                    yield parsed

            yield StreamingChunk(
                stage=ReplyStage.FINAL,
                chunk_type=ChunkType.STATUS,
                content="finished",
            )
        finally:
            if self._channel_manager:
                self._channel_manager.dec_active_runs(msg.channel_type, msg.account_id)

    # ── invoke_with_result (结构化结果) ─────────────────────

    async def invoke_with_result(
        self,
        msg: UnifiedMessage,
        dispatch_result: DispatchResult,
        agent_config_id: int,
        current_user,
        db,
        *,
        cancel_token: StreamCancellationToken | None = None,
        timeout: float = _DEFAULT_TIMEOUT_SECONDS,
    ) -> StreamResult:
        result = StreamResult(start_time=time.monotonic())
        accumulated: list[str] = []

        try:
            async for chunk in self.invoke_stream(
                msg,
                dispatch_result,
                agent_config_id,
                current_user,
                db,
                cancel_token=cancel_token,
                timeout=timeout,
            ):
                result.chunk_count += 1

                if chunk.chunk_type == ChunkType.TEXT_DELTA:
                    result.text_chunk_count += 1
                    if chunk.content:
                        accumulated.append(chunk.content)
                elif chunk.chunk_type == ChunkType.REASONING:
                    result.reasoning_chunk_count += 1
                elif chunk.chunk_type == ChunkType.TOOL_CALL:
                    result.tool_chunk_count += 1
                elif chunk.chunk_type == ChunkType.TOOL_RESULT:
                    result.tool_chunk_count += 1
                elif chunk.chunk_type == ChunkType.HEARTBEAT:
                    result.heartbeat_count += 1
                elif chunk.chunk_type == ChunkType.CITATION:
                    result.citation_count += 1
                elif chunk.chunk_type == ChunkType.ERROR:
                    result.error_count += 1
                    if chunk.content and not result.error_message:
                        result.error_message = chunk.content

                if chunk.metadata and "token_usage" in chunk.metadata:
                    result.token_usage = chunk.metadata["token_usage"]

            result.accumulated_text = "".join(accumulated)
            result.status = StreamStatus.COMPLETED

        except StreamCancelledError:
            result.status = StreamStatus.CANCELLED
            result.accumulated_text = "".join(accumulated)
            result.error_message = "stream cancelled"
            logger.debug("Agent stream cancelled for channel=%s peer=%s", msg.channel_type, msg.sender.id)
        except StreamTimeoutError:
            result.status = StreamStatus.TIMEOUT
            result.accumulated_text = "".join(accumulated)
            result.error_message = "stream timeout"
            logger.warning("Agent stream timeout for channel=%s peer=%s", msg.channel_type, msg.sender.id)
        except Exception as e:
            result.status = StreamStatus.ERROR
            result.accumulated_text = "".join(accumulated)
            result.error_message = f"{type(e).__name__}: {e}"
            logger.exception("Agent stream error for channel=%s peer=%s", msg.channel_type, msg.sender.id)
        finally:
            result.end_time = time.monotonic()

        return result

    # ── invoke_and_collect (兼容方法) ────────────────────────

    async def invoke_and_collect(
        self,
        msg: UnifiedMessage,
        dispatch_result: DispatchResult,
        agent_config_id: int,
        current_user,
        db,
    ) -> str | None:
        result = await self.invoke_with_result(msg, dispatch_result, agent_config_id, current_user, db)
        return result.accumulated_text or None

    # ── chunk 解析 ─────────────────────────────────────────

    def _parse_chunk(self, raw: bytes) -> StreamingChunk | None:
        try:
            data = json.loads(raw.decode("utf-8"))
        except (json.JSONDecodeError, UnicodeDecodeError):
            logger.debug("Failed to parse stream chunk (len=%d): %s", len(raw), raw[:200])
            return None

        status = data.get("status", "")
        content = data.get("response", "")

        # 流式文本块
        if status in _STREAMING_TEXT_STATUSES:
            return StreamingChunk(
                stage=ReplyStage.BLOCK,
                chunk_type=ChunkType.TEXT_DELTA,
                content=content,
            )

        # 推理/思考块
        if status in _STREAMING_REASONING_STATUSES:
            return StreamingChunk(
                stage=ReplyStage.TOOL,
                chunk_type=ChunkType.REASONING,
                content=content,
            )

        # 工具调用
        if status == "tool_call":
            tool_name = data.get("tool_name", "")
            tool_input = data.get("tool_input", {})
            return StreamingChunk(
                stage=ReplyStage.TOOL,
                chunk_type=ChunkType.TOOL_CALL,
                content=content,
                tool_name=tool_name,
                tool_input=tool_input if isinstance(tool_input, dict) else {},
            )

        # 工具结果
        if status == "tool_result":
            tool_name = data.get("tool_name", "")
            tool_output = data.get("tool_output", content)
            return StreamingChunk(
                stage=ReplyStage.TOOL,
                chunk_type=ChunkType.TOOL_RESULT,
                content=content,
                tool_name=tool_name,
                tool_output=tool_output,
            )

        # 心跳
        if status == "heartbeat":
            return StreamingChunk(
                stage=ReplyStage.BLOCK,
                chunk_type=ChunkType.HEARTBEAT,
                content=content,
            )

        # 引用
        if status == "citation":
            return StreamingChunk(
                stage=ReplyStage.BLOCK,
                chunk_type=ChunkType.CITATION,
                content=content,
            )

        # 错误
        if status == "error":
            return StreamingChunk(
                stage=ReplyStage.FINAL,
                chunk_type=ChunkType.ERROR,
                content=content,
            )

        # 终止态 / 交互态 — 不产出 chunk
        if status in _TERMINAL_STATUSES | _INTERACTIVE_STATUSES:
            return None

        logger.debug("Unknown chunk status: %s", status)
        return None