151 lines
4.2 KiB
Python
151 lines
4.2 KiB
Python
from fastapi import APIRouter, HTTPException, Request, status
|
|
|
|
from auth import CurrentUser
|
|
from schemas import (
|
|
CreateSessionRequest,
|
|
CreateSessionResponse,
|
|
LoginRequest,
|
|
LoginResponse,
|
|
RegisterRequest,
|
|
RegisterResponse,
|
|
SendMessageRequest,
|
|
SendMessageResponse,
|
|
SessionHistoryResponse,
|
|
UserUsageResponse,
|
|
)
|
|
from service import ChatProviderError, ChatService
|
|
from storage import SessionNotFoundError, SessionStorageError
|
|
from users import (
|
|
InvalidCredentialsError,
|
|
UserStore,
|
|
UserStoreError,
|
|
UsernameExistsError,
|
|
)
|
|
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
def get_chat_service(request: Request) -> ChatService:
|
|
return request.app.state.chat_service
|
|
|
|
|
|
def get_user_store(request: Request) -> UserStore:
|
|
return request.app.state.user_store
|
|
|
|
|
|
@router.post(
|
|
"/auth/register",
|
|
response_model=RegisterResponse,
|
|
status_code=status.HTTP_201_CREATED,
|
|
tags=["auth"],
|
|
)
|
|
async def register(
|
|
body: RegisterRequest,
|
|
request: Request,
|
|
) -> RegisterResponse:
|
|
store = get_user_store(request)
|
|
try:
|
|
user = await store.register(body.username, body.password)
|
|
except UsernameExistsError as exc:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_409_CONFLICT,
|
|
detail="username already exists",
|
|
) from exc
|
|
except UserStoreError as exc:
|
|
raise HTTPException(status_code=500, detail=str(exc)) from exc
|
|
return RegisterResponse(user_id=user.user_id, created_at=user.created_at)
|
|
|
|
|
|
@router.post(
|
|
"/auth/login",
|
|
response_model=LoginResponse,
|
|
tags=["auth"],
|
|
)
|
|
async def login(
|
|
body: LoginRequest,
|
|
request: Request,
|
|
) -> LoginResponse:
|
|
store = get_user_store(request)
|
|
try:
|
|
user = await store.login(body.username, body.password)
|
|
except InvalidCredentialsError as exc:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="invalid username or password",
|
|
) from exc
|
|
except UserStoreError as exc:
|
|
raise HTTPException(status_code=500, detail=str(exc)) from exc
|
|
return LoginResponse(user_id=user.user_id, api_key=user.api_key)
|
|
|
|
|
|
@router.post(
|
|
"/sessions",
|
|
response_model=CreateSessionResponse,
|
|
status_code=status.HTTP_201_CREATED,
|
|
tags=["sessions"],
|
|
)
|
|
async def create_session(
|
|
request: Request,
|
|
_user: CurrentUser,
|
|
body: CreateSessionRequest | None = None,
|
|
) -> CreateSessionResponse:
|
|
try:
|
|
return await get_chat_service(request).create_session(
|
|
_user.user_id,
|
|
body.system_prompt if body is not None else None,
|
|
)
|
|
except SessionStorageError as exc:
|
|
raise HTTPException(status_code=500, detail=str(exc)) from exc
|
|
|
|
|
|
@router.post(
|
|
"/sessions/{session_id}/messages",
|
|
response_model=SendMessageResponse,
|
|
tags=["sessions"],
|
|
)
|
|
async def send_message(
|
|
session_id: str,
|
|
body: SendMessageRequest,
|
|
request: Request,
|
|
user: CurrentUser,
|
|
) -> SendMessageResponse:
|
|
try:
|
|
response = await get_chat_service(request).generate_response(
|
|
session_id, body.content, user.user_id
|
|
)
|
|
return response
|
|
except SessionNotFoundError as exc:
|
|
raise HTTPException(status_code=404, detail="session not found") from exc
|
|
except ChatProviderError as exc:
|
|
raise HTTPException(status_code=502, detail=str(exc)) from exc
|
|
except SessionStorageError as exc:
|
|
raise HTTPException(status_code=500, detail=str(exc)) from exc
|
|
|
|
|
|
@router.get(
|
|
"/sessions/{session_id}/messages",
|
|
response_model=SessionHistoryResponse,
|
|
tags=["sessions"],
|
|
)
|
|
async def get_session_history(
|
|
session_id: str,
|
|
request: Request,
|
|
_user: CurrentUser,
|
|
) -> SessionHistoryResponse:
|
|
try:
|
|
return await get_chat_service(request).get_session_history(
|
|
session_id, _user.user_id
|
|
)
|
|
except SessionNotFoundError as exc:
|
|
raise HTTPException(status_code=404, detail="session not found") from exc
|
|
except SessionStorageError as exc:
|
|
raise HTTPException(status_code=500, detail=str(exc)) from exc
|
|
|
|
|
|
@router.get("/usage", response_model=UserUsageResponse, tags=["usage"])
|
|
async def get_current_user_usage(
|
|
request: Request,
|
|
user: CurrentUser,
|
|
) -> UserUsageResponse:
|
|
return await get_chat_service(request).get_user_usage(user.user_id) |