shudao-main/shudao-chat-py/routers/report_compat.py

"""
报告兼容路由
完全对齐 Go 版本的接口实现，保持外部一致性
"""
from fastapi import APIRouter, Request
from fastapi.responses import StreamingResponse, JSONResponse
import httpx
import json
import time
from typing import AsyncGenerator
from models.report import (
    ReportCompleteFlowRequest,
    UpdateAIMessageRequest,
    StopSSERequest,
    StreamChatRequest
)
from services.aichat_proxy import aichat_proxy
from utils.token import is_local_token
from utils.config import settings
from utils.logger import logger
from database import get_db
from sqlalchemy.orm import Session
from models.chat import AIMessage

router = APIRouter(prefix="/apiv1", tags=["报告兼容"])


def get_request_token(request: Request) -> str:
    """获取请求中的 token"""
    # 支持多种 header 名称
    for header_name in ["token", "Token", "Authorization"]:
        header_value = request.headers.get(header_name, "").strip()
        if header_value:
            if header_name == "Authorization" and header_value.startswith("Bearer "):
                return header_value.replace("Bearer ", "")
            return header_value
    return ""


def should_proxy_to_aichat(token: str) -> bool:
    """判断是否应该代理到 aichat"""
    if not token:
        return False
    
    # 本地 token 不代理，外部 token 代理
    return not is_local_token(token)


def _ensure_proxy_conversation_id(
    request_data: ReportCompleteFlowRequest,
    request: Request
) -> int:
    """
    为 aichat 生成可用的 ai_conversation_id。
    目的：避免 ai_conversation_id=0 触发 aichat 内部 transfer_session 崩溃分支。
    """
    conversation_id = int(request_data.ai_conversation_id or 0)
    if conversation_id > 0:
        return conversation_id

    # 生成一个正数会话ID（<= 2^53，前端 JS 可安全表示）
    generated_id = int(time.time() * 1_000_000)

    # 叠加少量用户特征，降低并发冲突概率
    user = getattr(request.state, "user", None)
    user_id = int(getattr(user, "user_id", 0) or 0)
    if user_id > 0:
        generated_id = generated_id - (generated_id % 1000) + (user_id % 1000)

    if generated_id <= 0:
        generated_id = int(time.time()) or 1

    return generated_id


def _build_aichat_complete_flow_body(
    request_data: ReportCompleteFlowRequest,
    request: Request
) -> bytes:
    """构建转发到 aichat 的请求体，并规避 ai_conversation_id=0 的崩溃场景。"""
    payload = request_data.dict()
    payload["ai_conversation_id"] = _ensure_proxy_conversation_id(request_data, request)
    return json.dumps(payload, ensure_ascii=False).encode("utf-8")


async def fallback_to_local_stream(
    request_data: ReportCompleteFlowRequest,
    request: Request
) -> StreamingResponse:
    """降级为本地 SSE 兼容输出（匹配前端 report/complete-flow 事件格式）。"""
    logger.info("[报告兼容] 降级为本地 SSE 兼容输出")

    stream_request = StreamChatRequest(
        message=request_data.user_question,
        ai_conversation_id=request_data.ai_conversation_id,
        business_type=0
    )

    local_url = f"http://127.0.0.1:{settings.app.port}/apiv1/stream/chat-with-db"

    async def stream_generator() -> AsyncGenerator[bytes, None]:
        ai_conversation_id = request_data.ai_conversation_id or 0
        ai_message_id = 0
        full_response = ""

        try:
            headers = {"Content-Type": "application/json"}
            for header_name in ["Authorization", "Token", "token"]:
                if header_value := request.headers.get(header_name):
                    headers[header_name] = header_value

            async with httpx.AsyncClient(timeout=600) as client:
                async with client.stream(
                    "POST",
                    local_url,
                    json=stream_request.dict(),
                    headers=headers
                ) as response:
                    if response.status_code != 200:
                        error_msg = f"data: {{\"type\": \"online_error\", \"message\": \"Local stream failed: {response.status_code}\"}}\n\n"
                        yield error_msg.encode("utf-8")
                        yield b"data: {\"type\": \"completed\"}\n\n"
                        return

                    async for raw_chunk in response.aiter_bytes(chunk_size=4096):
                        if not raw_chunk:
                            continue
                        chunk_text = raw_chunk.decode("utf-8", errors="ignore")
                        for line in chunk_text.split("\n"):
                            if not line.startswith("data: "):
                                continue
                            payload = line[6:].strip()
                            if not payload:
                                continue

                            if payload == "[DONE]":
                                if full_response:
                                    online_answer = {
                                        "type": "online_answer",
                                        "content": full_response,
                                        "ai_conversation_id": ai_conversation_id,
                                        "ai_message_id": ai_message_id,
                                    }
                                    yield f"data: {json.dumps(online_answer, ensure_ascii=False)}\n\n".encode("utf-8")
                                yield b"data: {\"type\": \"completed\"}\n\n"
                                return

                            if payload.startswith("{"):
                                try:
                                    data = json.loads(payload)
                                except Exception:
                                    data = None
                                if isinstance(data, dict) and data.get("type") == "initial":
                                    ai_conversation_id = data.get("ai_conversation_id") or ai_conversation_id
                                    ai_message_id = data.get("ai_message_id") or ai_message_id
                                    continue

                            full_response += payload.replace("\\n", "\n")

        except Exception as e:
            logger.error(f"[报告兼容] 本地 SSE 处理异常: {e}")
            error_msg = f"data: {{\"type\": \"online_error\", \"message\": \"Local stream error: {str(e)}\"}}\n\n"
            yield error_msg.encode("utf-8")
            yield b"data: {\"type\": \"completed\"}\n\n"

    return StreamingResponse(
        stream_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",
        }
    )


