wecom_it_smart_desk/backend/app/tasks/h5_ai_task.py

# =============================================================================
# 企微IT智能服务台 — H5 员工端 AI 回复后台任务
# =============================================================================
# 背景：原 h5_send_message 在同步 HTTP 请求内 await AI 推理（Dify 3~15s），
#       整条请求被阻塞，前端表现为"发送中"长时间卡顿。
#       本模块将 AI 推理移出请求，改为 asyncio 后台任务，结果经 WebSocket
#       流式推回（ai_reply_chunk / ai_reply），发送瞬时完成。
#
# 关键约束（详见 docs/02-需求分析/技术架构演进/员工端消息发送延时改造方案.md）：
#   1. 必须单 worker 运行（docker-compose --workers 1）：
#      ws_manager 是进程内单例，多 worker 时后台任务与员工 WS 连接可能不在
#      同进程，broadcast 会静默丢失（约 50%）。
#   2. 使用独立 DB session（_get_session_factory），不可复用请求的 db
#      （请求返回后该 session 会被关闭）。
# =============================================================================

import logging
from datetime import datetime

from app.database import _get_session_factory
from app.dependencies import get_shared_ai_handler
from app.models.conversation import Conversation
from app.models.message import Message
from app.services.ws_manager import manager as ws_manager

logger = logging.getLogger(__name__)


async def _persist_and_push(
    db,
    conversation: Conversation,
    employee_id: str,
    content: str,
    is_guidance: bool,
    should_count: bool,
    should_transfer: bool,
    dify_conversation_id,
):
    """持久化 AI 回复并推送给员工端 + 广播坐席端。

    做什么：
      1. 存 AI 消息到 DB
      2. 更新会话状态（dify 上下文 / 计数 / 转人工）
      3. 经 WS 向员工推 ai_reply 终态（前端据此替换打字机气泡）
      4. 经 WS 向坐席端广播 new_message + conversation_updated
    为什么：把"落库 + 推送"封装为单点，供同步路径与流式路径复用。
    """
    # 1. 存 AI 消息
    ai_message = Message(
        conversation_id=conversation.id,
        sender_type="ai",
        sender_id="ai_bot",
        sender_name="Duckula（达寇拉）",
        content=content,
        msg_type="text",
        is_read=True,
    )
    db.add(ai_message)
    await db.flush()

    # 2. 更新会话状态
    if dify_conversation_id:
        conversation.dify_conversation_id = dify_conversation_id
    if should_count:
        conversation.ai_substantive_reply_count += 1
    if should_transfer:
        conversation.status = "queued"
    conversation.updated_at = datetime.now()
    db.add(conversation)
    await db.flush()
    await db.commit()

    # 3. 推 ai_reply 终态给员工（前端替换打字机气泡）
    await ws_manager.broadcast_to_employees([employee_id], {
        "type": "ai_reply",
        "data": {
            "message_id": str(ai_message.id),
            "conversation_id": str(conversation.id),
            "sender_type": "ai",
            "sender_id": "ai_bot",
            "sender_name": "Duckula（达寇拉）",
            "content": content,
            "msg_type": "text",
            "is_guidance": is_guidance,
            "ai_reply_count": conversation.ai_substantive_reply_count,
            "can_call_agent": conversation.ai_substantive_reply_count >= 3,
            "conversation_status": conversation.status,
        },
    })

    # 4. 广播坐席端（new_message + conversation_updated）
    try:
        await ws_manager.broadcast({
            "type": "new_message",
            "data": {
                "conversation_id": str(conversation.id),
                "message_id": str(ai_message.id),
                "sender_type": "ai",
                "sender_id": "ai_bot",
                "sender_name": "Duckula（达寇拉）",
                "content": content,
                "msg_type": "text",
            },
        })
        await ws_manager.broadcast({
            "type": "conversation_updated",
            "data": {
                "conversation_id": str(conversation.id),
                "status": conversation.status,
                "assigned_agent_id": str(conversation.assigned_agent_id) if conversation.assigned_agent_id else None,
            },
        })
    except Exception as ws_err:
        # WS 广播失败不阻塞消息存储，只记录 warning
        logger.warning(f"WS 广播 AI 回复给坐席失败（消息已存储）: {ws_err}")


async def process_h5_ai_reply(
    conversation_id: str,
    employee_id: str,
    content: str,
    dify_conversation_id=None,
):
    """H5 发送消息后的 AI 回复处理（asyncio.create_task 入口）。

    流程：
      - 本地快判断（打招呼 / 呼叫人工）→ 同步结果，整段推送（不调 Dify）
      - 否则流式调 Dify，逐 chunk 推 ai_reply_chunk，流结束推 ai_reply 终态
      - 任意异常 → 推 ai_reply_failed，不阻塞用户
    """
    ai_handler = get_shared_ai_handler()
    factory = _get_session_factory()
    async with factory() as db:
        try:
            conversation = await db.get(Conversation, conversation_id)
            if not conversation:
                logger.warning(f"后台 AI 任务：会话不存在 {conversation_id}")
                return

            is_guidance = False
            should_count = False
            should_transfer = False
            new_dify_conv_id = dify_conversation_id
            full_parts: list = []

            # 本地快判断（不打 Dify）：打招呼 / 呼叫人工 → 同步路径
            if ai_handler.is_greeting(content) or ai_handler.is_call_human(content):
                result = await ai_handler.handle_message(
                    content=content,
                    dify_conversation_id=dify_conversation_id,
                    user_id=employee_id,
                )
                await _persist_and_push(
                    db, conversation, employee_id, result.content,
                    result.is_guidance, result.should_count,
                    result.should_transfer, result.dify_conversation_id,
                )
                return

            # 流式调 Dify（get_reply_stream 内部已处理真 SSE / 非流式 fallback）
            # 注意：首参是 message（用户文本），不是 content
            async for chunk in ai_handler.ai_service.get_reply_stream(
                message=content,
                conversation_id=dify_conversation_id,
                user_id=employee_id,
            ):
                delta = chunk.get("delta", "")
                if delta:
                    full_parts.append(delta)
                    await ws_manager.broadcast_to_employees([employee_id], {
                        "type": "ai_reply_chunk",
                        "data": {
                            "conversation_id": conversation_id,
                            "chunk": delta,
                        },
                    })
                if chunk.get("finished"):
                    new_dify_conv_id = chunk.get("conversation_id") or dify_conversation_id
                    hit = chunk.get("hit")
                    # 命中 → 计数；未命中 → 转人工
                    should_count = bool(hit)
                    should_transfer = not bool(hit)

            content_ai = "".join(full_parts)
            if not content_ai:
                # 流式无内容（极端情况），给降级提示，不转人工
                content_ai = "⚠️ AI 暂时没有返回内容，请输入「IT」转人工。"
                should_count = False
                should_transfer = False
            await _persist_and_push(
                db, conversation, employee_id, content_ai,
                is_guidance, should_count, should_transfer, new_dify_conv_id,
            )
        except Exception as e:
            logger.error(f"后台 AI 任务异常: {e}", exc_info=True)
            try:
                await ws_manager.broadcast_to_employees([employee_id], {
                    "type": "ai_reply_failed",
                    "data": {
                        "conversation_id": conversation_id,
                        "message": "⚠️ AI 服务异常，请输入「IT」转人工或稍后重试。",
                    },
                })
            except Exception:
                # 推送失败也无所谓，员工端 3 秒轮询兜底
                pass