Added proper HTTP error codes

.vscode/settings.json

@@ -0,0 +1,7 @@
+{
+    "python.testing.pytestArgs": [
+        "."
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
+}

app/main.py

@@ -47,6 +47,12 @@ class ChangeDateItem(BaseModel):
     user_id: str
     date: str
 
+class ErrorResponse(BaseModel):
+    status: str = "error"
+    code: int
+    message: str
+    timestamp: datetime = datetime.now()
+
 @app.get("/ok")
 def ok_endpoint():
     return {"message": "ok"}
@@ -56,8 +62,16 @@ def get_user_by_id(user_id: str, api_key: str = Security(get_api_key)):
     try:
         user = get_user(user_id)
         return {"user_info": user.user_info, "user_messages": user.get_messages()}
+    except LookupError:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"User with ID {user_id} not found"
+        )
     except Exception as e:
-        return {"message": f"An error occurred: {e}"}
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=str(e)
+        )
 
 @app.post("/create_user")
 def create_user(request: CreateUserItem, api_key: str = Security(get_api_key)):
@@ -67,22 +81,34 @@ def create_user(request: CreateUserItem, api_key: str = Security(get_api_key)):
         # check if user exists by looking for pickle file in users/data
         if os.path.exists(f'users/data/{request.user_id}.pkl'):
             return {"message": f"[OK] User already exists: {request.user_id}"}
+        
+        user_info, matters_most = get_user_info(request.user_id)
+        if not user_info:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail="Could not fetch user information from DB"
+            )
+
+        user = User(request.user_id, user_info, client, openai_assistant, matters_most)
+        user.save_user()
+
+        # upload user pickle file to s3 bucket 
+        filename = f'{user.user_id}.pkl'
+        upload = upload_file_to_s3(filename)
+
+        if upload == True:
+            return {"message": f"[OK] User created: {user.user_id}"}
         else:
-            # create user
-            user_info, matters_most = get_user_info(request.user_id)
-            user = User(request.user_id, user_info, client, openai_assistant, matters_most)
-            user.save_user()
-
-            # upload user pickle file to s3 bucket 
-            filename = f'{user.user_id}.pkl'
-            response = upload_file_to_s3(filename)
-
-            # move user pickle from to_upload to data
-            if response:
-                return {"message": f"[OK] User created: {user.user_id}"}
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail="Failed to upload user pickle to S3"
+            )        
 
     except Exception as e:
-        return {"message": f"[ERROR]: {e}"}
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=str(e)
+        )
 
 @app.post("/chat")
 def chat(request: ChatItem, api_key: str = Security(get_api_key)):
@@ -91,8 +117,16 @@ def chat(request: ChatItem, api_key: str = Security(get_api_key)):
         user = get_user(request.user_id)
         response = user.send_message(request.message)
         return {"response": response}
+    except LookupError:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"User with ID {request.user_id} not found"
+        )
     except Exception as e:
-        return {"response": f"An error occurred: {e}"}
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=str(e)
+        )
 
 @app.post("/change_date")
 # request = {"date": "2022-01-01 Wed 10:00:00"} i.e. %Y-%m-%d %a %H:%M:%S
@@ -115,13 +149,29 @@ def change_date(request: ChangeDateItem, api_key: str = Security(get_api_key)):
         
         # Update user
         update = update_user(user)
-        if update['response'] == True:
-            return {"response": response}          
-        else:
-            raise ValueError(f"Failed to change day for {user.user_id}")
+        if not update:
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail=f"Failed to update user [{user.user_id}]"
+            )
+            
+        return {"response": response}          
         
+    except ValueError as e:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=str(e)
+        )
+    except LookupError:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"User with ID {request.user_id} not found"
+        )
     except Exception as e:
-        return {"message": f"An error occurred: {e}"}
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=str(e)
+        )
     
 @app.post("/reset_user")
 def reset_user_messages(request: CreateUserItem, api_key: str = Security(get_api_key)):
@@ -131,5 +181,13 @@ def reset_user_messages(request: CreateUserItem, api_key: str = Security(get_api
         user.reset()
         update_user(user)
         return {"response": "ok"}
+    except LookupError:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"User with ID {request.user_id} not found"
+        )
     except Exception as e:
-        return {"message": f"An error occurred: {e}"}
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=str(e)
+        )

app/utils.py

@@ -55,7 +55,7 @@ def get_user(user_id):
             user = User.load_user(user_id, client)
             return user
         else:
-            raise ValueError(f"User {user_id} does not exist") 
+            raise LookupError(f"User [{user_id}] pickle does not exist in S3") 
 
 def get_api_key(api_key_header: str = Security(api_key_header)) -> str:
     if api_key_header in api_keys:
@@ -140,10 +140,10 @@ def upload_file_to_s3(filename):
         )
         return True
     except (FileNotFoundError, NoCredentialsError, PartialCredentialsError) as e:
-        print(f"Error uploading file to S3: {e}")
-        return False
+        return False, e
 
 def download_file_from_s3(filename, bucket):
+    file_path = os.path.join('users', 'data', filename)
     try:
         if AWS_ACCESS_KEY and AWS_SECRET_KEY:
             session = boto3.session.Session(
@@ -155,11 +155,13 @@ def download_file_from_s3(filename, bucket):
             session = boto3.session.Session()
 
         s3_client = session.client('s3')
-        with open(os.path.join('users', 'data', filename), 'wb') as f:
+        with open(file_path, 'wb') as f:
             s3_client.download_fileobj(bucket, f"staging/users/{filename}", f)
         return True
-    except (FileNotFoundError, NoCredentialsError, PartialCredentialsError) as e:
+    except Exception as e:
         print(f"Error downloading file from S3: {e}")
+        if os.path.exists(file_path):
+            os.remove(file_path)
         return False
 
 def update_user(user):
@@ -167,9 +169,9 @@ def update_user(user):
     filename = f'{user.user_id}.pkl'
     if upload_file_to_s3(filename):
         os.remove(os.path.join('users', 'data', filename))
-        return {"response": True}
+        return True
     else:
-        raise ValueError(f"Failed to update {user.user_id}")
+        return False
 
 def upload_mementos_to_db(user_id):
     # Database connection parameters