unstructured-ingest 0.3.11__py3-none-any.whl → 0.3.12__py3-none-any.whl

unstructured_ingest/embed/octoai.py

@@ -4,7 +4,15 @@ from typing import TYPE_CHECKING
 from pydantic import Field, SecretStr
 
 from unstructured_ingest.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
+from unstructured_ingest.logger import logger
 from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.errors import (
+    ProviderError,
+    QuotaError,
+    RateLimitError,
+    UserAuthError,
+    UserError,
+)
 
 if TYPE_CHECKING:
     from openai import OpenAI
@@ -30,12 +38,45 @@ class OctoAiEmbeddingConfig(EmbeddingConfig):
 class OctoAIEmbeddingEncoder(BaseEmbeddingEncoder):
     config: OctoAiEmbeddingConfig
 
+    def wrap_error(self, e: Exception) -> Exception:
+        # https://platform.openai.com/docs/guides/error-codes/api-errors
+        from openai import APIStatusError
+
+        if not isinstance(e, APIStatusError):
+            logger.error(f"unhandled exception from openai: {e}", exc_info=True)
+            raise e
+        error_code = e.code
+        if 400 <= e.status_code < 500:
+            # user error
+            if e.status_code == 401:
+                return UserAuthError(e.message)
+            if e.status_code == 429:
+                # 429 indicates rate limit exceeded and quote exceeded
+                if error_code == "insufficient_quota":
+                    return QuotaError(e.message)
+                else:
+                    return RateLimitError(e.message)
+            return UserError(e.message)
+        if e.status_code >= 500:
+            return ProviderError(e.message)
+        logger.error(f"unhandled exception from openai: {e}", exc_info=True)
+        return e
+
     def embed_query(self, query: str):
-        client = self.config.get_client()
-        response = client.embeddings.create(input=query, model=self.config.embedder_model_name)
+        try:
+            client = self.config.get_client()
+            response = client.embeddings.create(input=query, model=self.config.embedder_model_name)
+        except Exception as e:
+            raise self.wrap_error(e=e)
         return response.data[0].embedding
 
     def embed_documents(self, elements: list[dict]) -> list[dict]:
-        embeddings = [self.embed_query(e.get("text", "")) for e in elements]
+        texts = [e.get("text", "") for e in elements]
+        try:
+            client = self.config.get_client()
+            response = client.embeddings.create(input=texts, model=self.config.embedder_model_name)
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        embeddings = [data.embedding for data in response.data]
         elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
         return elements_with_embeddings

unstructured_ingest/embed/openai.py

@@ -4,7 +4,15 @@ from typing import TYPE_CHECKING
 from pydantic import Field, SecretStr
 
 from unstructured_ingest.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
+from unstructured_ingest.logger import logger
 from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.errors import (
+    ProviderError,
+    QuotaError,
+    RateLimitError,
+    UserAuthError,
+    UserError,
+)
 
 if TYPE_CHECKING:
     from openai import OpenAI
@@ -25,9 +33,37 @@ class OpenAIEmbeddingConfig(EmbeddingConfig):
 class OpenAIEmbeddingEncoder(BaseEmbeddingEncoder):
     config: OpenAIEmbeddingConfig
 
+    def wrap_error(self, e: Exception) -> Exception:
+        # https://platform.openai.com/docs/guides/error-codes/api-errors
+        from openai import APIStatusError
+
+        if not isinstance(e, APIStatusError):
+            logger.error(f"unhandled exception from openai: {e}", exc_info=True)
+            raise e
+        error_code = e.code
+        if 400 <= e.status_code < 500:
+            # user error
+            if e.status_code == 401:
+                return UserAuthError(e.message)
+            if e.status_code == 429:
+                # 429 indicates rate limit exceeded and quote exceeded
+                if error_code == "insufficient_quota":
+                    return QuotaError(e.message)
+                else:
+                    return RateLimitError(e.message)
+            return UserError(e.message)
+        if e.status_code >= 500:
+            return ProviderError(e.message)
+        logger.error(f"unhandled exception from openai: {e}", exc_info=True)
+        return e
+
     def embed_query(self, query: str) -> list[float]:
+
         client = self.config.get_client()
