error handler

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

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"}