refactor: Split entity related exceptions into graph and database exceptions

Move and split database entity related exceptions into graph and database exceptions

Refactor COG-502

cognee/api/v1/datasets/routers/get_datasets_router.py

@@ -9,7 +9,7 @@ from fastapi.responses import JSONResponse, FileResponse
 from pydantic import BaseModel
 
 from cognee.api.DTO import OutDTO
-from cognee.exceptions import EntityNotFoundError
+from cognee.infrastructure.databases.exceptions import EntityNotFoundError
 from cognee.modules.users.models import User
 from cognee.modules.users.methods import get_authenticated_user
 from cognee.modules.pipelines.models import PipelineRunStatus

cognee/exceptions/__init__.py

@@ -8,9 +8,6 @@ such as service failures, resource conflicts, and invalid operations.
 from .exceptions import (
     CogneeApiError,
     ServiceError,
-    EntityNotFoundError,
-    EntityAlreadyExistsError,
-    InvalidOperationError,
     InvalidValueError,
     InvalidAttributeError,
 )

cognee/exceptions/exceptions.py

@@ -34,42 +34,6 @@ class ServiceError(CogneeApiError):
         super().__init__(message, name, status_code)
 
 
-class EntityNotFoundError(CogneeApiError):
-    """Database returns nothing"""
-
-    def __init__(
-        self,
-        message: str = "The requested entity does not exist.",
-        name: str = "EntityNotFoundError",
-        status_code=status.HTTP_404_NOT_FOUND,
-    ):
-        super().__init__(message, name, status_code)
-
-
-class EntityAlreadyExistsError(CogneeApiError):
-    """Conflict detected, like trying to create a resource that already exists"""
-
-    def __init__(
-        self,
-        message: str = "The entity already exists.",
-        name: str = "EntityAlreadyExistsError",
-        status_code=status.HTTP_409_CONFLICT,
-    ):
-        super().__init__(message, name, status_code)
-
-
-class InvalidOperationError(CogneeApiError):
-    """Invalid operations like trying to delete a non-existing entity, etc."""
-
-    def __init__(
-        self,
-        message: str = "Invalid operation attempted.",
-        name: str = "InvalidOperationError",
-        status_code=status.HTTP_400_BAD_REQUEST,
-    ):
-        super().__init__(message, name, status_code)
-
-
 class InvalidValueError(CogneeApiError):
     def __init__(
             self,
@@ -79,6 +43,7 @@ class InvalidValueError(CogneeApiError):
     ):
         super().__init__(message, name, status_code)
 
+
 class InvalidAttributeError(CogneeApiError):
     def __init__(
             self,

cognee/infrastructure/databases/exceptions/__init__.py

@@ -0,0 +1,10 @@
+"""
+Custom exceptions for the Cognee API.
+
+This module defines a set of exceptions for handling various database errors
+"""
+
+from .exceptions import (
+    EntityNotFoundError,
+    EntityAlreadyExistsError,
+)

cognee/infrastructure/databases/exceptions/exceptions.py

@@ -0,0 +1,25 @@
+from cognee.exceptions import CogneeApiError
+from fastapi import status
+
+class EntityNotFoundError(CogneeApiError):
+    """Database returns nothing"""
+
+    def __init__(
+        self,
+        message: str = "The requested entity does not exist.",
+        name: str = "EntityNotFoundError",
+        status_code=status.HTTP_404_NOT_FOUND,
+    ):
+        super().__init__(message, name, status_code)
+
+
+class EntityAlreadyExistsError(CogneeApiError):
+    """Conflict detected, like trying to create a resource that already exists"""
+
+    def __init__(
+        self,
+        message: str = "The entity already exists.",
+        name: str = "EntityAlreadyExistsError",
+        status_code=status.HTTP_409_CONFLICT,
+    ):
+        super().__init__(message, name, status_code)

cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py

@@ -7,7 +7,7 @@ from sqlalchemy import text, select, MetaData, Table
 from sqlalchemy.orm import joinedload
 from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
 
-from cognee.exceptions import EntityNotFoundError
+from cognee.infrastructure.databases.exceptions import EntityNotFoundError
 from ..ModelBase import Base
 
 class SQLAlchemyAdapter():

cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py

@@ -6,7 +6,8 @@ from sqlalchemy.orm import Mapped, mapped_column
 from sqlalchemy import JSON, Column, Table, select, delete
 from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
 
-from cognee.exceptions import EntityNotFoundError, InvalidValueError
+from cognee.exceptions import InvalidValueError
+from cognee.infrastructure.databases.exceptions import EntityNotFoundError
 from cognee.infrastructure.engine import DataPoint
 
 from .serialize_data import serialize_data

cognee/modules/graph/cognee_graph/CogneeGraph.py

@@ -1,10 +1,11 @@
 from typing import List, Dict, Union
 
-from cognee.exceptions import EntityAlreadyExistsError, EntityNotFoundError, InvalidValueError
+from cognee.exceptions import InvalidValueError
+from cognee.modules.graph.exceptions import EntityNotFoundError, EntityAlreadyExistsError
 from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Node, Edge
 from cognee.modules.graph.cognee_graph.CogneeAbstractGraph import CogneeAbstractGraph
-from cognee.infrastructure.databases.graph import get_graph_engine
+
 
 class CogneeGraph(CogneeAbstractGraph):
     """

cognee/modules/graph/exceptions/__init__.py

@@ -0,0 +1,10 @@
+"""
+Custom exceptions for the Cognee API.
+
+This module defines a set of exceptions for handling various graph errors
+"""
+
+from .exceptions import (
+    EntityNotFoundError,
+    EntityAlreadyExistsError,
+)

cognee/modules/graph/exceptions/exceptions.py

@@ -0,0 +1,25 @@
+from cognee.exceptions import CogneeApiError
+from fastapi import status
+
+class EntityNotFoundError(CogneeApiError):
+    """Database returns nothing"""
+
+    def __init__(
+        self,
+        message: str = "The requested entity does not exist.",
+        name: str = "EntityNotFoundError",
+        status_code=status.HTTP_404_NOT_FOUND,
+    ):
+        super().__init__(message, name, status_code)
+
+
+class EntityAlreadyExistsError(CogneeApiError):
+    """Conflict detected, like trying to create a resource that already exists"""
+
+    def __init__(
+        self,
+        message: str = "The entity already exists.",
+        name: str = "EntityAlreadyExistsError",
+        status_code=status.HTTP_409_CONFLICT,
+    ):
+        super().__init__(message, name, status_code)

cognee/tasks/graph/infer_data_ontology.py

@@ -11,7 +11,7 @@ import aiofiles
 import pandas as pd
 from pydantic import BaseModel
 
-from cognee.exceptions import EntityNotFoundError
+from cognee.modules.graph.exceptions import EntityNotFoundError, EntityAlreadyExistsError
 from cognee.modules.ingestion.exceptions import IngestionError
 from cognee.infrastructure.llm.prompts import read_query_prompt
 from cognee.infrastructure.llm.get_llm_client import get_llm_client