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

unstructured_ingest/v2/processes/connectors/airtable.py

@@ -12,11 +12,11 @@ from unstructured_ingest.v2.interfaces import (
     ConnectionConfig,
     Downloader,
     DownloaderConfig,
+    DownloadResponse,
     FileData,
     Indexer,
     IndexerConfig,
     SourceIdentifiers,
-    download_responses,
 )
 from unstructured_ingest.v2.processes.connector_registry import (
     SourceRegistryEntry,
@@ -214,7 +214,7 @@ class AirtableDownloader(Downloader):
         row_dict.update(table_row["fields"])
         return row_dict
 
-    def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
+    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
         table_meta = AirtableTableMeta.model_validate(file_data.additional_metadata)
         table_contents = self.get_table_contents(table_meta=table_meta)
         df = pandas.DataFrame.from_dict(

unstructured_ingest/v2/processes/connectors/astradb.py

@@ -19,6 +19,7 @@ from unstructured_ingest.error import (
 )
 from unstructured_ingest.utils.data_prep import batch_generator
 from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.constants import RECORD_ID_LABEL
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
     ConnectionConfig,
@@ -129,11 +130,6 @@ class AstraDBIndexerConfig(IndexerConfig):
         "numbers, and underscores."
     )
     keyspace: Optional[str] = Field(default=None, description="The Astra DB connection keyspace.")
-    namespace: Optional[str] = Field(
-        default=None,
-        description="The Astra DB connection namespace.",
-        deprecated="Please use 'keyspace' instead.",
-    )
     batch_size: int = Field(default=20, description="Number of records per batch")
 
 
@@ -147,21 +143,17 @@ class AstraDBUploaderConfig(UploaderConfig):
         "Note that the collection name must only include letters, "
         "numbers, and underscores."
     )
-    embedding_dimension: int = Field(
-        default=384, description="The dimensionality of the embeddings"
-    )
     keyspace: Optional[str] = Field(default=None, description="The Astra DB connection keyspace.")
-    namespace: Optional[str] = Field(
-        default=None,
-        description="The Astra DB connection namespace.",
-        deprecated="Please use 'keyspace' instead.",
-    )
     requested_indexing_policy: Optional[dict[str, Any]] = Field(
         default=None,
         description="The indexing policy to use for the collection.",
         examples=['{"deny": ["metadata"]}'],
     )
     batch_size: int = Field(default=20, description="Number of records per batch")
+    record_id_key: str = Field(
+        default=RECORD_ID_LABEL,
+        description="searchable key to find entries for the same record on previous runs",
+    )
 
 
 @dataclass
@@ -173,12 +165,12 @@ class AstraDBIndexer(Indexer):
         return get_astra_collection(
             connection_config=self.connection_config,
             collection_name=self.index_config.collection_name,
-            keyspace=self.index_config.keyspace or self.index_config.namespace,
+            keyspace=self.index_config.keyspace,
         )
 
     def precheck(self) -> None:
         try:
-            self.get_collection()
+            self.get_collection().options()
         except Exception as e:
             logger.error(f"Failed to validate connection {e}", exc_info=True)
             raise SourceConnectionError(f"failed to validate connection: {e}")
@@ -223,7 +215,7 @@ class AstraDBIndexer(Indexer):
                 additional_metadata={
                     "ids": list(batch),
                     "collection_name": self.index_config.collection_name,
-                    "keyspace": self.index_config.keyspace or self.index_config.namespace,
+                    "keyspace": self.index_config.keyspace,
                 },
             )
             yield fd
@@ -309,10 +301,11 @@ class AstraDBUploadStager(UploadStager):
         default_factory=lambda: AstraDBUploadStagerConfig()
     )
 
-    def conform_dict(self, element_dict: dict) -> dict:
+    def conform_dict(self, element_dict: dict, file_data: FileData) -> dict:
         return {
             "$vector": element_dict.pop("embeddings", None),
             "content": element_dict.pop("text", None),
+            RECORD_ID_LABEL: file_data.identifier,
             "metadata": element_dict,
         }
 
