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

unstructured_ingest/v2/processes/connectors/kafka/kafka.py

@@ -0,0 +1,196 @@
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from time import time
+from typing import TYPE_CHECKING, Any, ContextManager, Generator, Optional
+
+from pydantic import Secret
+
+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,
+    FileData,
+    FileDataSourceMetadata,
+    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 confluent_kafka import Consumer
+
+CONNECTOR_TYPE = "kafka"
+
+
+class KafkaAccessConfig(AccessConfig, ABC):
+    pass
+
+
+class KafkaConnectionConfig(ConnectionConfig, ABC):
+    access_config: Secret[KafkaAccessConfig]
+    timeout: Optional[float] = 1.0
+    bootstrap_server: str
+    port: int
+
+    @abstractmethod
+    def get_consumer_configuration(self) -> dict:
+        pass
+
+    @contextmanager
+    @requires_dependencies(["confluent_kafka"], extras="kafka")
+    def get_consumer(self) -> ContextManager["Consumer"]:
+        from confluent_kafka import Consumer
+
+        consumer = Consumer(self.get_consumer_configuration())
+        try:
+            logger.debug("kafka consumer connected")
+            yield consumer
+        finally:
+            consumer.close()
+
+
+class KafkaIndexerConfig(IndexerConfig):
+    topic: str
+    num_messages_to_consume: Optional[int] = 100
+
+    def update_consumer(self, consumer: "Consumer") -> None:
+        consumer.subscribe([self.topic])
+
+
+@dataclass
+class KafkaIndexer(Indexer):
+    connection_config: KafkaConnectionConfig
+    index_config: KafkaIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @contextmanager
+    def get_consumer(self) -> ContextManager["Consumer"]:
+        with self.connection_config.get_consumer() as consumer:
+            self.index_config.update_consumer(consumer=consumer)
+            yield consumer
+
+    @requires_dependencies(["confluent_kafka"], extras="kafka")
+    def generate_messages(self) -> Generator[Any, None, None]:
+        from confluent_kafka import KafkaError, KafkaException
+
+        messages_consumed = 0
+        max_empty_polls = 10
+        empty_polls = 0
+        num_messages_to_consume = self.index_config.num_messages_to_consume
+        with self.get_consumer() as consumer:
+            while messages_consumed < num_messages_to_consume and empty_polls < max_empty_polls:
+                msg = consumer.poll(timeout=self.connection_config.timeout)
+                if msg is None:
+                    logger.debug("No Kafka messages found")
+                    empty_polls += 1
+                    continue
+                if msg.error():
+                    if msg.error().code() == KafkaError._PARTITION_EOF:
+                        logger.info(
+                            "Reached end of partition for topic %s [%d] at offset %d"
+                            % (msg.topic(), msg.partition(), msg.offset())
+                        )
+                        break
+                    else:
+                        raise KafkaException(msg.error())
+                try:
+                    empty_polls = 0
+                    messages_consumed += 1
+                    yield msg
+                finally:
+                    consumer.commit(asynchronous=False)
+
+    def generate_file_data(self, msg) -> FileData:
+        msg_content = msg.value().decode("utf8")
+        identifier = f"{msg.topic()}_{msg.partition()}_{msg.offset()}"
+        additional_metadata = {
+            "topic": msg.topic(),
+            "partition": msg.partition(),
+            "offset": msg.offset(),
+            "content": msg_content,
+        }
+        filename = f"{identifier}.txt"
+        return FileData(
+            identifier=identifier,
+            connector_type=self.connector_type,
+            source_identifiers=SourceIdentifiers(
+                filename=filename,
+                fullpath=filename,
+            ),
+            metadata=FileDataSourceMetadata(
+                date_processed=str(time()),
+            ),
+            additional_metadata=additional_metadata,
+            display_name=filename,
+        )
+
+    def run(self) -> Generator[FileData, None, None]:
+        for message in self.generate_messages():
+            yield self.generate_file_data(message)
+
+    async def run_async(self, file_data: FileData, **kwargs: Any) -> download_responses:
+        raise NotImplementedError()
+
+    def precheck(self):
+        try:
+            with self.get_consumer() as consumer:
+                cluster_meta = consumer.list_topics(timeout=self.connection_config.timeout)
+                current_topics = [
+                    topic for topic in cluster_meta.topics if topic != "__consumer_offsets"
+                ]
+                logger.info(f"successfully checked available topics: {current_topics}")
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+
+class KafkaDownloaderConfig(DownloaderConfig):
+    pass
+
+
+@dataclass
+class KafkaDownloader(Downloader):
+    connection_config: KafkaConnectionConfig
+    download_config: KafkaDownloaderConfig = field(default_factory=KafkaDownloaderConfig)
+    connector_type: str = CONNECTOR_TYPE
+    version: Optional[str] = None
+    source_url: Optional[str] = None
+
+    def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
+        source_identifiers = file_data.source_identifiers
+        if source_identifiers is None:
+            raise ValueError("FileData is missing source_identifiers")
+
+        # Build the download path using source_identifiers
+        download_path = Path(self.download_dir) / source_identifiers.relative_path
+        download_path.parent.mkdir(parents=True, exist_ok=True)
+
+        try:
+            content = file_data.additional_metadata["content"]
+            with open(download_path, "w") as file:
+                file.write(content)
+        except Exception as e:
+            logger.error(f"Failed to download file {file_data.identifier}: {e}")
+            raise SourceConnectionNetworkError(f"failed to download file {file_data.identifier}")
+
+        return self.generate_download_response(file_data=file_data, download_path=download_path)
+
+
+kafka_source_entry = SourceRegistryEntry(
+    connection_config=KafkaConnectionConfig,
+    indexer=KafkaIndexer,
+    indexer_config=KafkaIndexerConfig,
+    downloader=KafkaDownloader,
+    downloader_config=KafkaDownloaderConfig,
+)

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

