PBThuong's picture
Perf: add SSE streaming, disable Gemini thinking, skip query rewrite when no history
5354bc3
Raw
History Blame Contribute Delete
2.6 kB
"""
DeepMed-AI — api/v1/endpoints/chat.py
Chat-related endpoints: /chat, /chat/stream, /clear, /new-chat — with rate limiting.
"""
import uuid
from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import StreamingResponse
from slowapi import Limiter
from slowapi.util import get_remote_address
from app.schemas.chat import ChatRequest, ChatResponse
from app.services.chat_service import chat_service
router = APIRouter(tags=["Chat"])
limiter = Limiter(key_func=get_remote_address)
def _get_session_id(request: Request) -> str:
"""Get or create a session ID from X-Session-ID header or cookie session."""
session_id = request.headers.get("X-Session-ID")
if session_id:
return session_id
if "session_id" not in request.session:
request.session["session_id"] = str(uuid.uuid4())
return request.session["session_id"]
@router.post("/chat", response_model=ChatResponse)
@limiter.limit("15/minute") # Prevent abuse: max 15 chat requests per minute per IP
async def chat_endpoint(request: Request, body: ChatRequest):
"""Process a user message through the HybridRAG pipeline."""
if not chat_service.workflow_app:
raise HTTPException(status_code=503, detail="AI system not initialized. Please wait.")
session_id = _get_session_id(request)
return await chat_service.process_message(session_id, body.message)
@router.post("/chat/stream")
@limiter.limit("15/minute")
async def chat_stream_endpoint(request: Request, body: ChatRequest):
"""Stream AI response via Server-Sent Events (SSE)."""
if not chat_service.workflow_app:
raise HTTPException(status_code=503, detail="AI system not initialized. Please wait.")
session_id = _get_session_id(request)
return StreamingResponse(
chat_service.process_message_stream(session_id, body.message),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)
@router.post("/clear")
async def clear_endpoint(req: Request):
"""Clear the in-memory conversation state for the current session."""
chat_service.clear_conversation(_get_session_id(req))
return {"message": "Conversation cleared", "success": True}
@router.post("/new-chat")
async def new_chat_endpoint(req: Request):
"""Create a new chat session with a fresh session ID."""
new_id = str(uuid.uuid4())
req.session["session_id"] = new_id
return {"message": "New chat created", "session_id": new_id, "success": True}