| import asyncio |
| import json |
| import logging |
| from typing import Optional |
|
|
| from fastapi import APIRouter, HTTPException |
| from fastapi.responses import StreamingResponse |
| from pydantic import BaseModel |
|
|
| from app.service.rag_pipeline import get_answer_from_rag, get_answer_stream_from_rag |
| from app.service.session_manager import session_manager |
|
|
| logger = logging.getLogger(__name__) |
|
|
| router = APIRouter( |
| prefix="/api", |
| tags=["Ask"] |
| ) |
|
|
|
|
| class AskQuestionRequest(BaseModel): |
| question: str |
| context: Optional[str] = None |
| session_id: Optional[str] = None |
| model: Optional[str] = None |
| stream: Optional[bool] = False |
| use_rag: Optional[bool] = None |
|
|
|
|
| class StopAskRequest(BaseModel): |
| session_id: Optional[str] = None |
|
|
|
|
| @router.post("/ask/stream") |
| async def ask_question_stream(request: AskQuestionRequest): |
| """ |
| Stream tokens in real-time using Server-Sent Events (SSE). |
| """ |
| session_id = (request.session_id or "default").strip() |
| current_task = asyncio.current_task() |
| if current_task: |
| session_manager.register_task(session_id, current_task) |
|
|
| logger.info( |
| f"[ask-stream] session_id={session_id!r} " |
| f"model={request.model!r} " |
| f"use_rag={request.use_rag!r} " |
| f"question={request.question[:80]!r}" |
| ) |
|
|
| def event_stream(): |
| try: |
| for token in get_answer_stream_from_rag( |
| question=request.question, |
| context=request.context, |
| session_id=session_id, |
| model=request.model, |
| use_rag=request.use_rag, |
| ): |
| payload = json.dumps({"token": token}) |
| yield f"data: {payload}\n\n" |
| yield "data: [DONE]\n\n" |
| except asyncio.CancelledError: |
| logger.info(f"[ask-stream] Cancelled for session '{session_id}'.") |
| err_payload = json.dumps({"error": f"Generation stopped for session '{session_id}'."}) |
| yield f"data: {err_payload}\n\n" |
| except Exception as e: |
| logger.error(f"[ask-stream] error: {e}") |
| err_payload = json.dumps({"error": str(e)}) |
| yield f"data: {err_payload}\n\n" |
| finally: |
| if current_task: |
| session_manager.unregister_task(session_id, current_task) |
|
|
| return StreamingResponse( |
| event_stream(), |
| media_type="text/event-stream", |
| headers={ |
| "Cache-Control": "no-cache", |
| "Connection": "keep-alive", |
| "X-Accel-Buffering": "no", |
| } |
| ) |
|
|
|
|
| @router.post("/ask") |
| async def ask_question(request: AskQuestionRequest): |
| if request.stream: |
| return await ask_question_stream(request) |
|
|
| session_id = (request.session_id or "default").strip() |
| current_task = asyncio.current_task() |
| if current_task: |
| session_manager.register_task(session_id, current_task) |
|
|
| try: |
| logger.info( |
| f"[ask] session_id={session_id!r} " |
| f"model={request.model!r} " |
| f"use_rag={request.use_rag!r} " |
| f"question={request.question[:80]!r}" |
| ) |
|
|
| answer = get_answer_from_rag( |
| question=request.question, |
| context=request.context, |
| session_id=session_id, |
| model=request.model, |
| use_rag=request.use_rag, |
| ) |
|
|
| logger.info( |
| f"[ask] answer preview: {str(answer)[:120]!r}" |
| ) |
|
|
| return { |
| "status": "success", |
| "answer": answer, |
| } |
|
|
| except asyncio.CancelledError: |
| logger.info(f"[ask] Generation stopped for session '{session_id}'.") |
| raise HTTPException( |
| status_code=499, |
| detail=f"Generation terminated for session '{session_id}'." |
| ) |
| except Exception as e: |
| logger.error(f"[ask] error: {e}") |
| raise HTTPException( |
| status_code=500, |
| detail=str(e) |
| ) |
| finally: |
| if current_task: |
| session_manager.unregister_task(session_id, current_task) |
|
|
|
|
| @router.post("/ask/stop") |
| async def stop_ask( |
| request: Optional[StopAskRequest] = None, |
| session_id: Optional[str] = None, |
| ): |
| """ |
| Cancel active model generation task(s). |
| If session_id is provided, only halts that specific session. |
| """ |
| target_id = (request.session_id if request else None) or session_id |
| if target_id: |
| target_id = target_id.strip() |
| terminated = session_manager.terminate_session(target_id) |
| logger.info(f"[ask] Stop signal for session '{target_id}' (terminated={terminated}).") |
| return { |
| "status": "success", |
| "session_id": target_id, |
| "terminated": terminated, |
| "message": ( |
| f"Generation for session '{target_id}' stopped." |
| if terminated |
| else f"No active generation running for session '{target_id}'." |
| ), |
| } |
|
|
| count = session_manager.terminate_all() |
| return { |
| "status": "success", |
| "message": f"Stopped {count} active generation(s).", |
| "stopped_count": count, |
| } |
