68 lines
2.3 KiB
Python
68 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import AsyncIterator, Optional
|
|
|
|
from fastapi import APIRouter, HTTPException, status
|
|
from fastapi.responses import StreamingResponse
|
|
|
|
from app.deps import AuthDep
|
|
from app.services.lightrag_service import lightrag_service
|
|
from app.services.supabase_rest import supabase_rest
|
|
|
|
router = APIRouter(prefix="/chat")
|
|
|
|
|
|
@router.get("")
|
|
async def chat(
|
|
query: str,
|
|
auth: AuthDep,
|
|
document_id: Optional[str] = None,
|
|
) -> StreamingResponse:
|
|
"""
|
|
基于 LightRAG 的 SSE 流式回答。
|
|
query: 必填问题
|
|
document_id: 可选，指定所属 workspace
|
|
"""
|
|
if not query or not query.strip():
|
|
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="query 不能为空")
|
|
|
|
workspace_id: Optional[str] = None
|
|
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
|
|
)
|
|
|
|
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 未就绪"),
|
|
)
|
|
|
|
lightrag_result = await lightrag_service.stream_answer(
|
|
query_text=query, user_id=auth.user_id, workspace_id=workspace_id
|
|
)
|
|
|
|
async def event_stream() -> AsyncIterator[str]:
|
|
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")
|