@@ -328,10 +321,15 @@ class AstraDBUploadStager(UploadStager):
             elements_contents = json.load(elements_file)
         conformed_elements = []
         for element in elements_contents:
-            conformed_elements.append(self.conform_dict(element_dict=element))
-        output_path = Path(output_dir) / Path(f"{output_filename}.json")
+            conformed_elements.append(self.conform_dict(element_dict=element, file_data=file_data))
+        output_filename_path = Path(output_filename)
+        if output_filename_path.suffix == ".json":
+            output_path = Path(output_dir) / output_filename_path
+        else:
+            output_path = Path(output_dir) / output_filename_path.with_suffix(".json")
+        output_path.parent.mkdir(parents=True, exist_ok=True)
         with open(output_path, "w") as output_file:
-            json.dump(conformed_elements, output_file)
+            json.dump(conformed_elements, output_file, indent=2)
         return output_path
 
 
@@ -346,8 +344,8 @@ class AstraDBUploader(Uploader):
             get_astra_collection(
                 connection_config=self.connection_config,
                 collection_name=self.upload_config.collection_name,
-                keyspace=self.upload_config.keyspace or self.upload_config.namespace,
-            )
+                keyspace=self.upload_config.keyspace,
+            ).options()
         except Exception as e:
             logger.error(f"Failed to validate connection {e}", exc_info=True)
             raise DestinationConnectionError(f"failed to validate connection: {e}")
@@ -357,7 +355,19 @@ class AstraDBUploader(Uploader):
         return get_astra_collection(
             connection_config=self.connection_config,
             collection_name=self.upload_config.collection_name,
-            keyspace=self.upload_config.keyspace or self.upload_config.namespace,
+            keyspace=self.upload_config.keyspace,
+        )
+
+    def delete_by_record_id(self, collection: "AstraDBCollection", file_data: FileData):
+        logger.debug(
+            f"deleting records from collection {collection.name} "
+            f"with {self.upload_config.record_id_key} "
+            f"set to {file_data.identifier}"
+        )
+        delete_filter = {self.upload_config.record_id_key: {"$eq": file_data.identifier}}
+        delete_resp = collection.delete_many(filter=delete_filter)
+        logger.debug(
+            f"deleted {delete_resp.deleted_count} records from collection {collection.name}"
         )
 
     def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
@@ -371,6 +381,8 @@ class AstraDBUploader(Uploader):
         astra_db_batch_size = self.upload_config.batch_size
         collection = self.get_collection()
 
+        self.delete_by_record_id(collection=collection, file_data=file_data)
+
         for chunk in batch_generator(elements_dict, astra_db_batch_size):
             collection.insert_many(chunk)
 

unstructured_ingest/v2/processes/connectors/{azure_cognitive_search.py → azure_ai_search.py}

@@ -9,6 +9,7 @@ from pydantic import Field, Secret
 from unstructured_ingest.error import DestinationConnectionError, WriteError
 from unstructured_ingest.utils.data_prep import batch_generator
 from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.constants import RECORD_ID_LABEL
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
     ConnectionConfig,
@@ -26,18 +27,18 @@ from unstructured_ingest.v2.processes.connectors.utils import parse_datetime
 
 if TYPE_CHECKING:
     from azure.search.documents import SearchClient
+    from azure.search.documents.indexes import SearchIndexClient
 
+CONNECTOR_TYPE = "azure_ai_search"
 
-CONNECTOR_TYPE = "azure_cognitive_search"
 
