Repository: topoteretes/cognee
File: cognee/tasks/ingestion/save_data_item_to_storage.py (Python)

Code before:


class SaveDataSettings(BaseSettings):
    accept_local_file_path: bool = True

    model_config = SettingsConfigDict(env_file=".env", extra="allow")


settings = SaveDataSettings()


async def save_data_item_to_storage(data_item: Union[BinaryIO, str, Any]) -> str:
    if "llama_index" in str(type(data_item)):
        # Dynamic import is used because the llama_index module is optional.
        from .transform_data import get_data_from_llama_index

        return await get_data_from_llama_index(data_item)

    try:
        from docling_core.types import DoclingDocument

        if isinstance(data_item, DoclingDocument):
            # Convert DoclingDocument to plain text and continue processing file
            data_item = data_item.export_to_text()
    except ImportError:
        logger.debug(
            "Docling import failed. Docling optional dependency not installed, so we don't expect Docling documents. "
            "Continuing normal operation."
        )

    # data is a file object coming from upload.
    if hasattr(data_item, "file"):
        return await save_data_to_file(data_item.file, filename=data_item.filename)

    if isinstance(data_item, str):
        parsed_url = urlparse(data_item)

        try:
            # In case data item is a string with a relative path transform data item to absolute path and check
            # if the file exists
            abs_path = (Path.cwd() / Path(data_item)).resolve()
            abs_path.is_file()
        except (OSError, ValueError):
            # In case file path is too long it's most likely not a relative path
            abs_path = data_item
            logger.debug(f"Data item was too long to be a possible file path: {abs_path}")
            abs_path = Path("")

        # data is s3 file path
        if parsed_url.scheme == "s3":
            return data_item

Review comment (bug):
What if the data is a docling type but Docling depdency is not installed? The users will get no error about it and the debug log level is only used by developers. Check if this can be handled how LlamaIndex documents were handled in the same function

Code after:
logger = get_logger()


class SaveDataSettings(BaseSettings):
    accept_local_file_path: bool = True

    model_config = SettingsConfigDict(env_file=".env", extra="allow")


class HTMLContent(str):
    def __new__(cls, value: str):
        if not ("<" in value and ">" in value):
            raise ValueError("Not valid HTML-like content")
        return super().__new__(cls, value)


settings = SaveDataSettings()


async def save_data_item_to_storage(data_item: Union[BinaryIO, str, Any]) -> str:
    if "llama_index" in str(type(data_item)):
        # Dynamic import is used because the llama_index module is optional.
        from .transform_data import get_data_from_llama_index

        return await get_data_from_llama_index(data_item)

    if "docling" in str(type(data_item)):
        from docling_core.types import DoclingDocument

        if isinstance(data_item, DoclingDocument):
            data_item = data_item.export_to_text()

    # data is a file object coming from upload.
    if hasattr(data_item, "file"):
        return await save_data_to_file(data_item.file, filename=data_item.filename)

    if isinstance(data_item, str):
        parsed_url = urlparse(data_item)

        try:
            # In case data item is a string with a relative path transform data item to absolute path and check
            # if the file exists
            abs_path = (Path.cwd() / Path(data_item)).resolve()
            abs_path.is_file()
        except (OSError, ValueError):
            # In case file path is too long it's most likely not a relative path
            abs_path = data_item
            logger.debug(f"Data item was too long to be a possible file path: {abs_path}")
            abs_path = Path("")

        # data is s3 file path