-        response = client.embeddings.create(input=query, model=self.config.embedder_model_name)
+        try:
+            response = client.embeddings.create(input=query, model=self.config.embedder_model_name)
+        except Exception as e:
+            raise self.wrap_error(e=e)
         return response.data[0].embedding
 
     def embed_documents(self, elements: list[dict]) -> list[dict]:

unstructured_ingest/embed/togetherai.py

@@ -4,7 +4,15 @@ from typing import TYPE_CHECKING
 from pydantic import Field, SecretStr
 
 from unstructured_ingest.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
+from unstructured_ingest.logger import logger
 from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.errors import (
+    RateLimitError as CustomRateLimitError,
+)
+from unstructured_ingest.v2.errors import (
+    UserAuthError,
+    UserError,
+)
 
 if TYPE_CHECKING:
     from together import Together
@@ -27,6 +35,20 @@ class TogetherAIEmbeddingConfig(EmbeddingConfig):
 class TogetherAIEmbeddingEncoder(BaseEmbeddingEncoder):
     config: TogetherAIEmbeddingConfig
 
+    def wrap_error(self, e: Exception) -> Exception:
+        # https://docs.together.ai/docs/error-codes
+        from together.error import AuthenticationError, RateLimitError, TogetherException
+
+        if not isinstance(e, TogetherException):
+            logger.error(f"unhandled exception from openai: {e}", exc_info=True)
+            return e
+        message = e.args[0]
+        if isinstance(e, AuthenticationError):
+            return UserAuthError(message)
+        if isinstance(e, RateLimitError):
+            return CustomRateLimitError(message)
+        return UserError(message)
+
     def embed_query(self, query: str) -> list[float]:
         return self._embed_documents(elements=[query])[0]
 
@@ -36,5 +58,10 @@ class TogetherAIEmbeddingEncoder(BaseEmbeddingEncoder):
 
     def _embed_documents(self, elements: list[str]) -> list[list[float]]:
         client = self.config.get_client()
-        outputs = client.embeddings.create(model=self.config.embedder_model_name, input=elements)
+        try:
+            outputs = client.embeddings.create(
+                model=self.config.embedder_model_name, input=elements
+            )
+        except Exception as e:
+            raise self.wrap_error(e=e)
         return [outputs.data[i].embedding for i in range(len(elements))]

unstructured_ingest/embed/voyageai.py

@@ -4,7 +4,16 @@ from typing import TYPE_CHECKING, Optional
 from pydantic import Field, SecretStr
 
 from unstructured_ingest.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
+from unstructured_ingest.logger import logger
 from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.errors import (
+    ProviderError,
+    UserAuthError,
+    UserError,
+)
+from unstructured_ingest.v2.errors import (
+    RateLimitError as CustomRateLimitError,
+)
 
 if TYPE_CHECKING:
     from voyageai import Client as VoyageAIClient
@@ -38,9 +47,32 @@ class VoyageAIEmbeddingConfig(EmbeddingConfig):
 class VoyageAIEmbeddingEncoder(BaseEmbeddingEncoder):
     config: VoyageAIEmbeddingConfig
 
+    def wrap_error(self, e: Exception) -> Exception:
+        # https://docs.voyageai.com/docs/error-codes
+        from voyageai.error import AuthenticationError, RateLimitError, VoyageError
+
+        if not isinstance(e, VoyageError):
+            logger.error(f"unhandled exception from openai: {e}", exc_info=True)
+            raise e
+        http_code = e.http_status
+        message = e.user_message
+        if isinstance(e, AuthenticationError):
+            return UserAuthError(message)
+        if isinstance(e, RateLimitError):
+            return CustomRateLimitError(message)
+        if 400 <= http_code < 500:
+            return UserError(message)
+        if http_code >= 500:
+            return ProviderError(message)
+        logger.error(f"unhandled exception from openai: {e}", exc_info=True)
+        return e
+
     def _embed_documents(self, elements: list[str]) -> list[list[float]]:
         client: VoyageAIClient = self.config.get_client()
-        response = client.embed(texts=elements, model=self.config.embedder_model_name)
+        try:
+            response = client.embed(texts=elements, model=self.config.embedder_model_name)
+        except Exception as e:
+            self.wrap_error(e=e)
         return response.embeddings
 
     def embed_documents(self, elements: list[dict]) -> list[dict]:

