from __future__ import annotations import json from collections.abc import Iterator from fastapi import APIRouter, Depends, Query from fastapi.responses import StreamingResponse from sqlalchemy.orm import Session from app.core.database import get_db from app.core.dependencies import get_current_user from app.core.responses import api_success from app.models.user import User from app.schemas.chat import ( ChatCompletionRequest, ChatMessageRead, ChatSessionRead, CreateSessionResponse, StopChatRequest, UpdateSessionTitleRequest, ) from app.services.chat_service import ChatService router = APIRouter() @router.post("/session") def create_session(db: Session = Depends(get_db), current_user: User = Depends(get_current_user)) -> dict: session = ChatService.create_session(db, current_user) return api_success(CreateSessionResponse(sessionId=session.id).model_dump()) @router.get("/session/list") def list_sessions(db: Session = Depends(get_db), current_user: User = Depends(get_current_user)) -> dict: sessions = ChatService.list_sessions(db, current_user) return api_success([ChatSessionRead.model_validate(session).model_dump() for session in sessions]) @router.get("/history") def history( sessionId: int = Query(...), db: Session = Depends(get_db), current_user: User = Depends(get_current_user), ) -> dict: messages = ChatService.get_history(db, current_user, sessionId) return api_success([ChatMessageRead.model_validate(message).model_dump(mode="json") for message in messages]) @router.put("/session/title") def update_title( payload: UpdateSessionTitleRequest, db: Session = Depends(get_db), current_user: User = Depends(get_current_user), ) -> dict: session = ChatService.update_title(db, current_user, payload.sessionId, payload.title) return api_success(ChatSessionRead.model_validate(session).model_dump()) @router.delete("/session/{session_id}") def delete_session( session_id: int, db: Session = Depends(get_db), current_user: User = Depends(get_current_user), ) -> dict: ChatService.delete_session(db, current_user, session_id) return api_success() @router.post("/completions") def completions( payload: ChatCompletionRequest, db: Session = Depends(get_db), current_user: User = Depends(get_current_user), ) -> StreamingResponse: answer = ChatService.create_answer(db, current_user, payload.sessionId, payload.message) return StreamingResponse(_sse_chunks(answer), media_type="text/event-stream") @router.post("/stop") def stop( payload: StopChatRequest, db: Session = Depends(get_db), current_user: User = Depends(get_current_user), ) -> dict: ChatService.stop_generation(db, current_user, payload.sessionId) return api_success() def _sse_chunks(answer: str) -> Iterator[str]: chunk_size = 12 for index in range(0, len(answer), chunk_size): chunk = answer[index : index + chunk_size] yield f"data: {json.dumps({'content': chunk}, ensure_ascii=False)}\n\n" yield "data: [DONE]\n\n"