unstructured-ingest 0.0.0__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/onedrive.py

@@ -0,0 +1,216 @@
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from time import time
+from typing import TYPE_CHECKING, Any, Generator, Optional
+
+from dateutil import parser
+from unstructured.documents.elements import DataSourceMetadata
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.error import SourceConnectionError, SourceConnectionNetworkError
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    Downloader,
+    DownloaderConfig,
+    DownloadResponse,
+    FileData,
+    Indexer,
+    IndexerConfig,
+    SourceIdentifiers,
+    download_responses,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    SourceRegistryEntry,
+)
+
+if TYPE_CHECKING:
+    from office365.graph_client import GraphClient
+    from office365.onedrive.driveitems.driveItem import DriveItem
+
+CONNECTOR_TYPE = "onedrive"
+MAX_MB_SIZE = 512_000_000
+
+
+@dataclass
+class OnedriveAccessConfig(AccessConfig):
+    client_cred: str
+
+
+@dataclass
+class OnedriveConnectionConfig(ConnectionConfig):
+    client_id: str
+    user_pname: str
+    tenant: str = field(repr=False)
+    authority_url: Optional[str] = field(repr=False, default="https://login.microsoftonline.com")
+    access_config: OnedriveAccessConfig = enhanced_field(sensitive=True)
+
+    @requires_dependencies(["msal"], extras="onedrive")
+    def get_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_cred,
+            )
+            token = app.acquire_token_for_client(scopes=["https://graph.microsoft.com/.default"])
+        except ValueError as exc:
+            logger.error("Couldn't set up credentials for OneDrive")
+            raise exc
+        if "error" in token:
+            raise SourceConnectionNetworkError(
+                "failed to fetch token, {}: {}".format(token["error"], token["error_description"])
+            )
+        return token
+
+    @requires_dependencies(["office365"], extras="onedrive")
+    def get_client(self) -> "GraphClient":
+        from office365.graph_client import GraphClient
+
+        client = GraphClient(self.get_token)
+        return client
+
+
+@dataclass
+class OnedriveIndexerConfig(IndexerConfig):
+    path: Optional[str] = field(default="")
+    recursive: bool = False
+
+
+@dataclass
+class OnedriveIndexer(Indexer):
+    connection_config: OnedriveConnectionConfig
+    index_config: OnedriveIndexerConfig
+
+    def list_objects(self, folder, recursive) -> list["DriveItem"]:
+        drive_items = folder.children.get().execute_query()
+        files = [d for d in drive_items if d.is_file]
+        if not recursive:
+            return files
+        folders = [d for d in drive_items if d.is_folder]
+        for f in folders:
+            files.extend(self.list_objects(f, recursive))
+        return files
+
+    def get_root(self, client: "GraphClient") -> "DriveItem":
+        root = client.users[self.connection_config.user_pname].drive.get().execute_query().root
+        if fpath := self.index_config.path:
+            root = root.get_by_path(fpath).get().execute_query()
+            if root is None or not root.is_folder:
+                raise ValueError(f"Unable to find directory, given: {fpath}")
+        return root
+
+    def get_properties(self, drive_item: "DriveItem") -> dict:
+        properties = drive_item.properties
+        filtered_properties = {}
+        for k, v in properties.items():
+            try:
+                json.dumps(v)
+                filtered_properties[k] = v
+            except TypeError:
+                pass
+        return filtered_properties
+
+    def drive_item_to_file_data(self, drive_item: "DriveItem") -> FileData:
+        file_path = drive_item.parent_reference.path.split(":")[-1]
+        file_path = file_path[1:] if file_path and file_path[0] == "/" else file_path
+        filename = drive_item.name
+        server_path = file_path + "/" + filename
+        rel_path = server_path.replace(self.index_config.path, "").lstrip("/")
+        date_modified_dt = (
+            parser.parse(drive_item.last_modified_datetime)
+            if drive_item.last_modified_datetime
+            else None
+        )
+        date_created_at = (
+            parser.parse(drive_item.created_datetime) if drive_item.created_datetime else None
+        )
+        return FileData(
+            identifier=drive_item.id,
+            connector_type=CONNECTOR_TYPE,
+            source_identifiers=SourceIdentifiers(
+                fullpath=server_path, filename=drive_item.name, rel_path=rel_path
+            ),
+            metadata=DataSourceMetadata(
+                url=drive_item.parent_reference.path + "/" + drive_item.name,
+                version=drive_item.etag,
+                date_modified=str(date_modified_dt.timestamp()) if date_modified_dt else None,
+                date_created=str(date_created_at.timestamp()) if date_modified_dt else None,
+                date_processed=str(time()),
+                record_locator={
+                    "user_pname": self.connection_config.user_pname,
+                    "server_relative_path": server_path,
+                },
+            ),
+            additional_metadata=self.get_properties(drive_item=drive_item),
+        )
+
+    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+        client = self.connection_config.get_client()
+        root = self.get_root(client=client)
+        drive_items = self.list_objects(folder=root, recursive=self.index_config.recursive)
+        for drive_item in drive_items:
+            file_data = self.drive_item_to_file_data(drive_item=drive_item)
+            yield file_data
+
+
+@dataclass
+class OnedriveDownloaderConfig(DownloaderConfig):
+    pass
+
+
+@dataclass
+class OnedriveDownloader(Downloader):
+    connection_config: OnedriveConnectionConfig
+    download_config: OnedriveDownloaderConfig
+
+    @SourceConnectionNetworkError.wrap
+    def _fetch_file(self, file_data: FileData):
+        if file_data.source_identifiers is None or not file_data.source_identifiers.fullpath:
+            raise ValueError(
+                f"file data doesn't have enough information to get "
+                f"file content: {file_data.to_dict()}"
+            )
+
+        server_relative_path = file_data.source_identifiers.fullpath
+        client = self.connection_config.get_client()
+        root = client.users[self.connection_config.user_pname].drive.get().execute_query().root
+        file = root.get_by_path(server_relative_path).get().execute_query()
+        if not file:
+            raise FileNotFoundError(f"file not found: {server_relative_path}")
+        return file
+
+    def get_download_path(self, file_data: FileData) -> Optional[Path]:
+        rel_path = file_data.source_identifiers.relative_path
+        rel_path = rel_path[1:] if rel_path.startswith("/") else rel_path
+        return self.download_dir / Path(rel_path)
+
+    @SourceConnectionError.wrap
+    def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
+        file = self._fetch_file(file_data=file_data)
+        fsize = file.get_property("size", 0)
+        download_path = self.get_download_path(file_data=file_data)
+        download_path.parent.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Downloading {file_data.source_identifiers.fullpath} to {download_path}")
+        if fsize > MAX_MB_SIZE:
+            logger.info(f"Downloading file with size: {fsize} bytes in chunks")
+            with download_path.open(mode="wb") as f:
+                file.download_session(f, chunk_size=1024 * 1024 * 100).execute_query()
+        else:
+            with download_path.open(mode="wb") as f:
+                file.download(f).execute_query()
+        return DownloadResponse(file_data=file_data, path=download_path)
+
+
+onedrive_source_entry = SourceRegistryEntry(
+    connection_config=OnedriveConnectionConfig,
+    indexer_config=OnedriveIndexerConfig,
+    indexer=OnedriveIndexer,
+    downloader_config=OnedriveDownloaderConfig,
+    downloader=OnedriveDownloader,
+)