-
-class AzureCognitiveSearchAccessConfig(AccessConfig):
-    azure_cognitive_search_key: str = Field(
+class AzureAISearchAccessConfig(AccessConfig):
+    azure_ai_search_key: str = Field(
         alias="key", description="Credential that is used for authenticating to an Azure service"
     )
 
 
-class AzureCognitiveSearchConnectionConfig(ConnectionConfig):
+class AzureAISearchConnectionConfig(ConnectionConfig):
     endpoint: str = Field(
         description="The URL endpoint of an Azure AI (Cognitive) search service. "
         "In the form of https://{{service_name}}.search.windows.net"
@@ -45,10 +46,10 @@ class AzureCognitiveSearchConnectionConfig(ConnectionConfig):
     index: str = Field(
         description="The name of the Azure AI (Cognitive) Search index to connect to."
     )
-    access_config: Secret[AzureCognitiveSearchAccessConfig]
+    access_config: Secret[AzureAISearchAccessConfig]
 
-    @requires_dependencies(["azure.search", "azure.core"], extras="azure-cognitive-search")
-    def generate_client(self) -> "SearchClient":
+    @requires_dependencies(["azure.search", "azure.core"], extras="azure-ai-search")
+    def get_search_client(self) -> "SearchClient":
         from azure.core.credentials import AzureKeyCredential
         from azure.search.documents import SearchClient
 
@@ -56,27 +57,43 @@ class AzureCognitiveSearchConnectionConfig(ConnectionConfig):
             endpoint=self.endpoint,
             index_name=self.index,
             credential=AzureKeyCredential(
-                self.access_config.get_secret_value().azure_cognitive_search_key
+                self.access_config.get_secret_value().azure_ai_search_key
+            ),
+        )
+
+    @requires_dependencies(["azure.search", "azure.core"], extras="azure-ai-search")
+    def get_search_index_client(self) -> "SearchIndexClient":
+        from azure.core.credentials import AzureKeyCredential
+        from azure.search.documents.indexes import SearchIndexClient
+
+        return SearchIndexClient(
+            endpoint=self.endpoint,
+            credential=AzureKeyCredential(
+                self.access_config.get_secret_value().azure_ai_search_key
             ),
         )
 
 
-class AzureCognitiveSearchUploadStagerConfig(UploadStagerConfig):
+class AzureAISearchUploadStagerConfig(UploadStagerConfig):
     pass
 
 
-class AzureCognitiveSearchUploaderConfig(UploaderConfig):
+class AzureAISearchUploaderConfig(UploaderConfig):
     batch_size: int = Field(default=100, description="Number of records per batch")
+    record_id_key: str = Field(
+        default=RECORD_ID_LABEL,
+        description="searchable key to find entries for the same record on previous runs",
+    )
 
 
 @dataclass
-class AzureCognitiveSearchUploadStager(UploadStager):
-    upload_stager_config: AzureCognitiveSearchUploadStagerConfig = field(
-        default_factory=lambda: AzureCognitiveSearchUploadStagerConfig()
+class AzureAISearchUploadStager(UploadStager):
+    upload_stager_config: AzureAISearchUploadStagerConfig = field(
+        default_factory=lambda: AzureAISearchUploadStagerConfig()
     )
 
     @staticmethod
-    def conform_dict(data: dict) -> dict:
+    def conform_dict(data: dict, file_data: FileData) -> dict:
         """
         updates the dictionary that is from each Element being converted into a dict/json
         into a dictionary that conforms to the schema expected by the
@@ -84,6 +101,7 @@ class AzureCognitiveSearchUploadStager(UploadStager):
         """
 
         data["id"] = str(uuid.uuid4())
+        data[RECORD_ID_LABEL] = file_data.identifier
 
         if points := data.get("metadata", {}).get("coordinates", {}).get("points"):
             data["metadata"]["coordinates"]["points"] = json.dumps(points)
@@ -124,6 +142,7 @@ class AzureCognitiveSearchUploadStager(UploadStager):
 
     def run(
         self,
+        file_data: FileData,
         elements_filepath: Path,
         output_dir: Path,
         output_filename: str,
@@ -132,23 +151,63 @@ class AzureCognitiveSearchUploadStager(UploadStager):
         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]
+        conformed_elements = [
+            self.conform_dict(data=element, file_data=file_data) for element in elements_contents
+        ]
 
+        if Path(output_filename).suffix != ".json":
+            output_filename = f"{output_filename}.json"
+        else:
+            output_filename = f"{Path(output_filename).stem}.json"
         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)
+            json.dump(conformed_elements, output_file, indent=2)
         return output_path
 
 
 @dataclass
-class AzureCognitiveSearchUploader(Uploader):
-    upload_config: AzureCognitiveSearchUploaderConfig
-    connection_config: AzureCognitiveSearchConnectionConfig
+class AzureAISearchUploader(Uploader):
+    upload_config: AzureAISearchUploaderConfig
+    connection_config: AzureAISearchConnectionConfig
     connector_type: str = CONNECTOR_TYPE
 
+    def query_docs(self, record_id: str, index_key: str) -> list[str]:
+        client = self.connection_config.get_search_client()
+        results = list(client.search(filter=f"record_id eq '{record_id}'", select=[index_key]))
+        return [result[index_key] for result in results]
+
+    def delete_by_record_id(self, file_data: FileData, index_key: str) -> None:
+        logger.debug(
+            f"deleting any content with metadata "
+            f"{self.upload_config.record_id_key}={file_data.identifier} "
+            f"from azure cognitive search index: {self.connection_config.index}"
+        )
+        doc_ids_to_delete = self.query_docs(record_id=file_data.identifier, index_key=index_key)
+        if not doc_ids_to_delete:
+            return
+        client: SearchClient = self.connection_config.get_search_client()
+        results = client.delete_documents(
+            documents=[{index_key: doc_id} for doc_id in doc_ids_to_delete]
+        )
+        errors = []
+        success = []
+        for result in results:
+            if result.succeeded:
+                success.append(result)
+            else:
+                errors.append(result)
+        logger.debug(f"results: {len(success)} successes, {len(errors)} failures")
+        if errors:
+            raise WriteError(
+                ", ".join(
+                    [f"[{error.status_code}] {error.error_message}" for error in errors],
+                ),
+            )
+
     @DestinationConnectionError.wrap
-    @requires_dependencies(["azure"], extras="azure-cognitive-search")
-    def write_dict(self, *args, elements_dict: list[dict[str, Any]], **kwargs) -> None:
+    @requires_dependencies(["azure"], extras="azure-ai-search")
+    def write_dict(self, elements_dict: list[dict[str, Any]]) -> None:
         import azure.core.exceptions
 
         logger.info(
@@ -156,7 +215,7 @@ class AzureCognitiveSearchUploader(Uploader):
             f"index at {self.connection_config.index}",
         )
         try:
-            results = self.connection_config.generate_client().upload_documents(
+            results = self.connection_config.get_search_client().upload_documents(
                 documents=elements_dict
             )
 
@@ -174,24 +233,42 @@ class AzureCognitiveSearchUploader(Uploader):
             raise WriteError(
                 ", ".join(
                     [
-                        f"{error.azure_cognitive_search_key}: "
+                        f"{error.azure_ai_search_key}: "
                         f"[{error.status_code}] {error.error_message}"
                         for error in errors
                     ],
                 ),
             )
 
+    def can_delete(self) -> bool:
+        search_index_client = self.connection_config.get_search_index_client()
+        index = search_index_client.get_index(name=self.connection_config.index)
+        index_fields = index.fields
+        record_id_fields = [
+            field for field in index_fields if field.name == self.upload_config.record_id_key
+        ]
+        if not record_id_fields:
+            return False
+        record_id_field = record_id_fields[0]
+        return record_id_field.filterable
+
+    def get_index_key(self) -> str:
+        search_index_client = self.connection_config.get_search_index_client()
+        index = search_index_client.get_index(name=self.connection_config.index)
+        index_fields = index.fields
+        key_fields = [field for field in index_fields if field.key]
+        if not key_fields:
+            raise ValueError("no key field found in index fields")
+        return key_fields[0].name
+
     def precheck(self) -> None:
         try:
-            client = self.connection_config.generate_client()
+            client = self.connection_config.get_search_client()
             client.get_document_count()
         except Exception as e:
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise DestinationConnectionError(f"failed to validate connection: {e}")
 
-    def write_dict_wrapper(self, elements_dict):
-        return self.write_dict(elements_dict=elements_dict)
-
     def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
         with path.open("r") as file:
             elements_dict = json.load(file)
@@ -201,17 +278,21 @@ class AzureCognitiveSearchUploader(Uploader):
             f" index at {str(self.connection_config.index)}"
             f" with batch size {str(self.upload_config.batch_size)}"
         )
+        if self.can_delete():
+            index_key = self.get_index_key()
+            self.delete_by_record_id(file_data=file_data, index_key=index_key)
+        else:
+            logger.warning("criteria for deleting previous content not met, skipping")
 
         batch_size = self.upload_config.batch_size
-
         for chunk in batch_generator(elements_dict, batch_size):
             self.write_dict(elements_dict=chunk)  # noqa: E203
 
 
-azure_cognitive_search_destination_entry = DestinationRegistryEntry(
-    connection_config=AzureCognitiveSearchConnectionConfig,
-    uploader=AzureCognitiveSearchUploader,
-    uploader_config=AzureCognitiveSearchUploaderConfig,
-    upload_stager=AzureCognitiveSearchUploadStager,
-    upload_stager_config=AzureCognitiveSearchUploadStagerConfig,
+azure_ai_search_destination_entry = DestinationRegistryEntry(
+    connection_config=AzureAISearchConnectionConfig,
+    uploader=AzureAISearchUploader,
+    uploader_config=AzureAISearchUploaderConfig,
+    upload_stager=AzureAISearchUploadStager,
+    upload_stager_config=AzureAISearchUploadStagerConfig,
 )

unstructured_ingest/v2/processes/connectors/confluence.py

@@ -11,12 +11,12 @@ from unstructured_ingest.v2.interfaces import (
     ConnectionConfig,
     Downloader,
     DownloaderConfig,
+    DownloadResponse,
     FileData,
     FileDataSourceMetadata,
     Indexer,
     IndexerConfig,
     SourceIdentifiers,
-    download_responses,
 )
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import (
@@ -154,7 +154,7 @@ class ConfluenceDownloader(Downloader):
     download_config: ConfluenceDownloaderConfig = field(default_factory=ConfluenceDownloaderConfig)
     connector_type: str = CONNECTOR_TYPE
 
-    def run(self, file_data: FileData, **kwargs) -> download_responses:
+    def run(self, file_data: FileData, **kwargs) -> DownloadResponse:
         doc_id = file_data.identifier
         try:
             client = self.connection_config.get_client()

unstructured_ingest/v2/processes/connectors/couchbase.py

@@ -205,6 +205,7 @@ class CouchbaseIndexer(Indexer):
             yield FileData(
                 identifier=identified,
                 connector_type=CONNECTOR_TYPE,
+                doc_type="batch",
                 metadata=FileDataSourceMetadata(
                     url=f"{self.connection_config.connection_string}/"
                     f"{self.connection_config.bucket}",

unstructured_ingest/v2/processes/connectors/delta_table.py

@@ -1,9 +1,11 @@
 import json
 import os
+import traceback
 from dataclasses import dataclass, field
-from multiprocessing import Process
+from multiprocessing import Process, Queue
 from pathlib import Path
 from typing import Any, Optional
+from urllib.parse import urlparse
 
 import pandas as pd
 from pydantic import Field, Secret
@@ -26,6 +28,15 @@ from unstructured_ingest.v2.processes.connector_registry import DestinationRegis
 CONNECTOR_TYPE = "delta_table"
 
 
+def write_deltalake_with_error_handling(queue, **kwargs):
+    from deltalake.writer import write_deltalake
+
+    try:
+        write_deltalake(**kwargs)
+    except Exception:
+        queue.put(traceback.format_exc())
+
+
 class DeltaTableAccessConfig(AccessConfig):
     aws_access_key_id: Optional[str] = Field(default=None, description="AWS Access Key Id")
     aws_secret_access_key: Optional[str] = Field(default=None, description="AWS Secret Access Key")
@@ -94,7 +105,7 @@ class DeltaTableUploader(Uploader):
     connection_config: DeltaTableConnectionConfig
     connector_type: str = CONNECTOR_TYPE
 
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
+    @requires_dependencies(["boto3"], extras="delta-table")
     def precheck(self):
         secrets = self.connection_config.access_config.get_secret_value()
         if (
@@ -102,13 +113,24 @@ class DeltaTableUploader(Uploader):
             and secrets.aws_access_key_id
             and secrets.aws_secret_access_key
         ):
-            from fsspec import get_filesystem_class
+            from boto3 import client
+
+            url = urlparse(self.connection_config.table_uri)
+            bucket_name = url.netloc
+            dir_path = url.path.lstrip("/")
 
             try:
-                fs = get_filesystem_class("s3")(
-                    key=secrets.aws_access_key_id, secret=secrets.aws_secret_access_key
+                s3_client = client(
+                    "s3",
+                    aws_access_key_id=secrets.aws_access_key_id,
+                    aws_secret_access_key=secrets.aws_secret_access_key,
                 )
-                fs.write_bytes(path=self.connection_config.table_uri, value=b"")
+                s3_client.put_object(Bucket=bucket_name, Key=dir_path, Body=b"")
+
+                response = s3_client.get_bucket_location(Bucket=bucket_name)
+
+                if self.connection_config.aws_region != response.get("LocationConstraint"):
+                    raise ValueError("Wrong AWS Region was provided.")
 
             except Exception as e:
                 logger.error(f"failed to validate connection: {e}", exc_info=True)
@@ -145,7 +167,6 @@ class DeltaTableUploader(Uploader):
 
     @requires_dependencies(["deltalake"], extras="delta-table")
     def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        from deltalake.writer import write_deltalake
 
         df = self.read_dataframe(path)
         updated_upload_path = os.path.join(
@@ -164,17 +185,24 @@ class DeltaTableUploader(Uploader):
             "mode": "overwrite",
             "storage_options": storage_options,
         }
+        queue = Queue()
         # NOTE: deltalake writer on Linux sometimes can finish but still trigger a SIGABRT and cause
         # ingest to fail, even though all tasks are completed normally. Putting the writer into a
         # process mitigates this issue by ensuring python interpreter waits properly for deltalake's
         # rust backend to finish
         writer = Process(
-            target=write_deltalake,
-            kwargs=writer_kwargs,
+            target=write_deltalake_with_error_handling,
+            kwargs={"queue": queue, **writer_kwargs},
         )
         writer.start()
         writer.join()
 
+        # Check if the queue has any exception message
+        if not queue.empty():
+            error_message = queue.get()
+            logger.error(f"Exception occurred in write_deltalake: {error_message}")
+            raise RuntimeError(f"Error in write_deltalake: {error_message}")
+
 
 delta_table_destination_entry = DestinationRegistryEntry(
     connection_config=DeltaTableConnectionConfig,

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

@@ -0,0 +1,19 @@
+from unstructured_ingest.v2.processes.connector_registry import (
+    add_destination_entry,
+    add_source_entry,
+)
+
+from .elasticsearch import CONNECTOR_TYPE as ELASTICSEARCH_CONNECTOR_TYPE
+from .elasticsearch import elasticsearch_destination_entry, elasticsearch_source_entry
+from .opensearch import CONNECTOR_TYPE as OPENSEARCH_CONNECTOR_TYPE
+from .opensearch import opensearch_destination_entry, opensearch_source_entry
+
+add_source_entry(source_type=ELASTICSEARCH_CONNECTOR_TYPE, entry=elasticsearch_source_entry)
+add_destination_entry(
+    destination_type=ELASTICSEARCH_CONNECTOR_TYPE, entry=elasticsearch_destination_entry
+)
+
+add_source_entry(source_type=OPENSEARCH_CONNECTOR_TYPE, entry=opensearch_source_entry)
+add_destination_entry(
+    destination_type=OPENSEARCH_CONNECTOR_TYPE, entry=opensearch_destination_entry
+)