unstructured-ingest 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/onedrive.py

@@ -9,7 +9,11 @@ from typing import TYPE_CHECKING, Any, Generator, Optional
 from dateutil import parser
 from pydantic import Field, Secret
 
-from unstructured_ingest.error import SourceConnectionError, SourceConnectionNetworkError
+from unstructured_ingest.error import (
+    DestinationConnectionError,
+    SourceConnectionError,
+    SourceConnectionNetworkError,
+)
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
@@ -22,16 +26,20 @@ from unstructured_ingest.v2.interfaces import (
     Indexer,
     IndexerConfig,
     SourceIdentifiers,
+    Uploader,
+    UploaderConfig,
     download_responses,
 )
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
     SourceRegistryEntry,
 )
 
 if TYPE_CHECKING:
     from office365.graph_client import GraphClient
     from office365.onedrive.driveitems.driveItem import DriveItem
+    from office365.onedrive.drives.drive import Drive
 
 CONNECTOR_TYPE = "onedrive"
 MAX_MB_SIZE = 512_000_000
@@ -55,6 +63,11 @@ class OnedriveConnectionConfig(ConnectionConfig):
     )
     access_config: Secret[OnedriveAccessConfig]
 
+    def get_drive(self) -> "Drive":
+        client = self.get_client()
+        drive = client.users[self.user_pname].drive
+        return drive
+
     @requires_dependencies(["msal"], extras="onedrive")
     def get_token(self):
         from msal import ConfidentialClientApplication
@@ -100,7 +113,6 @@ class OnedriveIndexer(Indexer):
                 raise SourceConnectionError(
                     "{} ({})".format(error, token_resp.get("error_description"))
                 )
-            self.connection_config.get_client()
         except Exception as e:
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise SourceConnectionError(f"failed to validate connection: {e}")
@@ -224,6 +236,149 @@ class OnedriveDownloader(Downloader):
         return DownloadResponse(file_data=file_data, path=download_path)
 
 
