airbyte-cdk 0.0.0.dev0__py3-none-any.whl

airbyte_cdk/sources/file_based/file_types/unstructured_parser.py

@@ -0,0 +1,480 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+import logging
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO, IOBase
+from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple, Union
+
+import backoff
+import dpath
+import nltk
+import requests
+from unstructured.file_utils.filetype import (
+    EXT_TO_FILETYPE,
+    FILETYPE_TO_MIMETYPE,
+    STR_TO_FILETYPE,
+    FileType,
+    detect_filetype,
+)
+
+from airbyte_cdk.models import FailureType
+from airbyte_cdk.sources.file_based.config.file_based_stream_config import FileBasedStreamConfig
+from airbyte_cdk.sources.file_based.config.unstructured_format import (
+    APIParameterConfigModel,
+    APIProcessingConfigModel,
+    LocalProcessingConfigModel,
+    UnstructuredFormat,
+)
+from airbyte_cdk.sources.file_based.exceptions import FileBasedSourceError, RecordParseError
+from airbyte_cdk.sources.file_based.file_based_stream_reader import (
+    AbstractFileBasedStreamReader,
+    FileReadMode,
+)
+from airbyte_cdk.sources.file_based.file_types.file_type_parser import FileTypeParser
+from airbyte_cdk.sources.file_based.remote_file import RemoteFile
+from airbyte_cdk.sources.file_based.schema_helpers import SchemaType
+from airbyte_cdk.utils import is_cloud_environment
+from airbyte_cdk.utils.traced_exception import AirbyteTracedException
+
+unstructured_partition_pdf = None
+unstructured_partition_docx = None
+unstructured_partition_pptx = None
+
+AIRBYTE_NLTK_DATA_DIR = "/airbyte/nltk_data"
+TMP_NLTK_DATA_DIR = "/tmp/nltk_data"
+
+
+def get_nltk_temp_folder() -> str:
+    """
+    For non-root connectors /tmp is not currently writable, but we should allow it in the future.
+    It's safe to use /airbyte for now. Fallback to /tmp for local development.
+    """
+    try:
+        nltk_data_dir = AIRBYTE_NLTK_DATA_DIR
+        os.makedirs(nltk_data_dir, exist_ok=True)
+    except OSError:
+        nltk_data_dir = TMP_NLTK_DATA_DIR
+        os.makedirs(nltk_data_dir, exist_ok=True)
+    return nltk_data_dir
+
+
+try:
+    nltk_data_dir = get_nltk_temp_folder()
+    nltk.data.path.append(nltk_data_dir)
+    nltk.data.find("tokenizers/punkt.zip")
+    nltk.data.find("tokenizers/punkt_tab.zip")
+    nltk.data.find("tokenizers/averaged_perceptron_tagger_eng.zip")
+except LookupError:
+    nltk.download("punkt", download_dir=nltk_data_dir, quiet=True)
+    nltk.download("punkt_tab", download_dir=nltk_data_dir, quiet=True)
+    nltk.download("averaged_perceptron_tagger_eng", download_dir=nltk_data_dir, quiet=True)
+
+
+def optional_decode(contents: Union[str, bytes]) -> str:
+    if isinstance(contents, bytes):
+        return contents.decode("utf-8")
+    return contents
+
+
+def _import_unstructured() -> None:
+    """Dynamically imported as needed, due to slow import speed."""
+    global unstructured_partition_pdf
+    global unstructured_partition_docx
+    global unstructured_partition_pptx
+    from unstructured.partition.docx import partition_docx
+    from unstructured.partition.pdf import partition_pdf
+    from unstructured.partition.pptx import partition_pptx
+
+    # separate global variables to properly propagate typing
+    unstructured_partition_pdf = partition_pdf
+    unstructured_partition_docx = partition_docx
+    unstructured_partition_pptx = partition_pptx
+
+
+def user_error(e: Exception) -> bool:
+    """
+    Return True if this exception is caused by user error, False otherwise.
+    """
+    if not isinstance(e, RecordParseError):
+        return False
+    if not isinstance(e, requests.exceptions.RequestException):
+        return False
+    return bool(e.response and 400 <= e.response.status_code < 500)
+
+
+CLOUD_DEPLOYMENT_MODE = "cloud"
+
+
+class UnstructuredParser(FileTypeParser):
+    @property
+    def parser_max_n_files_for_schema_inference(self) -> Optional[int]:
+        """
+        Just check one file as the schema is static
+        """
+        return 1
+
+    @property
+    def parser_max_n_files_for_parsability(self) -> Optional[int]:
+        """
+        Do not check any files for parsability because it might be an expensive operation and doesn't give much confidence whether the sync will succeed.
+        """
+        return 0
+
+    def get_parser_defined_primary_key(self, config: FileBasedStreamConfig) -> Optional[str]:
+        """
+        Return the document_key field as the primary key.
+
+        his will pre-select the document key column as the primary key when setting up a connection, making it easier for the user to configure normalization in the destination.
+        """
+        return "document_key"
+
+    async def infer_schema(
+        self,
+        config: FileBasedStreamConfig,
+        file: RemoteFile,
+        stream_reader: AbstractFileBasedStreamReader,
+        logger: logging.Logger,
+    ) -> SchemaType:
+        format = _extract_format(config)
+        with stream_reader.open_file(file, self.file_read_mode, None, logger) as file_handle:
+            filetype = self._get_filetype(file_handle, file)
+            if filetype not in self._supported_file_types() and not format.skip_unprocessable_files:
+                raise self._create_parse_error(
+                    file,
+                    self._get_file_type_error_message(filetype),
+                )
+
+            return {
+                "content": {
+                    "type": "string",
+                    "description": "Content of the file as markdown. Might be null if the file could not be parsed",
+                },
+                "document_key": {
+                    "type": "string",
+                    "description": "Unique identifier of the document, e.g. the file path",
+                },
+                "_ab_source_file_parse_error": {
+                    "type": "string",
+                    "description": "Error message if the file could not be parsed even though the file is supported",
+                },
+            }
+
+    def parse_records(
+        self,
+        config: FileBasedStreamConfig,
+        file: RemoteFile,
+        stream_reader: AbstractFileBasedStreamReader,
+        logger: logging.Logger,
+        discovered_schema: Optional[Mapping[str, SchemaType]],
+    ) -> Iterable[Dict[str, Any]]:
+        format = _extract_format(config)
+        with stream_reader.open_file(file, self.file_read_mode, None, logger) as file_handle:
+            try:
+                markdown = self._read_file(file_handle, file, format, logger)
+                yield {
+                    "content": markdown,
+                    "document_key": file.uri,
+                    "_ab_source_file_parse_error": None,
+                }
+            except RecordParseError as e:
+                # RecordParseError is raised when the file can't be parsed because of a problem with the file content (either the file is not supported or the file is corrupted)
+                # if the skip_unprocessable_files flag is set, we log a warning and pass the error as part of the document
+                # otherwise, we raise the error to fail the sync
+                if format.skip_unprocessable_files:
+                    exception_str = str(e)
+                    logger.warn(f"File {file.uri} caused an error during parsing: {exception_str}.")
+                    yield {
+                        "content": None,
+                        "document_key": file.uri,
+                        "_ab_source_file_parse_error": exception_str,
+                    }
+                    logger.warn(f"File {file.uri} cannot be parsed. Skipping it.")
+                else:
+                    raise e
+            except Exception as e:
+                exception_str = str(e)
+                logger.error(f"File {file.uri} caused an error during parsing: {exception_str}.")
+                raise e
+
+    def _read_file(
+        self,
+        file_handle: IOBase,
+        remote_file: RemoteFile,
+        format: UnstructuredFormat,
+        logger: logging.Logger,
+    ) -> str:
+        _import_unstructured()
+        if (
+            (not unstructured_partition_pdf)
+            or (not unstructured_partition_docx)
+            or (not unstructured_partition_pptx)
+        ):
+            # check whether unstructured library is actually available for better error message and to ensure proper typing (can't be None after this point)
+            raise Exception("unstructured library is not available")
+
+        filetype: FileType | None = self._get_filetype(file_handle, remote_file)
+
+        if filetype is None or filetype not in self._supported_file_types():
+            raise self._create_parse_error(
+                remote_file,
+                self._get_file_type_error_message(filetype),
+            )
+        if filetype in {FileType.MD, FileType.TXT}:
+            file_content: bytes = file_handle.read()
+            decoded_content: str = optional_decode(file_content)
+            return decoded_content
+        if format.processing.mode == "local":
+            return self._read_file_locally(
+                file_handle,
+                filetype,
+                format.strategy,
+                remote_file,
+            )
+        elif format.processing.mode == "api":
+            try:
+                result: str = self._read_file_remotely_with_retries(
+                    file_handle,
+                    format.processing,
+                    filetype,
+                    format.strategy,
+                    remote_file,
+                )
+            except Exception as e:
+                # If a parser error happens during remotely processing the file, this means the file is corrupted. This case is handled by the parse_records method, so just rethrow.
+                #
+                # For other exceptions, re-throw as config error so the sync is stopped as problems with the external API need to be resolved by the user and are not considered part of the SLA.
+                # Once this parser leaves experimental stage, we should consider making this a system error instead for issues that might be transient.
+                if isinstance(e, RecordParseError):
+                    raise e
+                raise AirbyteTracedException.from_exception(
+                    e, failure_type=FailureType.config_error
+                )
+
+            return result
+
+    def _params_to_dict(
+        self, params: Optional[List[APIParameterConfigModel]], strategy: str
+    ) -> Dict[str, Union[str, List[str]]]:
+        result_dict: Dict[str, Union[str, List[str]]] = {"strategy": strategy}
+        if params is None:
+            return result_dict
+        for item in params:
+            key = item.name
+            value = item.value
+            if key in result_dict:
+                existing_value = result_dict[key]
+                # If the key already exists, append the new value to its list
+                if isinstance(existing_value, list):
+                    existing_value.append(value)
+                else:
+                    result_dict[key] = [existing_value, value]
+            else:
+                # If the key doesn't exist, add it to the dictionary
+                result_dict[key] = value
+
+        return result_dict
+
+    def check_config(self, config: FileBasedStreamConfig) -> Tuple[bool, Optional[str]]:
+        """
+        Perform a connection check for the parser config:
+        - Verify that encryption is enabled if the API is hosted on a cloud instance.
+        - Verify that the API can extract text from a file.
+
+        For local processing, we don't need to perform any additional checks, implicit pydantic validation is enough.
+        """
+        format_config = _extract_format(config)
+        if isinstance(format_config.processing, LocalProcessingConfigModel):
+            if format_config.strategy == "hi_res":
+                return False, "Hi-res strategy is not supported for local processing"
+            return True, None
+
+        if is_cloud_environment() and not format_config.processing.api_url.startswith("https://"):
+            return False, "Base URL must start with https://"
+
+        try:
+            self._read_file_remotely(
+                BytesIO(b"# Airbyte source connection test"),
+                format_config.processing,
+                FileType.MD,
+                "auto",
+                RemoteFile(uri="test", last_modified=datetime.now()),
+            )
+        except Exception:
+            return False, "".join(traceback.format_exc())
+
+        return True, None
+
+    @backoff.on_exception(
+        backoff.expo, requests.exceptions.RequestException, max_tries=5, giveup=user_error
+    )
+    def _read_file_remotely_with_retries(
+        self,
+        file_handle: IOBase,
+        format: APIProcessingConfigModel,
+        filetype: FileType,
+        strategy: str,
+        remote_file: RemoteFile,
+    ) -> str:
+        """
+        Read a file remotely, retrying up to 5 times if the error is not caused by user error. This is useful for transient network errors or the API server being overloaded temporarily.
+        """
+        return self._read_file_remotely(file_handle, format, filetype, strategy, remote_file)
+
+    def _read_file_remotely(
+        self,
+        file_handle: IOBase,
+        format: APIProcessingConfigModel,
+        filetype: FileType,
+        strategy: str,
+        remote_file: RemoteFile,
+    ) -> str:
+        headers = {"accept": "application/json", "unstructured-api-key": format.api_key}
+
+        data = self._params_to_dict(format.parameters, strategy)
+
+        file_data = {"files": ("filename", file_handle, FILETYPE_TO_MIMETYPE[filetype])}
+
+        response = requests.post(
+            f"{format.api_url}/general/v0/general", headers=headers, data=data, files=file_data
+        )
+
+        if response.status_code == 422:
+            # 422 means the file couldn't be processed, but the API is working. Treat this as a parsing error (passing an error record to the destination).
+            raise self._create_parse_error(remote_file, response.json())
+        else:
+            # Other error statuses are raised as requests exceptions (retry everything except user errors)
+            response.raise_for_status()
+
+        json_response = response.json()
+
+        return self._render_markdown(json_response)
+
+    def _read_file_locally(
+        self, file_handle: IOBase, filetype: FileType, strategy: str, remote_file: RemoteFile
+    ) -> str:
+        _import_unstructured()
+        if (
+            (not unstructured_partition_pdf)
+            or (not unstructured_partition_docx)
+            or (not unstructured_partition_pptx)
+        ):
+            # check whether unstructured library is actually available for better error message and to ensure proper typing (can't be None after this point)
+            raise Exception("unstructured library is not available")
+
+        file: Any = file_handle
+
+        # before the parsing logic is entered, the file is read completely to make sure it is in local memory
+        file_handle.seek(0)
+        file_handle.read()
+        file_handle.seek(0)
+
+        try:
+            if filetype == FileType.PDF:
+                # for PDF, read the file into a BytesIO object because some code paths in pdf parsing are doing an instance check on the file object and don't work with file-like objects
+                file_handle.seek(0)
+                with BytesIO(file_handle.read()) as file:
+                    file_handle.seek(0)
+                    elements = unstructured_partition_pdf(file=file, strategy=strategy)
+            elif filetype == FileType.DOCX:
+                elements = unstructured_partition_docx(file=file)
+            elif filetype == FileType.PPTX:
+                elements = unstructured_partition_pptx(file=file)
+        except Exception as e:
+            raise self._create_parse_error(remote_file, str(e))
+
+        return self._render_markdown([element.to_dict() for element in elements])
+
+    def _create_parse_error(
+        self,
+        remote_file: RemoteFile,
+        message: str,
+    ) -> RecordParseError:
+        return RecordParseError(
+            FileBasedSourceError.ERROR_PARSING_RECORD, filename=remote_file.uri, message=message
+        )
+
+    def _get_filetype(self, file: IOBase, remote_file: RemoteFile) -> Optional[FileType]:
+        """
+        Detect the file type based on the file name and the file content.
+
+        There are three strategies to determine the file type:
+        1. Use the mime type if available (only some sources support it)
+        2. Use the file name if available
+        3. Use the file content
+        """
+        if remote_file.mime_type and remote_file.mime_type in STR_TO_FILETYPE:
+            return STR_TO_FILETYPE[remote_file.mime_type]
+
+        # set name to none, otherwise unstructured will try to get the modified date from the local file system
+        if hasattr(file, "name"):
+            file.name = None
+
+        # detect_filetype is either using the file name or file content
+        # if possible, try to leverage the file name to detect the file type
+        # if the file name is not available, use the file content
+        file_type: FileType | None = None
+        try:
+            file_type = detect_filetype(
+                filename=remote_file.uri,
+            )
+        except Exception:
+            # Path doesn't exist locally. Try something else...
+            pass
+
+        if file_type and file_type != FileType.UNK:
+            return file_type
+
+        type_based_on_content = detect_filetype(file=file)
+        file.seek(0)  # detect_filetype is reading to read the file content, so we need to reset
+
+        if type_based_on_content and type_based_on_content != FileType.UNK:
+            return type_based_on_content
+
+        extension = "." + remote_file.uri.split(".")[-1].lower()
+        if extension in EXT_TO_FILETYPE:
+            return EXT_TO_FILETYPE[extension]
+
+        return None
+
+    def _supported_file_types(self) -> List[Any]:
+        return [FileType.MD, FileType.PDF, FileType.DOCX, FileType.PPTX, FileType.TXT]
+
+    def _get_file_type_error_message(
+        self,
+        file_type: FileType | None,
+    ) -> str:
+        supported_file_types = ", ".join([str(type) for type in self._supported_file_types()])
+        return f"File type {file_type or 'None'!s} is not supported. Supported file types are {supported_file_types}"
+
+    def _render_markdown(self, elements: List[Any]) -> str:
+        return "\n\n".join((self._convert_to_markdown(el) for el in elements))
+
+    def _convert_to_markdown(self, el: Dict[str, Any]) -> str:
+        if dpath.get(el, "type") == "Title":
+            category_depth = dpath.get(el, "metadata/category_depth", default=1) or 1
+            if not isinstance(category_depth, int):
+                category_depth = (
+                    int(category_depth) if isinstance(category_depth, (str, float)) else 1
+                )
+            heading_str = "#" * category_depth
+            return f"{heading_str} {dpath.get(el, 'text')}"
+        elif dpath.get(el, "type") == "ListItem":
+            return f"- {dpath.get(el, 'text')}"
+        elif dpath.get(el, "type") == "Formula":
+            return f"```\n{dpath.get(el, 'text')}\n```"
+        else:
+            return str(dpath.get(el, "text", default=""))
+
+    @property
+    def file_read_mode(self) -> FileReadMode:
+        return FileReadMode.READ_BINARY
+
+
+def _extract_format(config: FileBasedStreamConfig) -> UnstructuredFormat:
+    config_format = config.format
+    if not isinstance(config_format, UnstructuredFormat):
+        raise ValueError(f"Invalid format config: {config_format}")
+    return config_format

airbyte_cdk/sources/file_based/remote_file.py

@@ -0,0 +1,18 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from datetime import datetime
+from typing import Optional
+
+from pydantic.v1 import BaseModel
+
+
+class RemoteFile(BaseModel):
+    """
+    A file in a file-based stream.
+    """
+
+    uri: str
+    last_modified: datetime
+    mime_type: Optional[str] = None