Restore 0.1.5 version from stash

wolai-backend/app/routers/chat.py

@@ -1,55 +1,149 @@
 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__)
 
 
-@router.get("")
-async def chat(
+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] = None,
+    document_id: Optional[str],
+    workspace_id: Optional[str],
+    model: Optional[str],
+    use_web_search: bool,
 ) -> StreamingResponse:
-    """
-    基于 LightRAG 的 SSE 流式回答。
-    query: 必填问题
-    document_id: 可选，指定所属 workspace
-    """
+    """统一封装 GET/POST 的对话逻辑，便于同时支持长文本 POST。"""
     if not query or not query.strip():
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="query 不能为空")
 
-    workspace_id: Optional[str] = None
+    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")
-        workspace_id = (
-            str(document.get("workspace_id")) if document.get("workspace_id") else None
-        )
+        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
 
-    health = await lightrag_service.run_healthcheck()
-    if not health.get("ok"):
-        raise HTTPException(
-            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
-            detail=health.get("error", "LightRAG 未就绪"),
-        )
+    fallback_result: Optional[dict[str, object]] = None
+    lightrag_result: Optional[dict[str, object]] = None
+    web_search_result: Optional[dict[str, object]] = None
 
-    lightrag_result = await lightrag_service.stream_answer(
-        query_text=query, user_id=auth.user_id, workspace_id=workspace_id
-    )
+    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:
@@ -65,3 +159,41 @@ async def chat(
         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,
+    )