@router.post("/report/complete-flow")
async def complete_flow(request: Request):
    """
    完整报告生成流程（SSE）
    完全对齐 Go 版本的实现
    """
    # 解析请求体
    request_body = await request.body()
    
    try:
        request_data = ReportCompleteFlowRequest(**json.loads(request_body))
    except Exception as e:
        return StreamingResponse(
            iter([
                f"data: {{\"type\": \"online_error\", \"message\": \"Request parse error: {str(e)}\"}}\n\n".encode('utf-8'),
                b"data: {\"type\": \"completed\"}\n\n"
            ]),
            media_type="text/event-stream"
        )
    
    # 验证问题不为空
    if not request_data.user_question.strip():
        return StreamingResponse(
            iter([
                b"data: {\"type\": \"online_error\", \"message\": \"Question cannot be empty\"}\n\n",
                b"data: {\"type\": \"completed\"}\n\n"
            ]),
            media_type="text/event-stream"
        )

    token = get_request_token(request)

    if should_proxy_to_aichat(token):
        proxy_body = _build_aichat_complete_flow_body(request_data, request)

        if int(request_data.ai_conversation_id or 0) <= 0:
            proxy_data = json.loads(proxy_body.decode("utf-8"))
            logger.info(
                f"[报告兼容] ai_conversation_id=0 已重写为 {proxy_data.get('ai_conversation_id')}，代理到 aichat"
            )
        else:
            logger.info("[报告兼容] 代理 complete-flow 到 aichat")

        try:
            return await aichat_proxy.proxy_sse("/report/complete-flow", request, proxy_body)
        except Exception as e:
            logger.error(f"[报告兼容] 代理 complete-flow 到 aichat 失败，降级本地: {e}")

    # 本地 token 或代理失败：降级到本地兼容 SSE
    return await fallback_to_local_stream(request_data, request)


@router.post("/report/update-ai-message")
async def update_ai_message(request: Request):
    """
    更新 AI 消息内容
    完全对齐 Go 版本的实现
    """
    request_body = await request.body()
    
    # 获取 token 并判断路由策略
    token = get_request_token(request)
    
    if should_proxy_to_aichat(token):
        # 外部 token，代理到 aichat
        logger.info("[报告兼容] 代理更新消息到 aichat")
        try:
            return await aichat_proxy.proxy_json("/report/update-ai-message", request, request_body)
        except Exception as e:
            logger.error(f"[报告兼容] 代理更新消息失败: {e}")
            # 降级到本地处理
    
    # 本地处理
    try:
        request_data = UpdateAIMessageRequest(**json.loads(request_body))
    except Exception as e:
        return JSONResponse(
            content={"success": False, "message": f"Request parse error: {str(e)}"},
            status_code=400
        )
    
    if request_data.ai_message_id == 0:
        return JSONResponse(
            content={"success": False, "message": "ai_message_id cannot be empty"},
            status_code=400
        )
    
    # 更新数据库
    try:
        db: Session = next(get_db())
        db.query(AIMessage).filter(
            AIMessage.id == request_data.ai_message_id,
            AIMessage.is_deleted == 0
        ).update({"content": request_data.content})
        db.commit()
        
        return JSONResponse(
            content={"success": True, "message": "AI message updated"},
            status_code=200
        )
    except Exception as e:
        logger.error(f"[报告兼容] 更新消息失败: {e}")
        return JSONResponse(
            content={"success": False, "message": f"Update failed: {str(e)}"},
            status_code=500
        )


@router.post("/sse/stop")
async def stop_sse(request: Request):
    """
    停止 SSE 流
    完全对齐 Go 版本的实现
    """
    request_body = await request.body()
    
    # 获取 token 并判断路由策略
    token = get_request_token(request)
    
    if should_proxy_to_aichat(token):
        # 外部 token，代理到 aichat
        logger.info("[报告兼容] 代理停止请求到 aichat")
        try:
            return await aichat_proxy.proxy_json("/sse/stop", request, request_body)
        except Exception as e:
            logger.error(f"[报告兼容] 代理停止请求失败: {e}")
            # 降级到本地处理
    
    # 本地处理（简单返回成功）
    try:
        request_data = StopSSERequest(**json.loads(request_body))
        return JSONResponse(
            content={
                "success": True,
                "message": "Stop request received",
                "ai_conversation_id": request_data.ai_conversation_id
            },
            status_code=200
        )
    except Exception as e:
        return JSONResponse(
            content={"success": False, "message": f"Request parse error: {str(e)}"},
            status_code=400
        )