unstructured-ingest 0.6.2__py3-none-any.whl → 0.7.0__py3-none-any.whl

unstructured_ingest/connector/discord.py

@@ -1,180 +0,0 @@
-import datetime as dt
-import typing as t
-from dataclasses import dataclass
-from pathlib import Path
-
-from unstructured_ingest.enhanced_dataclass import enhanced_field
-from unstructured_ingest.error import SourceConnectionError, SourceConnectionNetworkError
-from unstructured_ingest.interfaces import (
-    AccessConfig,
-    BaseConnectorConfig,
-    BaseSingleIngestDoc,
-    BaseSourceConnector,
-    IngestDocCleanupMixin,
-    SourceConnectorCleanupMixin,
-    SourceMetadata,
-)
-from unstructured_ingest.logger import logger
-from unstructured_ingest.utils.dep_check import (
-    requires_dependencies,
-)
-
-
-@dataclass
-class DiscordAccessConfig(AccessConfig):
-    token: str = enhanced_field(sensitive=True)
-
-
-@dataclass
-class SimpleDiscordConfig(BaseConnectorConfig):
-    """Connector config where channels is a comma separated list of
-    Discord channels to pull messages from.
-    """
-
-    # Discord Specific Options
-    access_config: DiscordAccessConfig
-    channels: t.List[str]
-    period: t.Optional[int] = None
-
-
-@dataclass
-class DiscordIngestDoc(IngestDocCleanupMixin, BaseSingleIngestDoc):
-    """Class encapsulating fetching a doc and writing processed results (but not
-    doing the processing!).
-    Also includes a cleanup method. When things go wrong and the cleanup
-    method is not called, the file is left behind on the filesystem to assist debugging.
-    """
-
-    connector_config: SimpleDiscordConfig
-    channel: str
-    days: t.Optional[int] = None
-    registry_name: str = "discord"
-
-    # NOTE(crag): probably doesn't matter,  but intentionally not defining tmp_download_file
-    # __post_init__ for multiprocessing simplicity (no Path objects in initially
-    # instantiated object)
-    def _tmp_download_file(self):
-        channel_file = self.channel + ".txt"
-        return Path(self.read_config.download_dir) / channel_file
-
-    @property
-    def _output_filename(self):
-        output_file = self.channel + ".json"
-        return Path(self.processor_config.output_dir) / output_file
-
-    def _create_full_tmp_dir_path(self):
-        self._tmp_download_file().parent.mkdir(parents=True, exist_ok=True)
-
-    @SourceConnectionNetworkError.wrap
-    @requires_dependencies(dependencies=["discord"], extras="discord")
-    def _get_messages(self):
-        """Actually fetches the data from discord."""
-        import discord
-        from discord.ext import commands
-
-        messages: t.List[discord.Message] = []
-        jumpurl: t.List[str] = []
-        intents = discord.Intents.default()
-        intents.message_content = True
-        bot = commands.Bot(command_prefix=">", intents=intents)
-
-        @bot.event
-        async def on_ready():
-            try:
-                after_date = None
-                if self.days:
-                    after_date = dt.datetime.utcnow() - dt.timedelta(days=self.days)
-                channel = bot.get_channel(int(self.channel))
-                jumpurl.append(channel.jump_url)  # type: ignore
-                async for msg in channel.history(after=after_date):  # type: ignore
-                    messages.append(msg)
-                await bot.close()
-            except Exception:
-                logger.error("Error fetching messages")
-                await bot.close()
-                raise
-
-        bot.run(self.connector_config.access_config.token)
-        jump_url = None if len(jumpurl) < 1 else jumpurl[0]
-        return messages, jump_url
-
-    def update_source_metadata(self, **kwargs):
-        messages, jump_url = kwargs.get("messages_tuple", self._get_messages())
-        if messages == []:
-            self.source_metadata = SourceMetadata(
-                exists=False,
-            )
-            return
-        dates = [m.created_at for m in messages if m.created_at]
-        dates.sort()
-        self.source_metadata = SourceMetadata(
-            date_created=dates[0].isoformat(),
-            date_modified=dates[-1].isoformat(),
-            source_url=jump_url,
-            exists=True,
-        )
-
-    @SourceConnectionError.wrap
-    @BaseSingleIngestDoc.skip_if_file_exists
-    def get_file(self):
-        self._create_full_tmp_dir_path()
-
-        messages, jump_url = self._get_messages()
-        self.update_source_metadata(messages_tuple=(messages, jump_url))
-        if messages == []:
-            raise ValueError(f"Failed to retrieve messages from Discord channel {self.channel}")
-        self._tmp_download_file().parent.mkdir(parents=True, exist_ok=True)
-        with open(self._tmp_download_file(), "w") as f:
-            for m in messages:
-                f.write(m.content + "\n")
-
-    @property
-    def filename(self):
-        """The filename of the file created from a discord channel"""
-        return self._tmp_download_file()
-
-    @property
-    def version(self) -> t.Optional[str]:
-        return None
-
-    @property
-    def record_locator(self) -> t.Optional[t.Dict[str, t.Any]]:
-        return {
-            "channel": self.channel,
-        }
-
-
-class DiscordSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
-    """Objects of this class support fetching document(s) from"""
-
-    connector_config: SimpleDiscordConfig
-
-    def initialize(self):
-        pass
-
-    @requires_dependencies(dependencies=["discord"], extras="discord")
-    def check_connection(self):
-        import asyncio
-
-        import discord
-        from discord.client import Client
-
-        intents = discord.Intents.default()
-        try:
-            client = Client(intents=intents)
-            asyncio.run(client.start(token=self.connector_config.access_config.token))
-        except Exception as e:
-            logger.error(f"failed to validate connection: {e}", exc_info=True)
-            raise SourceConnectionError(f"failed to validate connection: {e}")
-
-    def get_ingest_docs(self):
-        return [
-            DiscordIngestDoc(
-                connector_config=self.connector_config,
-                processor_config=self.processor_config,
-                read_config=self.read_config,
-                channel=channel,
-                days=self.connector_config.period,
-            )
-            for channel in self.connector_config.channels
-        ]