unstructured_ingest/v2/processes/connectors/opensearch.py

@@ -0,0 +1,155 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Optional
+
+from unstructured_ingest.enhanced_dataclass import EnhancedDataClassJsonMixin, enhanced_field
+from unstructured_ingest.error import (
+    DestinationConnectionError,
+)
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.v2.processes.connectors.elasticsearch import (
+    ElasticsearchDownloader,
+    ElasticsearchDownloaderConfig,
+    ElasticsearchIndexer,
+    ElasticsearchIndexerConfig,
+    ElasticsearchUploader,
+    ElasticsearchUploaderConfig,
+    ElasticsearchUploadStager,
+    ElasticsearchUploadStagerConfig,
+)
+
+if TYPE_CHECKING:
+    from opensearchpy import OpenSearch
+
+CONNECTOR_TYPE = "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):
+    password: Optional[str] = enhanced_field(default=None, sensitive=True)
+    use_ssl: bool = False
+    verify_certs: bool = False
+    ssl_show_warn: bool = False
+    ca_certs: Optional[str] = None
+    client_cert: Optional[str] = None
+    client_key: Optional[str] = None
+
+
+@dataclass
+class OpenSearchClientInput(EnhancedDataClassJsonMixin):
+    http_auth: Optional[tuple[str, str]] = enhanced_field(sensitive=True, default=None)
+    hosts: Optional[list[str]] = None
+    use_ssl: bool = False
+    verify_certs: bool = False
+    ssl_show_warn: bool = False
+    ca_certs: Optional[str] = None
+    client_cert: Optional[str] = None
+    client_key: Optional[str] = None
+
+
+@dataclass
+class OpenSearchConnectionConfig(ConnectionConfig):
+    hosts: Optional[list[str]] = None
+    username: Optional[str] = None
+    access_config: OpenSearchAccessConfig = enhanced_field(sensitive=True)
+
+    def get_client_kwargs(self) -> dict:
+        # Update auth related fields to conform to what the SDK expects based on the
+        # supported methods:
+        # https://github.com/opensearch-project/opensearch-py/blob/main/opensearchpy/client/__init__.py
+        client_input = OpenSearchClientInput()
+        if self.hosts:
+            client_input.hosts = self.hosts
+        if self.access_config.use_ssl:
+            client_input.use_ssl = self.access_config.use_ssl
+        if self.access_config.verify_certs:
+            client_input.verify_certs = self.access_config.verify_certs
+        if self.access_config.ssl_show_warn:
+            client_input.ssl_show_warn = self.access_config.ssl_show_warn
+        if self.access_config.ca_certs:
+            client_input.ca_certs = self.access_config.ca_certs
+        if self.access_config.client_cert:
+            client_input.client_cert = self.access_config.client_cert
+        if self.access_config.client_key:
+            client_input.client_key = self.access_config.client_key
+        if self.username and self.access_config.password:
+            client_input.http_auth = (self.username, self.access_config.password)
+        logger.debug(
+            f"OpenSearch client inputs mapped to: {client_input.to_dict(redact_sensitive=True)}"
+        )
+        client_kwargs = client_input.to_dict(redact_sensitive=False)
+        client_kwargs = {k: v for k, v in client_kwargs.items() if v is not None}
+        return client_kwargs
+
+    @DestinationConnectionError.wrap
+    @requires_dependencies(["opensearchpy"], extras="opensearch")
+    def get_client(self) -> "OpenSearch":
+        from opensearchpy import OpenSearch
+
+        return OpenSearch(**self.get_client_kwargs())
+
+
+@dataclass
+class OpenSearchIndexer(ElasticsearchIndexer):
+    connection_config: OpenSearchConnectionConfig
+    client: "OpenSearch" = field(init=False)
+
+    @requires_dependencies(["opensearchpy"], extras="opensearch")
+    def load_scan(self):
+        from opensearchpy.helpers import scan
+
+        return scan
+
+
+@dataclass
+class OpenSearchDownloader(ElasticsearchDownloader):
+    connection_config: OpenSearchConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @requires_dependencies(["opensearchpy"], extras="opensearch")
+    def load_async(self):
+        from opensearchpy import AsyncOpenSearch
+        from opensearchpy.helpers import async_scan
+
+        return AsyncOpenSearch, async_scan
+
+
+@dataclass
+class OpenSearchUploader(ElasticsearchUploader):
+    connection_config: OpenSearchConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @requires_dependencies(["opensearchpy"], extras="opensearch")
+    def load_parallel_bulk(self):
+        from opensearchpy.helpers import parallel_bulk
+
+        return parallel_bulk
+
+
+opensearch_source_entry = SourceRegistryEntry(
+    connection_config=OpenSearchConnectionConfig,
+    indexer=OpenSearchIndexer,
+    indexer_config=ElasticsearchIndexerConfig,
+    downloader=OpenSearchDownloader,
+    downloader_config=ElasticsearchDownloaderConfig,
+)
+
+
+opensearch_destination_entry = DestinationRegistryEntry(
+    connection_config=OpenSearchConnectionConfig,
+    upload_stager_config=ElasticsearchUploadStagerConfig,
+    upload_stager=ElasticsearchUploadStager,
+    uploader_config=ElasticsearchUploaderConfig,
+    uploader=OpenSearchUploader,
+)

