diff --git a/simplylab/error.py b/simplylab/error.py new file mode 100644 index 0000000..0344d70 --- /dev/null +++ b/simplylab/error.py @@ -0,0 +1,16 @@ +class Error(Exception): + status_code = 500 + + +class MessageLimitedIn30SecondsError(Error): + + def __init__(self): + super().__init__("3 messages limited in 30 seconds") + self.status_code = 401 + + +class MessageLimitedInDailyError(Error): + + def __init__(self): + super().__init__("20 messages limited in daily") + self.status_code = 401 diff --git a/simplylab/main.py b/simplylab/main.py index 914a79b..e45da3b 100644 --- a/simplylab/main.py +++ b/simplylab/main.py @@ -1,7 +1,8 @@ from typing import Union -from fastapi import FastAPI +from fastapi import FastAPI, Request from dotenv import load_dotenv +from starlette.responses import JSONResponse from simplylab.entity import GetAiChatResponseInput from simplylab.entity import GetAiChatResponseOutput @@ -9,13 +10,21 @@ from simplylab.entity import GetUserChatHistoryInput from simplylab.entity import GetUserChatHistoryOutput from simplylab.entity import GetChatStatusTodayInput from simplylab.entity import GetChatStatusTodayOutput -from simplylab.entity import UserChatMessage +from simplylab.error import Error from simplylab.services import Services load_dotenv() app = FastAPI() +@app.exception_handler(Error) +async def error_handler(request: Request, exc: Error): + return JSONResponse( + status_code=exc.status_code, + content={"message": str(exc)} + ) + + @app.get("/") async def read_root(): return {"Hello": "World"}