95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
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_mock_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"
|