unstructured_ingest/connector/elasticsearch.py

@@ -1,396 +0,0 @@
-import copy
-import hashlib
-import typing as t
-import uuid
-from dataclasses import dataclass, field
-from pathlib import Path
-
-from dataclasses_json.core import Json
-
-from unstructured_ingest.enhanced_dataclass import enhanced_field
-from unstructured_ingest.enhanced_dataclass.core import _asdict
-from unstructured_ingest.error import DestinationConnectionError, SourceConnectionError
-from unstructured_ingest.interfaces import (
-    AccessConfig,
-    BaseConnectorConfig,
-    BaseDestinationConnector,
-    BaseIngestDocBatch,
-    BaseSingleIngestDoc,
-    BaseSourceConnector,
-    IngestDocCleanupMixin,
-    SourceConnectorCleanupMixin,
-    SourceMetadata,
-    WriteConfig,
-)
-from unstructured_ingest.logger import logger
-from unstructured_ingest.utils.data_prep import flatten_dict, generator_batching_wbytes
-from unstructured_ingest.utils.dep_check import requires_dependencies
-
-if t.TYPE_CHECKING:
-    from elasticsearch import Elasticsearch
-
-
-@dataclass
-class ElasticsearchAccessConfig(AccessConfig):
-    hosts: t.Optional[t.List[str]] = None
-    username: t.Optional[str] = None
-    password: t.Optional[str] = enhanced_field(default=None, sensitive=True)
-    cloud_id: t.Optional[str] = None
-    api_key: t.Optional[str] = enhanced_field(
-        default=None, sensitive=True, overload_name="es_api_key"
-    )
-    api_key_id: t.Optional[str] = None
-    bearer_auth: t.Optional[str] = enhanced_field(default=None, sensitive=True)
-    ca_certs: t.Optional[str] = None
-    ssl_assert_fingerprint: t.Optional[str] = enhanced_field(default=None, sensitive=True)
-
-    def to_dict(self, **kwargs) -> t.Dict[str, Json]:
-        d = super().to_dict(**kwargs)
-        # Update auth related fields to conform to what the SDK expects based on the
-        # supported methods:
-        # https://www.elastic.co/guide/en/elasticsearch/client/python-api/current/connecting.html
-        if not self.ca_certs:
-            # ES library already sets a default for this, don't want to
-            # introduce data by setting it to None
-            d.pop("ca_certs")
-        if self.password and (self.cloud_id or self.ca_certs or self.ssl_assert_fingerprint):
-            d.pop("password")
-            d["basic_auth"] = ("elastic", self.password)
-        elif not self.cloud_id and self.username and self.password:
-            d.pop("username", None)
-            d.pop("password", None)
-            d["basic_auth"] = (self.username, self.password)
-        elif self.api_key and self.api_key_id:
-            d.pop("api_key_id", None)
-            d.pop("api_key", None)
-            d["api_key"] = (self.api_key_id, self.api_key)
-        # This doesn't exist on the client init, remove:
-        d.pop("api_key_id", None)
-        return d
-
-
-@dataclass
-class SimpleElasticsearchConfig(BaseConnectorConfig):
-    """Connector config where:
-    url is the url to access the elasticsearch server,
-    index_name is the name of the index to reach to,
-    """
-
-    index_name: str
-    batch_size: int = 100
-    fields: t.List[str] = field(default_factory=list)
-    access_config: ElasticsearchAccessConfig = None
-
-
-@dataclass
-class ElasticsearchDocumentMeta:
-    """Metadata specifying:
-    name of the elasticsearch index that is being reached to,
-    and the id of document that is being reached to,
-    """
-
-    index_name: str
-    document_id: str
-
-
-@dataclass
-class ElasticsearchIngestDoc(IngestDocCleanupMixin, BaseSingleIngestDoc):
-    """Class encapsulating fetching a doc and writing processed results (but not
-    doing the processing!).
-
-    Current implementation creates a python Elasticsearch client to fetch each doc,
-    rather than creating a client for each thread.
-    """
-
-    connector_config: SimpleElasticsearchConfig
-    document_meta: ElasticsearchDocumentMeta
-    document: dict = field(default_factory=dict)
-    registry_name: str = "elasticsearch"
-
-    # TODO: remove one of filename or _tmp_download_file, using a wrapper
-    @property
-    def filename(self):
-        f = self.document_meta.document_id
-        if self.connector_config.fields:
-            f = "{}-{}".format(
-                f,
-                hashlib.sha256(",".join(self.connector_config.fields).encode()).hexdigest()[:8],
-            )
-        return (
-            Path(self.read_config.download_dir) / self.document_meta.index_name / f"{f}.txt"
-        ).resolve()
-
-    @property
-    def _output_filename(self):
-        """Create filename document id combined with a hash of the query to uniquely identify
-        the output file."""
-        # Generate SHA256 hash and take the first 8 characters
-        filename = self.document_meta.document_id
-        if self.connector_config.fields:
-            filename = "{}-{}".format(
-                filename,
-                hashlib.sha256(",".join(self.connector_config.fields).encode()).hexdigest()[:8],
-            )
-        output_file = f"{filename}.json"
-        return (
-            Path(self.processor_config.output_dir) / self.connector_config.index_name / output_file
-        )
-
-    def update_source_metadata(self, **kwargs):
-        if self.document is None:
-            self.source_metadata = SourceMetadata(
-                exists=False,
-            )
-            return
-        self.source_metadata = SourceMetadata(
-            version=self.document["_version"],
-            exists=True,
-        )
-
-    @SourceConnectionError.wrap
-    @requires_dependencies(["elasticsearch"], extras="elasticsearch")
-    @BaseSingleIngestDoc.skip_if_file_exists
-    def get_file(self):
-        pass
-
-    @property
-    def date_created(self) -> t.Optional[str]:
-        return None
-
-    @property
-    def date_modified(self) -> t.Optional[str]:
-        return None
-
-    @property
-    def source_url(self) -> t.Optional[str]:
-        return None
-
-    @property
-    def record_locator(self) -> t.Optional[t.Dict[str, t.Any]]:
-        return {
-            "hosts": self.connector_config.access_config.hosts,
-            "index_name": self.connector_config.index_name,
-            "document_id": self.document_meta.document_id,
-        }
-
-
-@dataclass
-class ElasticsearchIngestDocBatch(BaseIngestDocBatch):
-    connector_config: SimpleElasticsearchConfig
-    ingest_docs: t.List[ElasticsearchIngestDoc] = field(default_factory=list)
-    list_of_ids: t.List[str] = field(default_factory=list)
-    registry_name: str = "elasticsearch_batch"
-
-    def __post_init__(self):
-        # Until python3.8 is deprecated, this is a limitation of dataclass inheritance
-        # to make it a required field
-        if len(self.list_of_ids) == 0:
-            raise ValueError("list_of_ids is required")
-
-    @property
-    def unique_id(self) -> str:
-        return ",".join(sorted(self.list_of_ids))
-
-    @requires_dependencies(["elasticsearch"], extras="elasticsearch")
-    def _get_docs(self):
-        from elasticsearch import Elasticsearch
-        from elasticsearch.helpers import scan
-
-        es = Elasticsearch(**self.connector_config.access_config.to_dict(apply_name_overload=False))
-        scan_query = {
-            "_source": self.connector_config.fields,
-            "version": True,
-            "query": {"ids": {"values": self.list_of_ids}},
-        }
-
-        result = scan(
-            es,
-            query=scan_query,
-            scroll="1m",
-            index=self.connector_config.index_name,
-        )
-        return list(result)
-
-    @SourceConnectionError.wrap
-    @requires_dependencies(["elasticsearch"], extras="elasticsearch")
-    def get_files(self):
-        documents = self._get_docs()
-        for doc in documents:
-            ingest_doc = ElasticsearchIngestDoc(
-                processor_config=self.processor_config,
-                read_config=self.read_config,
-                connector_config=self.connector_config,
-                document=doc,
-                document_meta=ElasticsearchDocumentMeta(
-                    self.connector_config.index_name, doc["_id"]
-                ),
-            )
-            ingest_doc.update_source_metadata()
-            doc_body = doc["_source"]
-            filename = ingest_doc.filename
-            flattened_dict = flatten_dict(dictionary=doc_body)
-            str_values = [str(value) for value in flattened_dict.values()]
-            concatenated_values = "\n".join(str_values)
-
-            filename.parent.mkdir(parents=True, exist_ok=True)
-            with open(filename, "w", encoding="utf8") as f:
-                f.write(concatenated_values)
-            self.ingest_docs.append(ingest_doc)
-
-
-@dataclass
-class ElasticsearchSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
-    """Fetches particular fields from all documents in a given elasticsearch cluster and index"""
-
-    connector_config: SimpleElasticsearchConfig
-    _es: t.Optional["Elasticsearch"] = field(init=False, default=None)
-
-    @property
-    def es(self):
-        from elasticsearch import Elasticsearch
-
-        if self._es is None:
-            self._es = Elasticsearch(
-                **self.connector_config.access_config.to_dict(apply_name_overload=False)
-            )
-        return self._es
-
-    def check_connection(self):
-        try:
-            self.es.perform_request("HEAD", "/", headers={"accept": "application/json"})
-        except Exception as e:
-            logger.error(f"failed to validate connection: {e}", exc_info=True)
-            raise SourceConnectionError(f"failed to validate connection: {e}")
-
-    def __post_init__(self):
-        self.scan_query: dict = {"stored_fields": [], "query": {"match_all": {}}}
-
-    def initialize(self):
-        pass
-
-    @requires_dependencies(["elasticsearch"], extras="elasticsearch")
-    def _get_doc_ids(self):
-        """Fetches all document ids in an index"""
-        from elasticsearch.helpers import scan
-
-        hits = scan(
-            self.es,
-            query=self.scan_query,
-            scroll="1m",
-            index=self.connector_config.index_name,
-        )
-
-        return [hit["_id"] for hit in hits]
-
-    def get_ingest_docs(self):
-        """Fetches all documents in an index, using ids that are fetched with _get_doc_ids"""
-        ids = self._get_doc_ids()
-        id_batches = [
-            ids[
-                i
-                * self.connector_config.batch_size : (i + 1)  # noqa
-                * self.connector_config.batch_size
-            ]
-            for i in range(
-                (len(ids) + self.connector_config.batch_size - 1)
-                // self.connector_config.batch_size
-            )
-        ]
-        return [
-            ElasticsearchIngestDocBatch(
-                connector_config=self.connector_config,
-                processor_config=self.processor_config,
-                read_config=self.read_config,
-                list_of_ids=batched_ids,
-            )
-            for batched_ids in id_batches
-        ]
-
-
-@dataclass
-class ElasticsearchWriteConfig(WriteConfig):
-    batch_size_bytes: int = 15_000_000
-    num_processes: int = 1
-
-
-@dataclass
-class ElasticsearchDestinationConnector(BaseDestinationConnector):
-    write_config: ElasticsearchWriteConfig
-    connector_config: SimpleElasticsearchConfig
-    _client: t.Optional["Elasticsearch"] = field(init=False, default=None)
-
-    def to_dict(self, **kwargs):
-        """
-        The _client variable in this dataclass breaks deepcopy due to:
-        TypeError: cannot pickle '_thread.lock' object
-        When serializing, remove it, meaning client data will need to be reinitialized
-        when deserialized
-        """
-        self_cp = copy.copy(self)
-        if hasattr(self_cp, "_client"):
-            setattr(self_cp, "_client", None)
-        return _asdict(self_cp, **kwargs)
-
-    @DestinationConnectionError.wrap
-    @requires_dependencies(["elasticsearch"], extras="elasticsearch")
-    def generate_client(self) -> "Elasticsearch":
-        from elasticsearch import Elasticsearch
-
-        return Elasticsearch(
-            **self.connector_config.access_config.to_dict(apply_name_overload=False)
-        )
-
-    @property
-    def client(self):
-        if self._client is None:
-            self._client = self.generate_client()
-        return self._client
-
-    def initialize(self):
-        _ = self.client
-
-    @DestinationConnectionError.wrap
-    def check_connection(self):
-        try:
-            assert self.client.ping()
-        except Exception as e:
-            logger.error(f"failed to validate connection: {e}", exc_info=True)
-            raise DestinationConnectionError(f"failed to validate connection: {e}")
-
-    @requires_dependencies(["elasticsearch"], extras="elasticsearch")
-    def write_dict(self, *args, elements_dict: t.List[t.Dict[str, t.Any]], **kwargs) -> None:
-        logger.info(
-            f"writing document batches to destination"
-            f" index named {self.connector_config.index_name}"
-            f" at {self.connector_config.access_config.hosts}"
-            f" with batch size (in bytes) {self.write_config.batch_size_bytes}"
-            f" with {self.write_config.num_processes} (number of) processes"
-        )
-        from elasticsearch.helpers import parallel_bulk
-
-        for batch in generator_batching_wbytes(
-            elements_dict, batch_size_limit_bytes=self.write_config.batch_size_bytes
-        ):
-            for success, info in parallel_bulk(
-                self.client, batch, thread_count=self.write_config.num_processes
-            ):
-                if not success:
-                    logger.error(
-                        "upload failed for a batch in elasticsearch destination connector:", info
-                    )
-
-    def normalize_dict(self, element_dict: dict) -> dict:
-        return {
-            "_index": self.connector_config.index_name,
-            "_id": str(uuid.uuid4()),
-            "_source": {
-                "element_id": element_dict.pop("element_id", None),
-                "embeddings": element_dict.pop("embeddings", None),
-                "text": element_dict.pop("text", None),
-                "type": element_dict.pop("type", None),
-                "metadata": flatten_dict(
-                    element_dict.pop("metadata", None),
-                    separator="-",
-                ),
-            },
-        }

