unstructured-ingest 0.3.13__py3-none-any.whl

unstructured_ingest/connector/opensearch.py

@@ -0,0 +1,218 @@
+import typing as t
+from dataclasses import dataclass, field
+
+from dataclasses_json.core import Json
+
+from unstructured_ingest.connector.elasticsearch import (
+    ElasticsearchDestinationConnector,
+    ElasticsearchDocumentMeta,
+    ElasticsearchIngestDoc,
+    ElasticsearchIngestDocBatch,
+    ElasticsearchSourceConnector,
+    SimpleElasticsearchConfig,
+)
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.error import DestinationConnectionError, SourceConnectionError
+from unstructured_ingest.interfaces import AccessConfig, BaseSingleIngestDoc
+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 opensearchpy import OpenSearch
+
+"""Since the actual OpenSearch project is a fork of Elasticsearch, we are relying
+heavily on the Elasticsearch connector code, inheriting the functionality as much as possible."""
+
+
+@dataclass
+class OpenSearchAccessConfig(AccessConfig):
+    hosts: t.Optional[t.List[str]] = None
+    username: t.Optional[str] = None
+    password: t.Optional[str] = enhanced_field(default=None, sensitive=True)
+    use_ssl: bool = False
+    verify_certs: bool = False
+    ssl_show_warn: bool = False
+    ca_certs: t.Optional[str] = None
+    client_cert: t.Optional[str] = None
+    client_key: t.Optional[str] = None
+
+    def to_dict(self, **kwargs) -> t.Dict[str, Json]:
+        d = super().to_dict(**kwargs)
+        d["http_auth"] = (self.username, self.password)
+        return d
+
+
+@dataclass
+class SimpleOpenSearchConfig(SimpleElasticsearchConfig):
+    access_config: OpenSearchAccessConfig = None
+
+
+@dataclass
+class OpenSearchIngestDoc(ElasticsearchIngestDoc):
+    """Class encapsulating fetching a doc and writing processed results (but not
+    doing the processing!).
+
+    Current implementation creates a python OpenSearch client to fetch each doc,
+    rather than creating a client for each thread.
+    """
+
+    connector_config: SimpleOpenSearchConfig
+    registry_name: str = "opensearch"
+
+    @SourceConnectionError.wrap
+    @requires_dependencies(["opensearchpy"], extras="opensearch")
+    @BaseSingleIngestDoc.skip_if_file_exists
+    def get_file(self):
+        pass
+
+
+@dataclass
+class OpenSearchIngestDocBatch(ElasticsearchIngestDocBatch):
+    connector_config: SimpleOpenSearchConfig
+    ingest_docs: t.List[OpenSearchIngestDoc] = field(default_factory=list)
+    registry_name: str = "opensearch_batch"
+
+    @requires_dependencies(["opensearchpy"], extras="opensearch")
+    def _get_docs(self):
+        from opensearchpy import OpenSearch
+        from opensearchpy.helpers import scan
+
+        ops = OpenSearch(**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(
+            ops,
+            query=scan_query,
+            scroll="1m",
+            index=self.connector_config.index_name,
+        )
+        return list(result)
+
+    @SourceConnectionError.wrap
+    @requires_dependencies(["opensearchpy"], extras="opensearch")
+    def get_files(self):
+        documents = self._get_docs()
+        for doc in documents:
+            ingest_doc = OpenSearchIngestDoc(
+                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 OpenSearchSourceConnector(ElasticsearchSourceConnector):
+    """Fetches particular fields from all documents in a given opensearch cluster and index"""
+
+    connector_config: SimpleOpenSearchConfig
+    _ops: t.Optional["OpenSearch"] = field(init=False, default=None)
+
+    @property
+    def ops(self):
+        from opensearchpy import OpenSearch
+
+        if self._ops is None:
+            self._ops = OpenSearch(
+                **self.connector_config.access_config.to_dict(apply_name_overload=False)
+            )
+        return self._ops
+
+    def check_connection(self):
+        try:
+            assert self.ops.ping()
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+    @requires_dependencies(["opensearchpy"], extras="opensearch")
+    def _get_doc_ids(self):
+        """Fetches all document ids in an index"""
+        from opensearchpy.helpers import scan
+
+        hits = scan(
+            self.ops,
+            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 [
+            OpenSearchIngestDocBatch(
+                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 OpenSearchDestinationConnector(ElasticsearchDestinationConnector):
+    connector_config: SimpleOpenSearchConfig
+    _client: t.Optional["OpenSearch"] = field(init=False, default=None)
+
+    @DestinationConnectionError.wrap
+    @requires_dependencies(["opensearchpy"], extras="opensearch")
+    def generate_client(self) -> "OpenSearch":
+        from opensearchpy import OpenSearch
+
+        return OpenSearch(**self.connector_config.access_config.to_dict(apply_name_overload=False))
+
+    @requires_dependencies(["opensearchpy"], extras="opensearch")
+    def write_dict(self, *args, elements_dict: t.List[t.Dict[str, t.Any]]) -> 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 opensearchpy.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 opensearch destination connector:", info
+                    )

unstructured_ingest/connector/outlook.py

@@ -0,0 +1,285 @@
+import hashlib
+import os
+import typing as t
+from collections import defaultdict
+from dataclasses import dataclass, field
+from itertools import chain
+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
+
+MAX_NUM_EMAILS = 1000000  # Maximum number of emails per folder
+if t.TYPE_CHECKING:
+    from office365.graph_client import GraphClient
+
+
+class MissingFolderError(Exception):
+    """There are no root folders with those names."""
+
+
+@dataclass
+class OutlookAccessConfig(AccessConfig):
+    client_credential: str = enhanced_field(repr=False, sensitive=True, overload_name="client_cred")
+
+
+@dataclass
+class SimpleOutlookConfig(BaseConnectorConfig):
+    """This class is getting the token."""
+
+    access_config: OutlookAccessConfig
+    user_email: str
+    client_id: str
+    tenant: t.Optional[str] = field(repr=False, default="common")
+    authority_url: t.Optional[str] = field(repr=False, default="https://login.microsoftonline.com")
+    outlook_folders: t.List[str] = field(default_factory=list)
+    recursive: bool = False
+    registry_name: str = "outlook"
+
+    def __post_init__(self):
+        if not (self.client_id and self.access_config.client_credential and self.user_email):
+            raise ValueError(
+                "Please provide one of the following mandatory values:"
+                "\nclient_id\nclient_cred\nuser_email",
+            )
+        self.token_factory = self._acquire_token
+
+    @requires_dependencies(["msal"])
+    def _acquire_token(self):
+        from msal import ConfidentialClientApplication
+
+        try:
+            app = ConfidentialClientApplication(
+                authority=f"{self.authority_url}/{self.tenant}",
+                client_id=self.client_id,
+                client_credential=self.access_config.client_credential,
+            )
+            token = app.acquire_token_for_client(
+                scopes=["https://graph.microsoft.com/.default"],
+            )
+        except ValueError as exc:
+            logger.error("Couldn't set up credentials for Outlook")
+            raise exc
+        return token
+
+    @requires_dependencies(["office365"], extras="outlook")
+    def _get_client(self):
+        from office365.graph_client import GraphClient
+
+        return GraphClient(self.token_factory)
+
+
+@dataclass
+class OutlookIngestDoc(IngestDocCleanupMixin, BaseSingleIngestDoc):
+    connector_config: SimpleOutlookConfig
+    message_id: str
+    registry_name: str = "outlook"
+
+    def __post_init__(self):
+        self._set_download_paths()
+
+    def hash_mail_name(self, id):
+        """Outlook email ids are 152 char long. Hash to shorten to 16."""
+        return hashlib.sha256(id.encode("utf-8")).hexdigest()[:16]
+
+    def _set_download_paths(self) -> None:
+        """Creates paths for downloading and parsing."""
+        download_path = Path(f"{self.read_config.download_dir}")
+        output_path = Path(f"{self.processor_config.output_dir}")
+
+        self.download_dir = download_path
+        self.download_filepath = (
+            download_path / f"{self.hash_mail_name(self.message_id)}.eml"
+        ).resolve()
+        oname = f"{self.hash_mail_name(self.message_id)}.eml.json"
+        self.output_dir = output_path
+        self.output_filepath = (output_path / oname).resolve()
+
+    @property
+    def filename(self):
+        return Path(self.download_filepath).resolve()
+
+    @property
+    def _output_filename(self):
+        return Path(self.output_filepath).resolve()
+
+    @property
+    def record_locator(self) -> t.Optional[t.Dict[str, t.Any]]:
+        return {
+            "message_id": self.message_id,
+            "user_email": self.connector_config.user_email,
+        }
+
+    @requires_dependencies(["office365"], extras="outlook")
+    def update_source_metadata(self, **kwargs):
+        from office365.runtime.client_request_exception import ClientRequestException
+
+        try:
+            client = self.connector_config._get_client()
+            msg = (
+                client.users[self.connector_config.user_email]
+                .messages[self.message_id]
+                .get()
+                .execute_query()
+            )
+        except ClientRequestException as e:
+            if e.response.status_code == 404:
+                self.source_metadata = SourceMetadata(
+                    exists=False,
+                )
+                return
+            raise
+        self.source_metadata = SourceMetadata(
+            date_created=msg.created_datetime.isoformat(),
+            date_modified=msg.last_modified_datetime.isoformat(),
+            version=msg.get_property("changeKey"),
+            source_url=msg.get_property("webLink"),
+            exists=True,
+        )
+
+    @SourceConnectionNetworkError.wrap
+    def _run_download(self, local_file):
+        client = self.connector_config._get_client()
+        client.users[self.connector_config.user_email].messages[self.message_id].download(
+            local_file,
+        ).execute_query()
+
+    @SourceConnectionError.wrap
+    @BaseSingleIngestDoc.skip_if_file_exists
+    @requires_dependencies(["office365"], extras="outlook")
+    def get_file(self):
+        """Relies on Office365 python sdk message object to do the download."""
+        try:
+            self.connector_config._get_client()
+            self.update_source_metadata()
+            if not self.download_dir.is_dir():
+                logger.debug(f"creating directory: {self.download_dir}")
+                self.download_dir.mkdir(parents=True, exist_ok=True)
+
+            with open(
+                os.path.join(
+                    self.download_dir,
+                    self.hash_mail_name(self.message_id) + ".eml",
+                ),
+                "wb",
+            ) as local_file:
+                self._run_download(local_file=local_file)
+
+        except Exception as e:
+            logger.error(
+                f"Error while downloading and saving file: {self.hash_mail_name(self.message_id)}.",
+            )
+            logger.error(e)
+            return
+        logger.info(f"file downloaded: {self.hash_mail_name(self.message_id)}")
+        return
+
+
+@dataclass
+class OutlookSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
+    connector_config: SimpleOutlookConfig
+    _client: t.Optional["GraphClient"] = field(init=False, default=None)
+
+    @property
+    def client(self) -> "GraphClient":
+        if self._client is None:
+            self._client = self.connector_config._get_client()
+        return self._client
+
+    def initialize(self):
+        try:
+            self.get_folder_ids()
+        except Exception as e:
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+    def check_connection(self):
+        try:
+            _ = self.client
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+    def recurse_folders(self, folder_id, main_folder_dict):
+        """We only get a count of subfolders for any folder.
+        Have to make additional calls to get subfolder ids."""
+        subfolders = (
+            self.client.users[self.connector_config.user_email]
+            .mail_folders[folder_id]
+            .child_folders.get()
+            .execute_query()
+        )
+        for subfolder in subfolders:
+            for k, v in main_folder_dict.items():
+                if subfolder.get_property("parentFolderId") in v:
+                    v.append(subfolder.id)
+            if subfolder.get_property("childFolderCount") > 0:
+                self.recurse_folders(subfolder.id, main_folder_dict)
+
+    def get_folder_ids(self):
+        """Sets the mail folder ids and subfolder ids for requested root mail folders."""
+        self.root_folders = defaultdict(list)
+        root_folders_with_subfolders = []
+        get_root_folders = (
+            self.client.users[self.connector_config.user_email].mail_folders.get().execute_query()
+        )
+
+        for folder in get_root_folders:
+            self.root_folders[folder.display_name].append(folder.id)
+            if folder.get_property("childFolderCount") > 0:
+                root_folders_with_subfolders.append(folder.id)
+
+        for folder in root_folders_with_subfolders:
+            self.recurse_folders(folder, self.root_folders)
+
+        # Narrow down all mail folder ids (plus all subfolders) to the ones that were requested.
+        self.selected_folder_ids = list(
+            chain.from_iterable(
+                [
+                    v
+                    for k, v in self.root_folders.items()
+                    if k.lower() in [x.lower() for x in self.connector_config.outlook_folders]
+                ],
+            ),
+        )
+        if not self.selected_folder_ids:
+            raise MissingFolderError(
+                "There are no root folders with the names: "
+                f"{self.connector_config.outlook_folders}",
+            )
+
+    def get_ingest_docs(self):
+        """Returns a list of all the message objects that are in the requested root folder(s)."""
+        filtered_messages = []
+
+        # Get all the relevant messages in the selected folders/subfolders.
+        for folder_id in self.selected_folder_ids:
+            messages = (
+                self.client.users[self.connector_config.user_email]
+                .mail_folders[folder_id]
+                .messages.get()
+                .top(MAX_NUM_EMAILS)  # Prevents the return from paging
+                .execute_query()
+            )
+            # Skip empty list if there are no messages in folder.
+            if messages:
+                filtered_messages.append(messages)
+        return [
+            OutlookIngestDoc(
+                connector_config=self.connector_config,
+                processor_config=self.processor_config,
+                read_config=self.read_config,
+                message_id=message.id,
+            )
+            for message in list(chain.from_iterable(filtered_messages))
+        ]

unstructured_ingest/connector/pinecone.py

@@ -0,0 +1,140 @@
+import copy
+import json
+import multiprocessing as mp
+import typing as t
+import uuid
+from dataclasses import dataclass
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.enhanced_dataclass.core import _asdict
+from unstructured_ingest.error import DestinationConnectionError, WriteError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    BaseConnectorConfig,
+    BaseDestinationConnector,
+    ConfigSessionHandleMixin,
+    IngestDocSessionHandleMixin,
+    WriteConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.data_prep import batch_generator, flatten_dict
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if t.TYPE_CHECKING:
+    from pinecone import Index as PineconeIndex
+
+
+@dataclass
+class PineconeAccessConfig(AccessConfig):
+    api_key: str = enhanced_field(sensitive=True)
+
+
+@dataclass
+class SimplePineconeConfig(ConfigSessionHandleMixin, BaseConnectorConfig):
+    index_name: str
+    environment: str
+    access_config: PineconeAccessConfig
+
+
+@dataclass
+class PineconeWriteConfig(WriteConfig):
+    batch_size: int = 50
+    num_processes: int = 1
+
+
+@dataclass
+class PineconeDestinationConnector(IngestDocSessionHandleMixin, BaseDestinationConnector):
+    write_config: PineconeWriteConfig
+    connector_config: SimplePineconeConfig
+    _index: t.Optional["PineconeIndex"] = None
+
+    def to_dict(self, **kwargs):
+        """
+        The _index 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, "_index"):
+            setattr(self_cp, "_index", None)
+        return _asdict(self_cp, **kwargs)
+
+    @property
+    def pinecone_index(self):
+        if self._index is None:
+            self._index = self.create_index()
+        return self._index
+
+    def initialize(self):
+        pass
+
+    @requires_dependencies(["pinecone"], extras="pinecone")
+    def create_index(self) -> "PineconeIndex":
+        from pinecone import Pinecone
+        from unstructured import __version__ as unstructured_version
+
+        pc = Pinecone(
+            api_key=self.connector_config.access_config.api_key,
+            source_tag=f"unstructured=={unstructured_version}",
+        )
+
+        index = pc.Index(self.connector_config.index_name)
+        logger.debug(f"connected to index: {pc.describe_index(self.connector_config.index_name)}")
+        return index
+
+    @DestinationConnectionError.wrap
+    def check_connection(self):
+        _ = self.pinecone_index
+
+    @DestinationConnectionError.wrap
+    @requires_dependencies(["pinecone"], extras="pinecone")
+    def upsert_batch(self, batch):
+        import pinecone.core.client.exceptions
+
+        index = self.pinecone_index
+        try:
+            response = index.upsert(batch)
+        except pinecone.core.client.exceptions.ApiException as api_error:
+            raise WriteError(f"http error: {api_error}") from api_error
+        logger.debug(f"results: {response}")
+
+    def write_dict(self, *args, elements_dict: t.List[t.Dict[str, t.Any]], **kwargs) -> None:
+        logger.info(
+            f"Upserting {len(elements_dict)} elements to destination "
+            f"index at {self.connector_config.index_name}",
+        )
+
+        pinecone_batch_size = self.write_config.batch_size
+
+        logger.info(f"using {self.write_config.num_processes} processes to upload")
+        if self.write_config.num_processes == 1:
+            for chunk in batch_generator(elements_dict, pinecone_batch_size):
+                self.upsert_batch(chunk)  # noqa: E203
+
+        else:
+            with mp.Pool(
+                processes=self.write_config.num_processes,
+            ) as pool:
+                pool.map(
+                    self.upsert_batch, list(batch_generator(elements_dict, pinecone_batch_size))
+                )
+
+    def normalize_dict(self, element_dict: dict) -> dict:
+        # While flatten_dict enables indexing on various fields,
+        # element_serialized enables easily reloading the element object to memory.
+        # element_serialized is formed without text/embeddings to avoid data bloating.
+        return {
+            "id": str(uuid.uuid4()),
+            "values": element_dict.pop("embeddings", None),
+            "metadata": {
+                "text": element_dict.pop("text", None),
+                "element_serialized": json.dumps(element_dict),
+                **flatten_dict(
+                    element_dict,
+                    separator="-",
+                    flatten_lists=True,
+                    remove_none=True,
+                ),
+            },
+        }