From a3fd079fd6429cee7edea09b025c03f66626f12b Mon Sep 17 00:00:00 2001
From: Igor Ilic <igorilic03@gmail.com>
Date: Tue, 26 Nov 2024 19:26:38 +0100
Subject: [PATCH] feat: Add Exceptions and exception handlers

Add classes for exceptions and add exception handling

Feature COG-502
---
 cognee/api/client.py            | 19 ++++++++++
 cognee/exceptions/__init__.py   |  0
 cognee/exceptions/exceptions.py | 67 +++++++++++++++++++++++++++++++++
 3 files changed, 86 insertions(+)
 create mode 100644 cognee/exceptions/__init__.py
 create mode 100644 cognee/exceptions/exceptions.py

diff --git a/cognee/api/client.py b/cognee/api/client.py
index 53c9f9762..39dd72506 100644
--- a/cognee/api/client.py
+++ b/cognee/api/client.py
@@ -7,6 +7,9 @@ from fastapi import FastAPI
 from fastapi.responses import JSONResponse, Response
 from fastapi.middleware.cors import CORSMiddleware
 
+from cognee.exceptions.exceptions import CogneeApiError
+from traceback import format_exc
+
 # Set up logging
 logging.basicConfig(
     level=logging.INFO,  # Set the logging level (e.g., DEBUG, INFO, WARNING, ERROR, CRITICAL)
@@ -76,6 +79,22 @@ async def request_validation_exception_handler(request: Request, exc: RequestVal
         content = jsonable_encoder({"detail": exc.errors(), "body": exc.body}),
     )
 
+@app.exception_handler(CogneeApiError)
+async def exception_handler(_: Request, exc: CogneeApiError) -> JSONResponse:
+    #TODO: Add checking if all values exist for exception
+    detail = {}
+    if exc.message:
+        detail["message"] = exc.message
+
+    if exc.name:
+        detail["message"] = f"{detail['message']} [{exc.name}]"
+
+    # log the stack trace for easier serverside debugging
+    logger.error(format_exc())
+    return JSONResponse(
+        status_code=exc.status_code, content={"detail": detail["message"]}
+    )
+
 app.include_router(
     get_auth_router(),
     prefix = "/api/v1/auth",
diff --git a/cognee/exceptions/__init__.py b/cognee/exceptions/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/cognee/exceptions/exceptions.py b/cognee/exceptions/exceptions.py
new file mode 100644
index 000000000..ed90db735
--- /dev/null
+++ b/cognee/exceptions/exceptions.py
@@ -0,0 +1,67 @@
+from fastapi import status
+
+
+class CogneeApiError(Exception):
+    """Base exception class"""
+
+    def __init__(
+        self,
+        message: str = "Service is unavailable",
+        name: str = "Cognee",
+        status_code=status.HTTP_418_IM_A_TEAPOT,
+    ):
+        self.message = message
+        self.name = name
+        self.status_code = status_code
+        super().__init__(self.message, self.name)
+
+
+class ServiceError(CogneeApiError):
+    """Failures in external services or APIs, like a database or a third-party service"""
+
+    def __init__(
+        self,
+        message: str = "Service is unavailable",
+        name: str = "ServiceError",
+        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+    ):
+        self.message = message
+        self.name = name
+        self.status_code = status_code
+        super().__init__(self.message, self.name)
+
+
+class EntityDoesNotExistError(CogneeApiError):
+    """Database returns nothing"""
+
+    pass
+
+
+class GroupNotFound(CogneeApiError):
+    """User group not found"""
+
+    pass
+
+
+class EntityAlreadyExistsError(CogneeApiError):
+    """Conflict detected, like trying to create a resource that already exists"""
+
+    pass
+
+
+class InvalidOperationError(CogneeApiError):
+    """Invalid operations like trying to delete a non-existing entity, etc."""
+
+    pass
+
+
+class AuthenticationFailed(CogneeApiError):
+    """Invalid authentication credentials"""
+
+    pass
+
+
+class InvalidTokenError(CogneeApiError):
+    """Invalid token"""
+
+    pass