unstructured_ingest/v2/errors.py

@@ -0,0 +1,18 @@
+class UserError(Exception):
+    pass
+
+
+class UserAuthError(UserError):
+    pass
+
+
+class RateLimitError(UserError):
+    pass
+
+
+class QuotaError(UserError):
+    pass
+
+
+class ProviderError(Exception):
+    pass

unstructured_ingest/v2/processes/connectors/__init__.py

@@ -48,12 +48,16 @@ from .outlook import CONNECTOR_TYPE as OUTLOOK_CONNECTOR_TYPE
 from .outlook import outlook_source_entry
 from .pinecone import CONNECTOR_TYPE as PINECONE_CONNECTOR_TYPE
 from .pinecone import pinecone_destination_entry
+from .redisdb import CONNECTOR_TYPE as REDIS_CONNECTOR_TYPE
+from .redisdb import redis_destination_entry
 from .salesforce import CONNECTOR_TYPE as SALESFORCE_CONNECTOR_TYPE
 from .salesforce import salesforce_source_entry
 from .sharepoint import CONNECTOR_TYPE as SHAREPOINT_CONNECTOR_TYPE
 from .sharepoint import sharepoint_source_entry
 from .slack import CONNECTOR_TYPE as SLACK_CONNECTOR_TYPE
 from .slack import slack_source_entry
+from .vectara import CONNECTOR_TYPE as VECTARA_CONNECTOR_TYPE
+from .vectara import vectara_destination_entry
 
 add_source_entry(source_type=ASTRA_DB_CONNECTOR_TYPE, entry=astra_db_source_entry)
 add_destination_entry(destination_type=ASTRA_DB_CONNECTOR_TYPE, entry=astra_db_destination_entry)
@@ -101,4 +105,7 @@ add_source_entry(source_type=GITLAB_CONNECTOR_TYPE, entry=gitlab_source_entry)
 
 add_source_entry(source_type=SLACK_CONNECTOR_TYPE, entry=slack_source_entry)
 
+add_destination_entry(destination_type=VECTARA_CONNECTOR_TYPE, entry=vectara_destination_entry)
 add_source_entry(source_type=CONFLUENCE_CONNECTOR_TYPE, entry=confluence_source_entry)
+
+add_destination_entry(destination_type=REDIS_CONNECTOR_TYPE, entry=redis_destination_entry)

unstructured_ingest/v2/processes/connectors/chroma.py

@@ -138,7 +138,6 @@ class ChromaUploader(Uploader):
 
     @DestinationConnectionError.wrap
     def upsert_batch(self, collection, batch):
-
         try:
             # Chroma wants lists even if there is only one element
             # Upserting to prevent duplicates

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

@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING
 
 from pydantic import Field, Secret, SecretStr
 
+from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
     SourceRegistryEntry,
@@ -50,6 +51,7 @@ class CloudKafkaConnectionConfig(KafkaConnectionConfig):
             "sasl.password": access_config.secret.get_secret_value(),
             "sasl.mechanism": "PLAIN",
             "security.protocol": "SASL_SSL",
+            "logger": logger,
         }
 
         return conf
@@ -61,10 +63,11 @@ class CloudKafkaConnectionConfig(KafkaConnectionConfig):
 
         conf = {
             "bootstrap.servers": f"{bootstrap}:{port}",
-            "sasl.username": access_config.kafka_api_key,
-            "sasl.password": access_config.secret,
+            "sasl.username": access_config.kafka_api_key.get_secret_value(),
+            "sasl.password": access_config.secret.get_secret_value(),
             "sasl.mechanism": "PLAIN",
             "security.protocol": "SASL_SSL",
+            "logger": logger,
         }
 
         return conf

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

@@ -170,7 +170,7 @@ class KafkaIndexer(Indexer, ABC):
                 ]
                 if self.index_config.topic not in current_topics:
                     raise SourceConnectionError(
-                        "expected topic {} not detected in cluster: {}".format(
+                        "expected topic '{}' not detected in cluster: '{}'".format(
                             self.index_config.topic, ", ".join(current_topics)
                         )
                     )
@@ -232,6 +232,13 @@ class KafkaUploader(Uploader, ABC):
                     topic for topic in cluster_meta.topics if topic != "__consumer_offsets"
                 ]
                 logger.info(f"successfully checked available topics: {current_topics}")
