Merge pull request #43 from Dynamite2003/feature/streaming

feat: add streaming output of LLM

Author: Sskift

backend/app/api/v1/endpoints/academic.py

@@ -1,4 +1,10 @@
+import asyncio
+import json
+import logging
+from typing import Any, AsyncIterator, Awaitable
+
 from fastapi import APIRouter, Depends, HTTPException
+from fastapi.responses import StreamingResponse
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from app.db.conversation_repository import ConversationRepository
@@ -16,6 +22,7 @@
 from app.services.academic.search_service import AcademicSearchService
 
 router = APIRouter()
+logger = logging.getLogger(__name__)
 
 @router.post("/search", response_model=PaperSearchResponse)
 async def search_papers(
@@ -42,54 +49,17 @@ async def intelligent_search(
     # 对 IntelligentSearchService.__init__ 的可选依赖参数生成 OpenAPI schema。
     service = IntelligentSearchService()
 
-    try:
-        response = await service.search(request)
-
-        # 如果用户已登录，保存对话历史
-        if current_user:
-            repo = ConversationRepository(db)
-
-            # 如果没有指定conversation_id，创建新对话
-            conversation_id = request.conversation_id
-            if not conversation_id:
-                # 使用查询的前50个字符作为标题
-                title = request.query[:50] if len(request.query) <= 50 else request.query[:47] + "..."
-                conversation = await repo.create_conversation(
-                    current_user.id,
-                    ConversationCreate(title=title)
-                )
-                conversation_id = conversation.id
-
-            # 保存用户消息
-            await repo.add_message(
-                conversation_id,
-                ConversationMessageCreate(
-                    role="user",
-                    content=request.query,
-                )
-            )
+    if request.stream:
+        return StreamingResponse(
+            _stream_intelligent_search(service, request, current_user, db),
+            media_type="text/event-stream",
+        )
 
-            # 保存AI回复
-            assistant_content = response.answer.response if response.answer else ""
-            await repo.add_message(
-                conversation_id,
-                ConversationMessageCreate(
-                    role="assistant",
-                    content=assistant_content,
-                    extra_data={
-                        "papers_count": len(response.papers) if response.papers else 0,
-                        "search_performed": response.search_performed,
-                    }
-                )
-            )
-
-            # 将conversation_id添加到响应中（扩展响应模型）
-            # 注意：这需要修改IntelligentSearchResponse schema
-            response_dict = response.model_dump()
-            response_dict["conversation_id"] = conversation_id
-            return response_dict
-
-        return response
+    try:
+        payload, persist = await _prepare_stream_payload(response, request, current_user, db)
+        if persist:
+            await persist()
+        return payload
     except ValueError as exc:
         raise HTTPException(status_code=400, detail=str(exc)) from exc
 
@@ -102,3 +72,102 @@ async def get_paper_details(paper_id: str):
 async def get_recommendations(domain: str):
     """根据研究领域推荐论文"""
     pass
+
+
+async def _prepare_stream_payload(
+    response: IntelligentSearchResponse,
+    request: IntelligentSearchRequest,
+    current_user: User | None,
+    db: AsyncSession,
+) -> tuple[Any, Awaitable[None] | None]:
+    """Return a payload ready for serialization plus a persistence coroutine."""
+
+    if not current_user:
+        return response, None
+
+    repo = ConversationRepository(db)
+    conversation_id = request.conversation_id
+    if conversation_id:
+        conversation = await repo.get_conversation(conversation_id, current_user.id)
+        if not conversation or conversation.category != "search":
+            raise HTTPException(
+                status_code=404,
+                detail="Conversation not found for intelligent search",
+            )
+    else:
+        title = request.query[:50] if len(request.query) <= 50 else request.query[:47] + "..."
+        conversation = await repo.create_conversation(
+            current_user.id,
+            ConversationCreate(title=title, category="search"),
+        )
+        conversation_id = conversation.id
+
+    payload = response.model_dump()
+    payload["conversation_id"] = conversation_id
+
+    async def persist() -> None:
+        await repo.add_message(
+            conversation_id,
+            ConversationMessageCreate(
+                role="user",
+                content=request.query,
+            ),
+        )
+        assistant_content = response.answer.response if response.answer else ""
+        await repo.add_message(
+            conversation_id,
+            ConversationMessageCreate(
+                role="assistant",
+                content=assistant_content,
+                extra_data={
+                    "papers_count": len(response.papers) if response.papers else 0,
+                    "search_performed": response.search_performed,
+                },
+            ),
+        )
+
+    return payload, persist
+
+
+def _format_sse(event: dict[str, Any]) -> bytes:
+    serialized = json.dumps(event, ensure_ascii=False)
+    return f"data: {serialized}\n\n".encode("utf-8")
+
+
+async def _stream_intelligent_search(
+    service: IntelligentSearchService,
+    request: IntelligentSearchRequest,
+    current_user: User | None,
+    db: AsyncSession,
+) -> AsyncIterator[bytes]:
+    queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue()
+
+    async def push_delta(token: str) -> None:
+        await queue.put({"type": "token", "content": token})
+
+    async def produce() -> None:
+        try:
+            response = await service.search(request, stream_callback=push_delta)
+            payload, persist = await _prepare_stream_payload(response, request, current_user, db)
+            payload_dict = payload if isinstance(payload, dict) else payload.model_dump()
+            await queue.put({"type": "result", "payload": payload_dict})
+            if persist:
+                try:
+                    await persist()
+                except Exception as exc:  # pragma: no cover
+                    logger.exception("Failed to persist streaming conversation: %s", exc)
+        except Exception as exc:  # pragma: no cover - streaming path
+            logger.exception("Intelligent search streaming失败: %s", exc)
+            await queue.put({"type": "error", "message": str(exc)})
+        finally:
+            await queue.put(None)
+
+    producer = asyncio.create_task(produce())
+    try:
+        while True:
+            event = await queue.get()
+            if event is None:
+                break
+            yield _format_sse(event)
+    finally:
+        await producer

backend/app/api/v1/endpoints/conversations.py

@@ -1,5 +1,5 @@
 """Conversation API endpoints."""
-from typing import List
+from typing import List, Literal, Optional
 
 from fastapi import APIRouter, Depends, HTTPException, Query, status
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -37,6 +37,9 @@ async def create_conversation(
 async def list_conversations(
     skip: int = Query(0, ge=0),
     limit: int = Query(50, ge=1, le=100),
+    category: Optional[Literal["search", "reading"]] = Query(
+        None, description="对话类别过滤：search 或 reading"
+    ),
     current_user: User = Depends(get_current_user),
     db: AsyncSession = Depends(get_db),
 ):
@@ -46,6 +49,7 @@ async def list_conversations(
         current_user.id,
         skip=skip,
         limit=limit,
+        category=category,
     )
 
     # 为每个对话添加额外信息
@@ -67,6 +71,7 @@ async def list_conversations(
             "id": conv.id,
             "user_id": conv.user_id,
             "title": conv.title,
+            "category": conv.category,
             "created_at": conv.created_at,
             "updated_at": conv.updated_at,
             "is_deleted": conv.is_deleted,
@@ -81,14 +86,17 @@ async def list_conversations(
 @router.get("/{conversation_id}", response_model=ConversationDetail)
 async def get_conversation(
     conversation_id: int,
+    category: Optional[Literal["search", "reading"]] = Query(
+        None, description="过滤对话类别"
+    ),
     current_user: User = Depends(get_current_user),
     db: AsyncSession = Depends(get_db),
 ):
     """获取特定对话详情"""
     repo = ConversationRepository(db)
     conversation = await repo.get_conversation(conversation_id, current_user.id)
 
-    if not conversation:
+    if not conversation or (category and conversation.category != category):
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
             detail="对话不存在或无权访问",
@@ -101,14 +109,17 @@ async def get_conversation(
 async def update_conversation(
     conversation_id: int,
     data: ConversationUpdate,
+    category: Optional[Literal["search", "reading"]] = Query(
+        None, description="过滤对话类别"
+    ),
     current_user: User = Depends(get_current_user),
     db: AsyncSession = Depends(get_db),
 ):
     """更新对话信息"""
     repo = ConversationRepository(db)
     conversation = await repo.update_conversation(conversation_id, current_user.id, data)
 
-    if not conversation:
+    if not conversation or (category and conversation.category != category):
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
             detail="对话不存在或无权访问",
@@ -128,6 +139,7 @@ async def update_conversation(
         "id": conversation.id,
         "user_id": conversation.user_id,
         "title": conversation.title,
+        "category": conversation.category,
         "created_at": conversation.created_at,
         "updated_at": conversation.updated_at,
         "is_deleted": conversation.is_deleted,
@@ -141,26 +153,32 @@ async def update_conversation(
 @router.delete("/{conversation_id}", status_code=status.HTTP_204_NO_CONTENT)
 async def delete_conversation(
     conversation_id: int,
+    category: Optional[Literal["search", "reading"]] = Query(
+        None, description="过滤对话类别"
+    ),
     current_user: User = Depends(get_current_user),
     db: AsyncSession = Depends(get_db),
 ):
     """删除对话"""
     repo = ConversationRepository(db)
-    success = await repo.delete_conversation(conversation_id, current_user.id)
-
-    if not success:
+    conversation = await repo.get_conversation(conversation_id, current_user.id)
+    if not conversation or (category and conversation.category != category):
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
             detail="对话不存在或无权访问",
         )
 
+    await repo.delete_conversation(conversation_id, current_user.id)
     return None
 
 
 @router.post("/{conversation_id}/messages", response_model=ConversationMessage, status_code=status.HTTP_201_CREATED)
 async def add_message(
     conversation_id: int,
     data: ConversationMessageCreate,
+    category: Optional[Literal["search", "reading"]] = Query(
+        None, description="过滤对话类别"
+    ),
     current_user: User = Depends(get_current_user),
     db: AsyncSession = Depends(get_db),
 ):
@@ -169,7 +187,7 @@ async def add_message(
 
     # 验证对话存在且属于当前用户
     conversation = await repo.get_conversation(conversation_id, current_user.id)
-    if not conversation:
+    if not conversation or (category and conversation.category != category):
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
             detail="对话不存在或无权访问",