File size: 2,598 Bytes
8eaa451
 
5354bc3
8eaa451
 
 
 
 
5354bc3
8eaa451
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5354bc3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8eaa451
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
"""
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}