+                if self.upload_config.topic not in current_topics:
+                    raise DestinationConnectionError(
+                        "expected topic '{}' not detected in cluster: '{}'".format(
+                            self.upload_config.topic, ", ".join(current_topics)
+                        )
+                    )
+
         except Exception as e:
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise DestinationConnectionError(f"failed to validate connection: {e}")
@@ -243,8 +250,10 @@ class KafkaUploader(Uploader, ABC):
         failed_producer = False
 
         def acked(err, msg):
+            nonlocal failed_producer
             if err is not None:
-                logger.error("Failed to deliver message: %s: %s" % (str(msg), str(err)))
+                failed_producer = True
+                logger.error("Failed to deliver kafka message: %s: %s" % (str(msg), str(err)))
 
         for element in elements:
             producer.produce(
@@ -253,7 +262,9 @@ class KafkaUploader(Uploader, ABC):
                 callback=acked,
             )
 
-        producer.flush(timeout=self.upload_config.timeout)
+        while producer_len := len(producer):
+            logger.debug(f"another iteration of kafka producer flush. Queue length: {producer_len}")
+            producer.flush(timeout=self.upload_config.timeout)
         if failed_producer:
             raise KafkaException("failed to produce all messages in batch")
 

unstructured_ingest/v2/processes/connectors/milvus.py

@@ -156,11 +156,18 @@ class MilvusUploader(Uploader):
 
     @DestinationConnectionError.wrap
     def precheck(self):
-        with self.get_client() as client:
-            if not client.has_collection(self.upload_config.collection_name):
-                raise DestinationConnectionError(
-                    f"Collection '{self.upload_config.collection_name}' does not exist"
-                )
+        from pymilvus import MilvusException
+
+        try:
+            with self.get_client() as client:
+                if not client.has_collection(self.upload_config.collection_name):
+                    raise DestinationConnectionError(
+                        f"Collection '{self.upload_config.collection_name}' does not exist"
+                    )
+        except MilvusException as milvus_exception:
+            raise DestinationConnectionError(
+                f"failed to precheck Milvus: {str(milvus_exception.message)}"
+            ) from milvus_exception
 
     @contextmanager
     def get_client(self) -> Generator["MilvusClient", None, None]:
@@ -197,7 +204,9 @@ class MilvusUploader(Uploader):
             try:
                 res = client.insert(collection_name=self.upload_config.collection_name, data=data)
             except MilvusException as milvus_exception:
-                raise WriteError("failed to upload records to milvus") from milvus_exception
+                raise WriteError(
+                    f"failed to upload records to Milvus: {str(milvus_exception.message)}"
+                ) from milvus_exception
             if "err_count" in res and isinstance(res["err_count"], int) and res["err_count"] > 0:
                 err_count = res["err_count"]
                 raise WriteError(f"failed to upload {err_count} docs")

unstructured_ingest/v2/processes/connectors/neo4j.py

@@ -378,6 +378,8 @@ class Neo4jUploader(Uploader):
 
 neo4j_destination_entry = DestinationRegistryEntry(
     connection_config=Neo4jConnectionConfig,
+    upload_stager=Neo4jUploadStager,
+    upload_stager_config=Neo4jUploadStagerConfig,
     uploader=Neo4jUploader,
     uploader_config=Neo4jUploaderConfig,
 )

unstructured_ingest/v2/processes/connectors/onedrive.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
+import asyncio
 import json
 from dataclasses import dataclass
 from pathlib import Path
 from time import time
-from typing import TYPE_CHECKING, Any, Generator, Optional
+from typing import TYPE_CHECKING, Any, AsyncIterator, Generator, Iterator, Optional, TypeVar
 
 from dateutil import parser
 from pydantic import Field, Secret
@@ -100,6 +101,27 @@ class OnedriveIndexerConfig(IndexerConfig):
     recursive: bool = False
 
 
+T = TypeVar("T")
+
+
+def async_iterable_to_sync_iterable(iterator: AsyncIterator[T]) -> Iterator[T]:
+    # This version works on Python 3.9 by manually handling the async iteration.
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    try:
+        while True:
+            try:
+                # Instead of anext(iterator), we directly call __anext__().
+                # __anext__ returns a coroutine that we must run until complete.
+                future = iterator.__anext__()
+                result = loop.run_until_complete(future)
+                yield result
+            except StopAsyncIteration:
+                break
+    finally:
+        loop.close()
+
+
 @dataclass
 class OnedriveIndexer(Indexer):
     connection_config: OnedriveConnectionConfig
