feat: Add UnstructuredLibraryImportError
Added exception when unstructured libary is called but not installed Feature COG-685
This commit is contained in:
parent
53b7806ccb
commit
07d9330e4a
3 changed files with 27 additions and 1 deletions
9
cognee/modules/data/exceptions/__init__.py
Normal file
9
cognee/modules/data/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 (
|
||||
UnstructuredLibraryImportError,
|
||||
)
|
||||
11
cognee/modules/data/exceptions/exceptions.py
Normal file
11
cognee/modules/data/exceptions/exceptions.py
Normal 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)
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue