feat: adds WrongDataDocumentError to classify documents
This commit is contained in:
parent
c75f017eb9
commit
c99b453d96
3 changed files with 31 additions and 0 deletions
|
|
@ -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.models.node_set import NodeSet
|
||||||
from cognee.modules.engine.utils.generate_node_id import generate_node_id
|
from cognee.modules.engine.utils.generate_node_id import generate_node_id
|
||||||
|
from cognee.tasks.documents.exceptions import WrongDataDocumentInputError
|
||||||
|
|
||||||
EXTENSION_TO_DOCUMENT_CLASS = {
|
EXTENSION_TO_DOCUMENT_CLASS = {
|
||||||
"pdf": PdfDocument, # Text documents
|
"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
|
- list[Document]: A list of Document objects created based on the classified data
|
||||||
documents.
|
documents.
|
||||||
"""
|
"""
|
||||||
|
if not isinstance(data_documents, list):
|
||||||
|
raise WrongDataDocumentInputError("data_documents")
|
||||||
|
|
||||||
documents = []
|
documents = []
|
||||||
for data_item in data_documents:
|
for data_item in data_documents:
|
||||||
|
|
||||||
document = EXTENSION_TO_DOCUMENT_CLASS[data_item.extension](
|
document = EXTENSION_TO_DOCUMENT_CLASS[data_item.extension](
|
||||||
id=data_item.id,
|
id=data_item.id,
|
||||||
title=f"{data_item.name}.{data_item.extension}",
|
title=f"{data_item.name}.{data_item.extension}",
|
||||||
|
|
|
||||||
9
cognee/tasks/documents/exceptions/__init__.py
Normal file
9
cognee/tasks/documents/exceptions/__init__.py
Normal file
|
|
@ -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,
|
||||||
|
)
|
||||||
17
cognee/tasks/documents/exceptions/exceptions.py
Normal file
17
cognee/tasks/documents/exceptions/exceptions.py
Normal file
|
|
@ -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)
|
||||||
Loading…
Add table
Reference in a new issue