+class OnedriveUploaderConfig(UploaderConfig):
+    remote_url: str = Field(
+        description="URL of the destination in OneDrive, e.g., 'onedrive://Documents/Folder'"
+    )
+    prefix: str = "onedrive://"
+
+    @property
+    def root_folder(self) -> str:
+        url = (
+            self.remote_url.replace(self.prefix, "", 1)
+            if self.remote_url.startswith(self.prefix)
+            else self.remote_url
+        )
+        return url.split("/")[0]
+
+    @property
+    def url(self) -> str:
+        url = (
+            self.remote_url.replace(self.prefix, "", 1)
+            if self.remote_url.startswith(self.prefix)
+            else self.remote_url
+        )
+        return url
+
+
+@dataclass
+class OnedriveUploader(Uploader):
+    connection_config: OnedriveConnectionConfig
+    upload_config: OnedriveUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @requires_dependencies(["office365"], extras="onedrive")
+    def precheck(self) -> None:
+        from office365.runtime.client_request_exception import ClientRequestException
+
+        try:
+            token_resp: dict = self.connection_config.get_token()
+            if error := token_resp.get("error"):
+                raise SourceConnectionError(
+                    "{} ({})".format(error, token_resp.get("error_description"))
+                )
+            drive = self.connection_config.get_drive()
+            root = drive.root
+            root_folder = self.upload_config.root_folder
+            folder = root.get_by_path(root_folder)
+            try:
+                folder.get().execute_query()
+            except ClientRequestException as e:
+                if e.message != "The resource could not be found.":
+                    raise e
+                folder = root.create_folder(root_folder).execute_query()
+                logger.info(f"successfully created folder: {folder.name}")
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        drive = self.connection_config.get_drive()
+
+        # Use the remote_url from upload_config as the base destination folder
+        base_destination_folder = self.upload_config.url
+
+        # Use the file's relative path to maintain directory structure, if needed
+        if file_data.source_identifiers and file_data.source_identifiers.rel_path:
+            # Combine the base destination folder with the file's relative path
+            destination_path = Path(base_destination_folder) / Path(
+                file_data.source_identifiers.rel_path
+            )
+        else:
+            # If no relative path is provided, upload directly to the base destination folder
+            destination_path = Path(base_destination_folder) / path.name
+
+        destination_folder = destination_path.parent
+        file_name = destination_path.name
+
+        # Convert destination folder to a string suitable for OneDrive API
+        destination_folder_str = str(destination_folder).replace("\\", "/")
+
+        # Resolve the destination folder in OneDrive, creating it if necessary
+        try:
+            # Attempt to get the folder
+            folder = drive.root.get_by_path(destination_folder_str)
+            folder.get().execute_query()
+        except Exception:
+            # Folder doesn't exist, create it recursively
+            current_folder = drive.root
+            for part in destination_folder.parts:
+                # Use filter to find the folder by name
+                folders = (
+                    current_folder.children.filter(f"name eq '{part}' and folder ne null")
+                    .get()
+                    .execute_query()
+                )
+                if folders:
+                    current_folder = folders[0]
+                else:
+                    # Folder doesn't exist, create it
+                    current_folder = current_folder.create_folder(part).execute_query()
+            folder = current_folder
+
+        # Check the size of the file
+        file_size = path.stat().st_size
+
+        if file_size < MAX_MB_SIZE:
+            # Use simple upload for small files
+            with path.open("rb") as local_file:
+                content = local_file.read()
+                logger.info(f"Uploading {path} to {destination_path} using simple upload")
+                try:
+                    uploaded_file = folder.upload(file_name, content).execute_query()
+                    if not uploaded_file or uploaded_file.name != file_name:
+                        raise DestinationConnectionError(f"Upload failed for file '{file_name}'")
+                    # Log details about the uploaded file
+                    logger.info(
+                        f"Uploaded file '{uploaded_file.name}' with ID '{uploaded_file.id}'"
+                    )
+                except Exception as e:
+                    logger.error(f"Failed to upload file '{file_name}': {e}", exc_info=True)
+                    raise DestinationConnectionError(
+                        f"Failed to upload file '{file_name}': {e}"
+                    ) from e
+        else:
+            # Use resumable upload for large files
+            destination_fullpath = f"{destination_folder_str}/{file_name}"
+            destination_drive_item = drive.root.item_with_path(destination_fullpath)
+
+            logger.info(f"Uploading {path} to {destination_fullpath} using resumable upload")
+            try:
+                uploaded_file = destination_drive_item.resumable_upload(
+                    source_path=str(path)
+                ).execute_query()
+                # Validate the upload
+                if not uploaded_file or uploaded_file.name != file_name:
+                    raise DestinationConnectionError(f"Upload failed for file '{file_name}'")
+                # Log details about the uploaded file
+                logger.info(f"Uploaded file {uploaded_file.name} with ID {uploaded_file.id}")
+            except Exception as e:
+                logger.error(f"Failed to upload file '{file_name}' using resumable upload: {e}")
+                raise DestinationConnectionError(
+                    f"Failed to upload file '{file_name}' using resumable upload: {e}"
+                ) from e
+
+
 onedrive_source_entry = SourceRegistryEntry(
     connection_config=OnedriveConnectionConfig,
     indexer_config=OnedriveIndexerConfig,
@@ -231,3 +386,9 @@ onedrive_source_entry = SourceRegistryEntry(
     downloader_config=OnedriveDownloaderConfig,
     downloader=OnedriveDownloader,
 )
+
+onedrive_destination_entry = DestinationRegistryEntry(
+    connection_config=OnedriveConnectionConfig,
+    uploader=OnedriveUploader,
+    uploader_config=OnedriveUploaderConfig,
+)

unstructured_ingest/v2/processes/connectors/qdrant/__init__.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from unstructured_ingest.v2.processes.connector_registry import (
+    add_destination_entry,
+)
+
+from .cloud import CONNECTOR_TYPE as CLOUD_CONNECTOR_TYPE
+from .cloud import qdrant_cloud_destination_entry
+from .local import CONNECTOR_TYPE as LOCAL_CONNECTOR_TYPE
+from .local import qdrant_local_destination_entry
+from .server import CONNECTOR_TYPE as SERVER_CONNECTOR_TYPE
+from .server import qdrant_server_destination_entry
+
+add_destination_entry(destination_type=CLOUD_CONNECTOR_TYPE, entry=qdrant_cloud_destination_entry)
+add_destination_entry(destination_type=SERVER_CONNECTOR_TYPE, entry=qdrant_server_destination_entry)
+add_destination_entry(destination_type=LOCAL_CONNECTOR_TYPE, entry=qdrant_local_destination_entry)

unstructured_ingest/v2/processes/connectors/qdrant/cloud.py

@@ -0,0 +1,59 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connectors.qdrant.qdrant import (
+    QdrantAccessConfig,
+    QdrantConnectionConfig,
+    QdrantUploader,
+    QdrantUploaderConfig,
+    QdrantUploadStager,
+    QdrantUploadStagerConfig,
+)
+
+CONNECTOR_TYPE = "qdrant-cloud"
+
+
+class CloudQdrantAccessConfig(QdrantAccessConfig):
+    api_key: str = Field(description="Qdrant API key")
+
+
+class CloudQdrantConnectionConfig(QdrantConnectionConfig):
+    url: str = Field(default=None, description="url of Qdrant Cloud")
+    access_config: Secret[CloudQdrantAccessConfig]
+
+    def get_client_kwargs(self) -> dict:
+        return {
+            "api_key": self.access_config.get_secret_value().api_key,
+            "url": self.url,
+        }
+
+
+class CloudQdrantUploadStagerConfig(QdrantUploadStagerConfig):
+    pass
+
+
+@dataclass
+class CloudQdrantUploadStager(QdrantUploadStager):
+    upload_stager_config: CloudQdrantUploadStagerConfig
+
+
+class CloudQdrantUploaderConfig(QdrantUploaderConfig):
+    pass
+
+
+@dataclass
+class CloudQdrantUploader(QdrantUploader):
+    connection_config: CloudQdrantConnectionConfig
+    upload_config: CloudQdrantUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+qdrant_cloud_destination_entry = DestinationRegistryEntry(
+    connection_config=CloudQdrantConnectionConfig,
+    uploader=CloudQdrantUploader,
+    uploader_config=CloudQdrantUploaderConfig,
+    upload_stager=CloudQdrantUploadStager,
+    upload_stager_config=CloudQdrantUploadStagerConfig,
+)

unstructured_ingest/v2/processes/connectors/qdrant/local.py

@@ -0,0 +1,58 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connectors.qdrant.qdrant import (
+    QdrantAccessConfig,
+    QdrantConnectionConfig,
+    QdrantUploader,
+    QdrantUploaderConfig,
+    QdrantUploadStager,
+    QdrantUploadStagerConfig,
+)
+
+CONNECTOR_TYPE = "qdrant-local"
+
+
+class LocalQdrantAccessConfig(QdrantAccessConfig):
+    pass
+
+
+class LocalQdrantConnectionConfig(QdrantConnectionConfig):
+    path: str = Field(default=None, description="Persistence path for QdrantLocal.")
+    access_config: Secret[LocalQdrantAccessConfig] = Field(
+        default_factory=LocalQdrantAccessConfig, validate_default=True
+    )
+
+    def get_client_kwargs(self) -> dict:
+        return {"path": self.path}
+
+
+class LocalQdrantUploadStagerConfig(QdrantUploadStagerConfig):
+    pass
+
+
+@dataclass
+class LocalQdrantUploadStager(QdrantUploadStager):
+    upload_stager_config: LocalQdrantUploadStagerConfig
+
+
+class LocalQdrantUploaderConfig(QdrantUploaderConfig):
+    pass
+
+
+@dataclass
+class LocalQdrantUploader(QdrantUploader):
+    connection_config: LocalQdrantConnectionConfig
+    upload_config: LocalQdrantUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+qdrant_local_destination_entry = DestinationRegistryEntry(
+    connection_config=LocalQdrantConnectionConfig,
+    uploader=LocalQdrantUploader,
+    uploader_config=LocalQdrantUploaderConfig,
+    upload_stager=LocalQdrantUploadStager,
+    upload_stager_config=LocalQdrantUploadStagerConfig,
+)

unstructured_ingest/v2/processes/connectors/qdrant/qdrant.py

@@ -0,0 +1,168 @@
+import asyncio
+import json
+import uuid
+from abc import ABC, abstractmethod
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.error import DestinationConnectionError, WriteError
+from unstructured_ingest.utils.data_prep import batch_generator, flatten_dict
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    FileData,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+from unstructured_ingest.v2.logger import logger
+
+if TYPE_CHECKING:
+    from qdrant_client import AsyncQdrantClient
+
+
+class QdrantAccessConfig(AccessConfig, ABC):
+    pass
+
+
+class QdrantConnectionConfig(ConnectionConfig, ABC):
+    access_config: Secret[QdrantAccessConfig] = Field(
+        default_factory=QdrantAccessConfig, validate_default=True, description="Access Config"
+    )
+
+    @abstractmethod
+    def get_client_kwargs(self) -> dict:
+        pass
+
+    @requires_dependencies(["qdrant_client"], extras="qdrant")
+    @asynccontextmanager
+    async def get_client(self) -> AsyncGenerator["AsyncQdrantClient", None]:
+        from qdrant_client.async_qdrant_client import AsyncQdrantClient
+
+        client_kwargs = self.get_client_kwargs()
+        client = AsyncQdrantClient(**client_kwargs)
+        try:
+            yield client
+        finally:
+            await client.close()
+
+
+class QdrantUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class QdrantUploadStager(UploadStager, ABC):
+    upload_stager_config: QdrantUploadStagerConfig = field(
+        default_factory=lambda: QdrantUploadStagerConfig()
+    )
+
+    @staticmethod
+    def conform_dict(data: dict) -> dict:
+        """Prepares dictionary in the format that Chroma requires"""
+        return {
+            "id": str(uuid.uuid4()),
+            "vector": data.pop("embeddings", {}),
+            "payload": {
+                "text": data.pop("text", None),
+                "element_serialized": json.dumps(data),
+                **flatten_dict(
+                    data,
+                    separator="-",
+                    flatten_lists=True,
+                ),
+            },
+        }
+
+    def run(
+        self,
+        elements_filepath: Path,
+        file_data: FileData,
+        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(data=element) for element in elements_contents]
+        output_path = Path(output_dir) / Path(f"{output_filename}.json")
+
+        with open(output_path, "w") as output_file:
+            json.dump(conformed_elements, output_file)
+        return output_path
+
+
+class QdrantUploaderConfig(UploaderConfig):
+    collection_name: str = Field(description="Name of the collection.")
+    batch_size: int = Field(default=50, description="Number of records per batch.")
+    num_processes: Optional[int] = Field(
+        default=1,
+        description="Optional limit on number of threads to use for upload.",
+        deprecated=True,
+    )
+
+
+@dataclass
+class QdrantUploader(Uploader, ABC):
+    upload_config: QdrantUploaderConfig
+    connection_config: QdrantConnectionConfig
+
+    @DestinationConnectionError.wrap
+    def precheck(self) -> None:
+        async def check_connection():
+            async with self.connection_config.get_client() as async_client:
+                await async_client.get_collections()
+
+        asyncio.run(check_connection())
+
+    def is_async(self):
+        return True
+
+    async def run_async(
+        self,
+        path: Path,
+        file_data: FileData,
+        **kwargs: Any,
+    ) -> None:
+        with path.open("r") as file:
+            elements: list[dict] = json.load(file)
+
+        logger.debug("Loaded %i elements from %s", len(elements), path)
+
+        batches = list(batch_generator(elements, batch_size=self.upload_config.batch_size))
+        logger.debug(
+            "Elements split into %i batches of size %i.",
+            len(batches),
+            self.upload_config.batch_size,
+        )
+        await asyncio.gather(*[self._upsert_batch(batch) for batch in batches])
+
+    async def _upsert_batch(self, batch: list[dict]) -> None:
+        from qdrant_client import models
+
+        points: list[models.PointStruct] = [models.PointStruct(**item) for item in batch]
+        try:
+            logger.debug(
+                "Upserting %i points to the '%s' collection.",
+                len(points),
+                self.upload_config.collection_name,
+            )
+            async with self.connection_config.get_client() as async_client:
+                await async_client.upsert(
+                    self.upload_config.collection_name, points=points, wait=True
+                )
+        except Exception as api_error:
+            logger.error(
+                "Failed to upsert points to the collection due to the following error %s", api_error
+            )
+
+            raise WriteError(f"Qdrant error: {api_error}") from api_error
+
+        logger.debug("Successfully upsert points to the collection.")

unstructured_ingest/v2/processes/connectors/qdrant/server.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connectors.qdrant.qdrant import (
+    QdrantAccessConfig,
+    QdrantConnectionConfig,
+    QdrantUploader,
+    QdrantUploaderConfig,
+    QdrantUploadStager,
+    QdrantUploadStagerConfig,
+)
+
+CONNECTOR_TYPE = "qdrant-server"
+
+
+class ServerQdrantAccessConfig(QdrantAccessConfig):
+    pass
+
+
+class ServerQdrantConnectionConfig(QdrantConnectionConfig):
+    url: str = Field(default=None, description="url of Qdrant server")
+    access_config: Secret[ServerQdrantAccessConfig] = Field(
+        default_factory=ServerQdrantAccessConfig, validate_default=True
+    )
+
+    def get_client_kwargs(self) -> dict:
+        return {
+            "url": self.url,
+        }
+
+
+class ServerQdrantUploadStagerConfig(QdrantUploadStagerConfig):
+    pass
+
+
+@dataclass
+class ServerQdrantUploadStager(QdrantUploadStager):
+    upload_stager_config: ServerQdrantUploadStagerConfig
+
+
+class ServerQdrantUploaderConfig(QdrantUploaderConfig):
+    pass
+
+
+@dataclass
+class ServerQdrantUploader(QdrantUploader):
+    connection_config: ServerQdrantConnectionConfig
+    upload_config: ServerQdrantUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+qdrant_server_destination_entry = DestinationRegistryEntry(
+    connection_config=ServerQdrantConnectionConfig,
+    uploader=ServerQdrantUploader,
+    uploader_config=ServerQdrantUploaderConfig,
+    upload_stager=ServerQdrantUploadStager,
+    upload_stager_config=ServerQdrantUploadStagerConfig,
+)

unstructured_ingest/v2/processes/connectors/sql/__init__.py

@@ -7,6 +7,8 @@ from unstructured_ingest.v2.processes.connector_registry import (
 
 from .postgres import CONNECTOR_TYPE as POSTGRES_CONNECTOR_TYPE
 from .postgres import postgres_destination_entry, postgres_source_entry
+from .singlestore import CONNECTOR_TYPE as SINGLESTORE_CONNECTOR_TYPE
+from .singlestore import singlestore_destination_entry
 from .snowflake import CONNECTOR_TYPE as SNOWFLAKE_CONNECTOR_TYPE
 from .snowflake import snowflake_destination_entry, snowflake_source_entry
 from .sqlite import CONNECTOR_TYPE as SQLITE_CONNECTOR_TYPE
@@ -19,3 +21,6 @@ add_source_entry(source_type=SNOWFLAKE_CONNECTOR_TYPE, entry=snowflake_source_en
 add_destination_entry(destination_type=SQLITE_CONNECTOR_TYPE, entry=sqlite_destination_entry)
 add_destination_entry(destination_type=POSTGRES_CONNECTOR_TYPE, entry=postgres_destination_entry)
 add_destination_entry(destination_type=SNOWFLAKE_CONNECTOR_TYPE, entry=snowflake_destination_entry)
+add_destination_entry(
+    destination_type=SINGLESTORE_CONNECTOR_TYPE, entry=singlestore_destination_entry
+)

unstructured_ingest/v2/processes/connectors/sql/postgres.py

@@ -1,6 +1,6 @@
 from contextlib import contextmanager
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Generator, Optional
+from typing import TYPE_CHECKING, Generator, Optional
 
 from pydantic import Field, Secret
 
@@ -12,7 +12,6 @@ from unstructured_ingest.v2.processes.connector_registry import (
     SourceRegistryEntry,
 )
 from unstructured_ingest.v2.processes.connectors.sql.sql import (
-    _DATE_COLUMNS,
     SQLAccessConfig,
     SQLConnectionConfig,
     SQLDownloader,
@@ -23,7 +22,6 @@ from unstructured_ingest.v2.processes.connectors.sql.sql import (
     SQLUploaderConfig,
     SQLUploadStager,
     SQLUploadStagerConfig,
-    parse_date_string,
 )
 
 if TYPE_CHECKING:
@@ -138,23 +136,6 @@ class PostgresUploader(SQLUploader):
     connector_type: str = CONNECTOR_TYPE
     values_delimiter: str = "%s"
 
-    def prepare_data(
-        self, columns: list[str], data: tuple[tuple[Any, ...], ...]
-    ) -> list[tuple[Any, ...]]:
-        output = []
-        for row in data:
-            parsed = []
-            for column_name, value in zip(columns, row):
-                if column_name in _DATE_COLUMNS:
-                    if value is None:
-                        parsed.append(None)
-                    else:
-                        parsed.append(parse_date_string(value))
-                else:
-                    parsed.append(value)
-            output.append(tuple(parsed))
-        return output
-
 
 postgres_source_entry = SourceRegistryEntry(
     connection_config=PostgresConnectionConfig,