mnote/wolai-backend/app/routers/chat.py

from __future__ import annotations

import json
import logging
from typing import AsyncIterator, Optional

from fastapi import APIRouter, HTTPException, status
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from app.deps import AuthDep
from app.services.lightrag_service import lightrag_service
from app.services.simple_ai_service import simple_ai_service
from app.services.searxng_client import searxng_client
from app.services.supabase_rest import supabase_rest

router = APIRouter(prefix="/chat")
logger = logging.getLogger(__name__)


class ChatRequest(BaseModel):
    query: str
    document_id: Optional[str] = None
    workspace_id: Optional[str] = None
    model: Optional[str] = None
    use_web_search: bool = False


async def _build_chat_response(
    query: str,
    auth: AuthDep,
    document_id: Optional[str],
    workspace_id: Optional[str],
    model: Optional[str],
    use_web_search: bool,
) -> StreamingResponse:
    """统一封装 GET/POST 的对话逻辑，便于同时支持长文本 POST。"""
    if not query or not query.strip():
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="query 不能为空")

    resolved_workspace: Optional[str] = None
    if workspace_id:
        membership = supabase_rest.select_one(
            "workspace_members", {"workspace_id": workspace_id, "user_id": auth.user_id}
        )
        if not membership:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN, detail="无权访问该 workspace"
            )
        resolved_workspace = workspace_id

    if document_id:
        document = supabase_rest.select_one(
            "documents", {"id": document_id, "user_id": auth.user_id}
        )
        if not document:
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found")
        doc_workspace = str(document.get("workspace_id")) if document.get("workspace_id") else None
        if resolved_workspace and doc_workspace and resolved_workspace != doc_workspace:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST, detail="workspace 与文档不匹配"
            )
        resolved_workspace = resolved_workspace or doc_workspace

    fallback_result: Optional[dict[str, object]] = None
    lightrag_result: Optional[dict[str, object]] = None
    web_search_result: Optional[dict[str, object]] = None

    fallback_reason: Optional[str] = None
    # 联网搜索模式：直接用搜索上下文 + DeepSeek/Ollama 生成
    if use_web_search:
        search_results = searxng_client.search(query)
        context = "\n\n".join(
            [f"[{idx+1}] {item['title']}\n{item.get('snippet','')}\n{item['url']}" for idx, item in enumerate(search_results)]
        )
        web_search_result = await lightrag_service.answer_with_context(
            query_text=query,
            context=context or "未获取到搜索结果",
            prefer_deepseek=True,
        )
        # 将搜索结果作为引用
        if web_search_result is not None:
            web_search_result["references"] = search_results
    else:
        health = await lightrag_service.run_healthcheck()
        if not health.get("ok"):
            fallback_reason = health.get("error", "LightRAG 未就绪")
        else:
            try:
                lightrag_result = await lightrag_service.stream_answer(
                    query_text=query,
                    user_id=auth.user_id,
                    workspace_id=resolved_workspace,
                    model_choice=model,
                    document_id=document_id,
                )
            except Exception as exc:  # pragma: no cover - 运行时保护
                fallback_reason = str(exc)
                logger.exception("LightRAG stream_answer failed, fallback to simple summary: %s", fallback_reason)

        if fallback_reason:
            fallback_result = simple_ai_service.summarize_document(
                query=query,
                user_id=auth.user_id,
                document_id=document_id,
                workspace_id=resolved_workspace,
            )
            logger.warning("LightRAG unavailable, use simple summary. reason=%s", fallback_reason)

    async def event_stream() -> AsyncIterator[str]:
        if web_search_result is not None:
            if web_search_result.get("is_streaming") and web_search_result.get("iterator"):
                iterator = web_search_result["iterator"]
                async for chunk in iterator:
                    payload = json.dumps({"type": "chunk", "content": chunk})
                    yield f"data: {payload}\n\n"
            else:
                payload = json.dumps(
                    {"type": "chunk", "content": web_search_result.get("content", "")}
                )
                yield f"data: {payload}\n\n"
            references = web_search_result.get("references", []) if isinstance(web_search_result, dict) else []
            yield f"data: {json.dumps({'type': 'references', 'data': references})}\n\n"
            yield "data: [DONE]\n\n"
            return

        if fallback_result is not None:
            payload = json.dumps(
                {
                    "type": "chunk",
                    "content": fallback_result.get("content", ""),
                }
            )
            yield f"data: {payload}\n\n"
            references = fallback_result.get("references", [])
            yield f"data: {json.dumps({'type': 'references', 'data': references})}\n\n"
            yield "data: [DONE]\n\n"
            return

        if lightrag_result is None:
            payload = json.dumps({"type": "chunk", "content": "LightRAG 暂不可用"})
            yield f"data: {payload}\n\n"
            yield f"data: {json.dumps({'type': 'references', 'data': []})}\n\n"
            yield "data: [DONE]\n\n"
            return

        if lightrag_result.get("is_streaming") and lightrag_result.get("iterator"):
            iterator = lightrag_result["iterator"]
            async for chunk in iterator:
                payload = json.dumps({"type": "chunk", "content": chunk})
                yield f"data: {payload}\n\n"
        else:
            payload = json.dumps(
                {"type": "chunk", "content": lightrag_result.get("content", "")}
            )
            yield f"data: {payload}\n\n"
        references = lightrag_result.get("references", [])
        yield f"data: {json.dumps({'type': 'references', 'data': references})}\n\n"
        yield "data: [DONE]\n\n"

    return StreamingResponse(event_stream(), media_type="text/event-stream")


@router.get("")
async def chat(
    query: str,
    auth: AuthDep,
    document_id: Optional[str] = None,
    workspace_id: Optional[str] = None,
    model: Optional[str] = None,
    use_web_search: bool = False,
) -> StreamingResponse:
    """
    基于 LightRAG 的 SSE 流式回答。
    query: 必填问题
    document_id: 可选，指定所属 workspace
    workspace_id: 可选，直接指定 workspace，优先级高于 document_id
    """
    return await _build_chat_response(
        query=query,
        auth=auth,
        document_id=document_id,
        workspace_id=workspace_id,
        model=model,
        use_web_search=use_web_search,
    )


@router.post("")
async def chat_post(payload: ChatRequest, auth: AuthDep) -> StreamingResponse:
    """POST 版本，适配长问题与 JSON 传参。"""
    return await _build_chat_response(
        query=payload.query,
        auth=auth,
        document_id=payload.document_id,
        workspace_id=payload.workspace_id,
        model=payload.model,
        use_web_search=payload.use_web_search,
    )