From c99b453d965df490b02e370710bb94070f97993d Mon Sep 17 00:00:00 2001
From: hajdul88 <52442977+hajdul88@users.noreply.github.com>
Date: Thu, 14 Aug 2025 10:57:16 +0200
Subject: [PATCH] feat: adds WrongDataDocumentError to classify documents

---
 cognee/tasks/documents/classify_documents.py    |  5 +++++
 cognee/tasks/documents/exceptions/__init__.py   |  9 +++++++++
 cognee/tasks/documents/exceptions/exceptions.py | 17 +++++++++++++++++
 3 files changed, 31 insertions(+)
 create mode 100644 cognee/tasks/documents/exceptions/__init__.py
 create mode 100644 cognee/tasks/documents/exceptions/exceptions.py

diff --git a/cognee/tasks/documents/classify_documents.py b/cognee/tasks/documents/classify_documents.py
index 673e17c75..6567e5825 100644
--- a/cognee/tasks/documents/classify_documents.py
+++ b/cognee/tasks/documents/classify_documents.py
@@ -10,6 +10,7 @@ from cognee.modules.data.processing.document_types import (
 )
 from cognee.modules.engine.models.node_set import NodeSet
 from cognee.modules.engine.utils.generate_node_id import generate_node_id
+from cognee.tasks.documents.exceptions import WrongDataDocumentInputError
 
 EXTENSION_TO_DOCUMENT_CLASS = {
     "pdf": PdfDocument,  # Text documents
@@ -111,8 +112,12 @@ async def classify_documents(data_documents: list[Data]) -> list[Document]:
         - list[Document]: A list of Document objects created based on the classified data
           documents.
     """
+    if not isinstance(data_documents, list):
+        raise WrongDataDocumentInputError("data_documents")
+
     documents = []
     for data_item in data_documents:
+
         document = EXTENSION_TO_DOCUMENT_CLASS[data_item.extension](
             id=data_item.id,
             title=f"{data_item.name}.{data_item.extension}",
diff --git a/cognee/tasks/documents/exceptions/__init__.py b/cognee/tasks/documents/exceptions/__init__.py
new file mode 100644
index 000000000..cdd50c6b0
--- /dev/null
+++ b/cognee/tasks/documents/exceptions/__init__.py
@@ -0,0 +1,9 @@
+"""
+Custom exceptions for the Cognee API.
+
+This module defines a set of exceptions for handling various data errors
+"""
+
+from .exceptions import (
+    WrongDataDocumentInputError,
+)
diff --git a/cognee/tasks/documents/exceptions/exceptions.py b/cognee/tasks/documents/exceptions/exceptions.py
new file mode 100644
index 000000000..a1fcb1d4d
--- /dev/null
+++ b/cognee/tasks/documents/exceptions/exceptions.py
@@ -0,0 +1,17 @@
+from cognee.exceptions import (
+    CogneeValidationError,
+    CogneeConfigurationError,
+)
+from fastapi import status
+
+
+class WrongDataDocumentInputError(CogneeValidationError):
+    """Raised when a wrong data document is provided."""
+    def __init__(
+        self,
+        field: str,
+        name: str = "WrongDataDocumentInputError",
+        status_code: int = status.HTTP_422_UNPROCESSABLE_ENTITY,
+    ):
+        message = f"Missing of invalid parameter: '{field}'."
+        super().__init__(message, name, status_code)
\ No newline at end of file