unstructured_ingest/v2/processes/connectors/pinecone.py

@@ -0,0 +1,178 @@
+import json
+import multiprocessing as mp
+import uuid
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional
+
+from unstructured.ingest.v2.logger import logger
+from unstructured.ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+)
+from unstructured.staging.base import flatten_dict
+from unstructured.utils import requires_dependencies
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.utils.data_prep import batch_generator
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    UploadContent,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+
+if TYPE_CHECKING:
+    from pinecone import Index as PineconeIndex
+
+
+CONNECTOR_TYPE = "pinecone"
+
+
+@dataclass
+class PineconeAccessConfig(AccessConfig):
+    api_key: Optional[str] = enhanced_field(default=None, overload_name="pinecone_api_key")
+
+
+@dataclass
+class PineconeConnectionConfig(ConnectionConfig):
+    index_name: str
+    environment: str
+    access_config: PineconeAccessConfig = enhanced_field(sensitive=True)
+
+    @requires_dependencies(["pinecone"], extras="pinecone")
+    def get_index(self) -> "PineconeIndex":
+        from pinecone import Pinecone
+        from unstructured import __version__ as unstructured_version
+
+        pc = Pinecone(
+            api_key=self.access_config.api_key,
+            source_tag=f"unstructured=={unstructured_version}",
+        )
+
+        index = pc.Index(self.index_name)
+        logger.debug(f"Connected to index: {pc.describe_index(self.index_name)}")
+        return index
+
+
+@dataclass
+class PineconeUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class PineconeUploaderConfig(UploaderConfig):
+    batch_size: int = 100
+    num_of_processes: int = 4
+
+
+@dataclass
+class PineconeUploadStager(UploadStager):
+    upload_stager_config: PineconeUploadStagerConfig = field(
+        default_factory=lambda: PineconeUploadStagerConfig()
+    )
+
+    @staticmethod
+    def conform_dict(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,
+                ),
+            },
+        }
+
+    def run(
+        self,
+        elements_filepath: Path,
+        output_dir: Path,
+        output_filename: str,
+        **kwargs: Any,
+    ) -> Path:
+        with open(elements_filepath) as elements_file:
+            elements_contents = json.load(elements_file)
+
+        conformed_elements = [
+            self.conform_dict(element_dict=element) for element in elements_contents
+        ]
+
+        output_path = Path(output_dir) / Path(f"{output_filename}.json")
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(output_path, "w") as output_file:
+            json.dump(conformed_elements, output_file)
+        return output_path
+
+
+@dataclass
+class PineconeUploader(Uploader):
+    upload_config: PineconeUploaderConfig
+    connection_config: PineconeConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @DestinationConnectionError.wrap
+    def check_connection(self):
+        _ = self.connection_config.get_index()
+
+    @requires_dependencies(["pinecone"], extras="pinecone")
+    def upsert_batch(self, batch):
+        from pinecone.exceptions import PineconeApiException
+
+        try:
+            index = self.connection_config.get_index()
+            response = index.upsert(batch)
+        except PineconeApiException as api_error:
+            raise DestinationConnectionError(f"http error: {api_error}") from api_error
+        logger.debug(f"results: {response}")
+
+    def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
+
+        elements_dict = []
+        for content in contents:
+            with open(content.path) as elements_file:
+                elements = json.load(elements_file)
+                elements_dict.extend(elements)
+
+        logger.info(
+            f"writing document batches to destination"
+            f" index named {self.connection_config.index_name}"
+            f" environment named {self.connection_config.environment}"
+            f" with batch size {self.upload_config.batch_size}"
+            f" with {self.upload_config.num_of_processes} (number of) processes"
+        )
+
+        pinecone_batch_size = self.upload_config.batch_size
+
+        if self.upload_config.num_of_processes == 1:
+            for batch in batch_generator(elements_dict, pinecone_batch_size):
+                self.upsert_batch(batch)  # noqa: E203
+
+        else:
+            with mp.Pool(
+                processes=self.upload_config.num_of_processes,
+            ) as pool:
+                pool.map(
+                    self.upsert_batch, list(batch_generator(elements_dict, pinecone_batch_size))
+                )
+
+
+pinecone_destination_entry = DestinationRegistryEntry(
+    connection_config=PineconeConnectionConfig,
+    uploader=PineconeUploader,
+    uploader_config=PineconeUploaderConfig,
+    upload_stager=PineconeUploadStager,
+    upload_stager_config=PineconeUploadStagerConfig,
+)