app/api/routes/chat.py

from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
from app.models.schemas import ChatRequest, ChatResponse
from app.services.chat_service import chat_service
from app.utils.logger import logger
from app.utils.errors import RagChatbotError
import json

router = APIRouter(prefix="/chat", tags=["chat"])


@router.post("/", response_model=ChatResponse)
async def chat(request: ChatRequest):
    try:
        response = await chat_service.generate_response(
            message=request.message,
            session_id=request.session_id,
            use_rag=request.use_rag,
            use_history=request.use_history
        )
        
        return ChatResponse(
            message=response,
            session_id=request.session_id
        )
        
    except RagChatbotError as e:
        logger.error(f"RAG error: {str(e)}")
        raise HTTPException(status_code=e.status_code, detail=str(e))
    except Exception as e:
        logger.error(f"Chat error: {str(e)}")
        raise HTTPException(status_code=500, detail="Internal server error")


@router.post("/stream")
async def chat_stream(request: ChatRequest):
    try:
        async def generate():
            try:
                async for chunk in chat_service.stream_response(
                    message=request.message,
                    session_id=request.session_id,
                    use_rag=request.use_rag,
                    use_history=request.use_history
                ):
                    yield f"data: {json.dumps({'chunk': chunk})}\n\n"
                
                yield f"data: {json.dumps({'done': True})}\n\n"
                
            except Exception as e:
                logger.error(f"Stream error: {str(e)}")
                yield f"data: {json.dumps({'error': str(e)})}\n\n"
        
        return StreamingResponse(
            generate(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
            }
        )
        
    except RagChatbotError as e:
        logger.error(f"RAG error: {str(e)}")
        raise HTTPException(status_code=e.status_code, detail=str(e))
    except Exception as e:
        logger.error(f"Chat stream error: {str(e)}")
        raise HTTPException(status_code=500, detail="Internal server error")