@@ -0,0 +1,75 @@
+import socket
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.processes.connector_registry import SourceRegistryEntry
+from unstructured_ingest.v2.processes.connectors.kafka.kafka import (
+    KafkaAccessConfig,
+    KafkaConnectionConfig,
+    KafkaDownloader,
+    KafkaDownloaderConfig,
+    KafkaIndexer,
+    KafkaIndexerConfig,
+)
+
+if TYPE_CHECKING:
+    pass
+
+CONNECTOR_TYPE = "kafka-local"
+
+
+class LocalKafkaAccessConfig(KafkaAccessConfig):
+    pass
+
+
+class LocalKafkaConnectionConfig(KafkaConnectionConfig):
+    access_config: Secret[LocalKafkaAccessConfig] = Field(
+        default=LocalKafkaAccessConfig(), validate_default=True
+    )
+
+    def get_consumer_configuration(self) -> dict:
+        bootstrap = self.bootstrap_server
+        port = self.port
+
+        conf = {
+            "bootstrap.servers": f"{bootstrap}:{port}",
+            "client.id": socket.gethostname(),
+            "group.id": "default_group_id",
+            "enable.auto.commit": "false",
+            "auto.offset.reset": "earliest",
+            "message.max.bytes": 10485760,
+        }
+        return conf
+
+
+class LocalKafkaIndexerConfig(KafkaIndexerConfig):
+    pass
+
+
+@dataclass
+class LocalKafkaIndexer(KafkaIndexer):
+    connection_config: LocalKafkaConnectionConfig
+    index_config: LocalKafkaIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class LocalKafkaDownloaderConfig(KafkaDownloaderConfig):
+    pass
+
+
+@dataclass
+class LocalKafkaDownloader(KafkaDownloader):
+    connection_config: LocalKafkaConnectionConfig
+    download_config: LocalKafkaDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+kafka_local_source_entry = SourceRegistryEntry(
+    connection_config=LocalKafkaConnectionConfig,
+    indexer=LocalKafkaIndexer,
+    indexer_config=LocalKafkaIndexerConfig,
+    downloader=LocalKafkaDownloader,
+    downloader_config=LocalKafkaDownloaderConfig,
+)

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,
+)