feat: Add UnstructuredLibraryImportError

Added exception when unstructured libary is called but not installed

Feature COG-685
This commit is contained in:
Igor Ilic 2024-12-08 14:53:19 +01:00
parent 53b7806ccb
commit 07d9330e4a
3 changed files with 27 additions and 1 deletions

View 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 (
UnstructuredLibraryImportError,
)

View file

@ -0,0 +1,11 @@
from cognee.exceptions import CogneeApiError
from fastapi import status
class UnstructuredLibraryImportError(CogneeApiError):
def __init__(
self,
message: str = "Import error. Unstructured library is not installed.",
name: str = "UnstructuredModuleImportError",
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
):
super().__init__(message, name, status_code)

View file

@ -2,13 +2,19 @@ from io import StringIO
from cognee.modules.chunking.TextChunker import TextChunker
from .Document import Document
from cognee.modules.data.exceptions import UnstructuredLibraryImportError
class UnstructuredDocument(Document):
type: str = "unstructured"
def read(self, chunk_size: int):
def get_text():
from unstructured.partition.auto import partition
try:
from unstructured.partition.auto import partition
except ModuleNotFoundError:
raise UnstructuredLibraryImportError
elements = partition(self.raw_data_location, content_type=self.mime_type)
in_memory_file = StringIO("\n\n".join([str(el) for el in elements]))
in_memory_file.seek(0)