unstructured_ingest/connector/fsspec/azure.py

@@ -1,78 +0,0 @@
-import typing as t
-from dataclasses import dataclass
-
-from unstructured_ingest.connector.fsspec.fsspec import (
-    FsspecDestinationConnector,
-    FsspecIngestDoc,
-    FsspecSourceConnector,
-    FsspecWriteConfig,
-    SimpleFsspecConfig,
-    WriteTextConfig,
-)
-from unstructured_ingest.enhanced_dataclass import enhanced_field
-from unstructured_ingest.error import DestinationConnectionError, SourceConnectionError
-from unstructured_ingest.interfaces import AccessConfig
-from unstructured_ingest.logger import logger
-from unstructured_ingest.utils.dep_check import requires_dependencies
-
-
-@dataclass
-class AzureWriteTextConfig(WriteTextConfig):
-    overwrite: bool = False
-
-
-@dataclass
-class AzureWriteConfig(FsspecWriteConfig):
-    write_text_config: t.Optional[AzureWriteTextConfig] = None
-
-
-@dataclass
-class AzureAccessConfig(AccessConfig):
-    account_name: t.Optional[str] = enhanced_field(default=None, sensitive=True)
-    account_key: t.Optional[str] = enhanced_field(default=None, sensitive=True)
-    connection_string: t.Optional[str] = enhanced_field(default=None, sensitive=True)
-    sas_token: t.Optional[str] = enhanced_field(default=None, sensitive=True)
-
-
-@dataclass
-class SimpleAzureBlobStorageConfig(SimpleFsspecConfig):
-    access_config: AzureAccessConfig = None
-
-
-@dataclass
-class AzureBlobStorageIngestDoc(FsspecIngestDoc):
-    connector_config: SimpleAzureBlobStorageConfig
-    registry_name: str = "azure"
-
-    @SourceConnectionError.wrap
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    def get_file(self):
-        super().get_file()
-
-
-@dataclass
-class AzureBlobStorageSourceConnector(FsspecSourceConnector):
-    connector_config: SimpleAzureBlobStorageConfig
-
-    def __post_init__(self):
-        self.ingest_doc_cls: t.Type[AzureBlobStorageIngestDoc] = AzureBlobStorageIngestDoc
-
-
-@dataclass
-class AzureBlobStorageDestinationConnector(FsspecDestinationConnector):
-    connector_config: SimpleAzureBlobStorageConfig
-    write_config: AzureWriteConfig
-
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    def initialize(self):
-        super().initialize()
-
-    @requires_dependencies(["adlfs"], extras="azure")
-    def check_connection(self):
-        from adlfs import AzureBlobFileSystem
-
-        try:
-            AzureBlobFileSystem(**self.connector_config.get_access_config())
-        except ValueError as connection_error:
-            logger.error(f"failed to validate connection: {connection_error}", exc_info=True)
-            raise DestinationConnectionError(f"failed to validate connection: {connection_error}")