@@ -116,17 +138,21 @@ class OnedriveIndexer(Indexer):
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise SourceConnectionError(f"failed to validate connection: {e}")
 
-    def list_objects(self, folder: DriveItem, recursive: bool) -> list["DriveItem"]:
+    def list_objects_sync(self, folder: DriveItem, recursive: bool) -> 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))
+            files.extend(self.list_objects_sync(f, recursive))
         return files
 
-    def get_root(self, client: "GraphClient") -> "DriveItem":
+    async def list_objects(self, folder: "DriveItem", recursive: bool) -> list["DriveItem"]:
+        return await asyncio.to_thread(self.list_objects_sync, folder, recursive)
+
+    def get_root_sync(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()
@@ -134,7 +160,10 @@ class OnedriveIndexer(Indexer):
                 raise ValueError(f"Unable to find directory, given: {fpath}")
         return root
 
-    def get_properties(self, drive_item: "DriveItem") -> dict:
+    async def get_root(self, client: "GraphClient") -> "DriveItem":
+        return await asyncio.to_thread(self.get_root_sync, client)
+
+    def get_properties_sync(self, drive_item: "DriveItem") -> dict:
         properties = drive_item.properties
         filtered_properties = {}
         for k, v in properties.items():
@@ -145,7 +174,10 @@ class OnedriveIndexer(Indexer):
                 pass
         return filtered_properties
 
-    def drive_item_to_file_data(self, drive_item: "DriveItem") -> FileData:
+    async def get_properties(self, drive_item: "DriveItem") -> dict:
+        return await asyncio.to_thread(self.get_properties_sync, drive_item)
+
+    def drive_item_to_file_data_sync(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
@@ -176,17 +208,34 @@ class OnedriveIndexer(Indexer):
                     "server_relative_path": server_path,
                 },
             ),
-            additional_metadata=self.get_properties(drive_item=drive_item),
+            additional_metadata=self.get_properties_sync(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)
+    async def drive_item_to_file_data(self, drive_item: "DriveItem") -> FileData:
+        # Offload the file data creation if it's not guaranteed async
+        return await asyncio.to_thread(self.drive_item_to_file_data_sync, drive_item)
+
+    async def _run_async(self, **kwargs: Any) -> AsyncIterator[FileData]:
+        token_resp = await asyncio.to_thread(self.connection_config.get_token)
+        if "error" in token_resp:
+            raise SourceConnectionError(
+                f"[{CONNECTOR_TYPE}]: {token_resp['error']} ({token_resp.get('error_description')})"
+            )
+
+        client = await asyncio.to_thread(self.connection_config.get_client)
+        root = await self.get_root(client=client)
+        drive_items = await 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)
+            file_data = await self.drive_item_to_file_data(drive_item=drive_item)
             yield file_data
 
+    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+        # Convert the async generator to a sync generator without loading all data into memory
+        async_gen = self._run_async(**kwargs)
+        for item in async_iterable_to_sync_iterable(async_gen):
+            yield item
+
 
 class OnedriveDownloaderConfig(DownloaderConfig):
     pass
@@ -220,19 +269,24 @@ class OnedriveDownloader(Downloader):
 
     @SourceConnectionError.wrap
     def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        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 self.generate_download_response(file_data=file_data, download_path=download_path)
+        try:
+            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 self.generate_download_response(file_data=file_data, download_path=download_path)
+        except Exception as e:
+            logger.error(f"[{CONNECTOR_TYPE}] Exception during downloading: {e}", exc_info=True)
+            # Re-raise to see full stack trace locally
+            raise
 
 
 class OnedriveUploaderConfig(UploaderConfig):

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

@@ -128,7 +128,6 @@ class QdrantUploader(Uploader, ABC):
         file_data: FileData,
         **kwargs: Any,
     ) -> None:
-
         batches = list(batch_generator(data, batch_size=self.upload_config.batch_size))
         logger.debug(
             "Elements split into %i batches of size %i.",