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

unstructured_ingest/v2/pipeline/steps/chunk.py

@@ -6,6 +6,7 @@ from pathlib import Path
 from typing import Callable, Optional, TypedDict
 
 from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.interfaces.file_data import file_data_from_file
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
 from unstructured_ingest.v2.processes.chunker import Chunker
@@ -51,7 +52,7 @@ class ChunkStep(PipelineStep):
         self, fn: Callable, path: str, file_data_path: str, **kwargs
     ) -> ChunkStepResponse:
         path = Path(path)
-        file_data = FileData.from_file(path=file_data_path)
+        file_data = file_data_from_file(path=file_data_path)
         output_filepath = self.get_output_filepath(filename=path)
         if not self.should_chunk(filepath=output_filepath, file_data=file_data):
             logger.debug(f"skipping chunking, output already exists: {output_filepath}")

unstructured_ingest/v2/pipeline/steps/download.py

@@ -8,6 +8,7 @@ from typing import Callable, Optional, TypedDict, TypeVar
 
 from unstructured_ingest.v2.interfaces import FileData, download_responses
 from unstructured_ingest.v2.interfaces.downloader import Downloader
+from unstructured_ingest.v2.interfaces.file_data import file_data_from_file
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
 from unstructured_ingest.v2.utils import serialize_base_model_json
@@ -87,12 +88,12 @@ class DownloadStep(PipelineStep):
                 f"match size of local file: {file_size_bytes}, updating"
             )
             file_data.metadata.filesize_bytes = file_size_bytes
-        logger.debug(f"updating file data with new content: {file_data.to_dict()}")
+        logger.debug(f"updating file data with new content: {file_data.model_dump()}")
         with file_data_path.open("w") as file:
-            json.dump(file_data.to_dict(), file, indent=2)
+            json.dump(file_data.model_dump(), file, indent=2)
 
     async def _run_async(self, fn: Callable, file_data_path: str) -> list[DownloadStepResponse]:
-        file_data = FileData.from_file(path=file_data_path)
+        file_data = file_data_from_file(path=file_data_path)
         download_path = self.process.get_download_path(file_data=file_data)
         if not self.should_download(file_data=file_data, file_data_path=file_data_path):
             logger.debug(f"skipping download, file already exists locally: {download_path}")
@@ -172,7 +173,7 @@ class DownloadStep(PipelineStep):
         filepath = (self.cache_dir / filename).resolve()
         filepath.parent.mkdir(parents=True, exist_ok=True)
         with open(str(filepath), "w") as f:
-            json.dump(file_data.to_dict(), f, indent=2)
+            json.dump(file_data.model_dump(), f, indent=2)
         return str(filepath)
 
     def get_hash(self, extras: Optional[list[str]]) -> str:

unstructured_ingest/v2/pipeline/steps/embed.py

@@ -6,6 +6,7 @@ from pathlib import Path
 from typing import Callable, Optional, TypedDict
 
 from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.interfaces.file_data import file_data_from_file
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
 from unstructured_ingest.v2.processes.embedder import Embedder
@@ -49,7 +50,7 @@ class EmbedStep(PipelineStep):
 
     async def _run_async(self, fn: Callable, path: str, file_data_path: str) -> EmbedStepResponse:
         path = Path(path)
-        file_data = FileData.from_file(path=file_data_path)
+        file_data = file_data_from_file(path=file_data_path)
         output_filepath = self.get_output_filepath(filename=path)
         if not self.should_embed(filepath=output_filepath, file_data=file_data):
             logger.debug(f"skipping embedding, output already exists: {output_filepath}")

unstructured_ingest/v2/pipeline/steps/filter.py

@@ -2,7 +2,7 @@ import asyncio
 from dataclasses import dataclass
 from typing import Callable, Optional
 
-from unstructured_ingest.v2.interfaces.file_data import FileData
+from unstructured_ingest.v2.interfaces.file_data import file_data_from_file
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
 from unstructured_ingest.v2.processes.filter import Filterer
@@ -20,7 +20,7 @@ class FilterStep(PipelineStep):
         logger.info(f"created {self.identifier} with configs: {config}")
 
     async def _run_async(self, fn: Callable, file_data_path: str, **kwargs) -> Optional[dict]:
-        file_data = FileData.from_file(path=file_data_path)
+        file_data = file_data_from_file(path=file_data_path)
         fn_kwargs = {"file_data": file_data}
         if not asyncio.iscoroutinefunction(fn):
             resp = fn(**fn_kwargs)

unstructured_ingest/v2/pipeline/steps/index.py

@@ -37,14 +37,14 @@ class IndexStep(PipelineStep):
     @instrument(span_name=STEP_ID)
     def run(self) -> Generator[str, None, None]:
         for file_data in self.process.run():
-            logger.debug(f"generated file data: {file_data.to_dict()}")
+            logger.debug(f"generated file data: {file_data.model_dump()}")
             try:
                 record_hash = self.get_hash(extras=[file_data.identifier])
                 filename = f"{record_hash}.json"
                 filepath = (self.cache_dir / filename).resolve()
                 filepath.parent.mkdir(parents=True, exist_ok=True)
                 with open(str(filepath), "w") as f:
-                    json.dump(file_data.to_dict(), f, indent=2)
+                    json.dump(file_data.model_dump(), f, indent=2)
                 yield str(filepath)
             except Exception as e:
                 logger.error(f"failed to create index for file data: {file_data}", exc_info=True)
@@ -54,14 +54,14 @@ class IndexStep(PipelineStep):
 
     async def run_async(self) -> AsyncGenerator[str, None]:
         async for file_data in self.process.run_async():
-            logger.debug(f"generated file data: {file_data.to_dict()}")
+            logger.debug(f"generated file data: {file_data.model_dump()}")
             try:
                 record_hash = self.get_hash(extras=[file_data.identifier])
                 filename = f"{record_hash}.json"
                 filepath = (self.cache_dir / filename).resolve()
                 filepath.parent.mkdir(parents=True, exist_ok=True)
                 with open(str(filepath), "w") as f:
-                    json.dump(file_data.to_dict(), f, indent=2)
+                    json.dump(file_data.model_dump(), f, indent=2)
                 yield str(filepath)
             except Exception as e:
                 logger.error(f"failed to create index for file data: {file_data}", exc_info=True)

unstructured_ingest/v2/pipeline/steps/partition.py

@@ -6,6 +6,7 @@ from pathlib import Path
 from typing import Callable, Optional, TypedDict
 
 from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.interfaces.file_data import file_data_from_file
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
 from unstructured_ingest.v2.processes.partitioner import Partitioner
@@ -51,12 +52,12 @@ class PartitionStep(PipelineStep):
         self, fn: Callable, path: str, file_data_path: str
     ) -> Optional[PartitionStepResponse]:
         path = Path(path)
-        file_data = FileData.from_file(path=file_data_path)
+        file_data = file_data_from_file(path=file_data_path)
         output_filepath = self.get_output_filepath(filename=Path(file_data_path))
         if not self.should_partition(filepath=output_filepath, file_data=file_data):
             logger.debug(f"skipping partitioning, output already exists: {output_filepath}")
             return PartitionStepResponse(file_data_path=file_data_path, path=str(output_filepath))
-        fn_kwargs = {"filename": path, "metadata": file_data.metadata.to_dict()}
+        fn_kwargs = {"filename": path, "metadata": file_data.metadata.model_dump()}
         if not asyncio.iscoroutinefunction(fn):
             partitioned_content = fn(**fn_kwargs)
         elif semaphore := self.context.semaphore:

unstructured_ingest/v2/pipeline/steps/stage.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Callable, Optional, TypedDict
 
-from unstructured_ingest.v2.interfaces.file_data import FileData
+from unstructured_ingest.v2.interfaces.file_data import file_data_from_file
 from unstructured_ingest.v2.interfaces.upload_stager import UploadStager
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
@@ -43,7 +43,7 @@ class UploadStageStep(PipelineStep):
         output_filename = f"{self.get_hash(extras=[path.name])}{path.suffix}"
         fn_kwargs = {
             "elements_filepath": path,
-            "file_data": FileData.from_file(path=file_data_path),
+            "file_data": file_data_from_file(path=file_data_path),
             "output_dir": self.cache_dir,
             "output_filename": output_filename,
         }

unstructured_ingest/v2/pipeline/steps/uncompress.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Callable, TypedDict
 
-from unstructured_ingest.v2.interfaces.file_data import FileData
+from unstructured_ingest.v2.interfaces.file_data import file_data_from_file
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
 from unstructured_ingest.v2.processes.uncompress import Uncompressor
@@ -28,7 +28,7 @@ class UncompressStep(PipelineStep):
     async def _run_async(
         self, fn: Callable, path: str, file_data_path: str
     ) -> list[UncompressStepResponse]:
-        file_data = FileData.from_file(path=file_data_path)
+        file_data = file_data_from_file(path=file_data_path)
         fn_kwargs = {"file_data": file_data}
         if not asyncio.iscoroutinefunction(fn):
             new_file_data = fn(**fn_kwargs)

unstructured_ingest/v2/pipeline/steps/upload.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Callable, Optional, TypedDict
 
-from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.interfaces.file_data import file_data_from_file
 from unstructured_ingest.v2.interfaces.uploader import UploadContent
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.pipeline.interfaces import BatchPipelineStep
@@ -41,14 +41,14 @@ class UploadStep(BatchPipelineStep):
     @instrument(span_name=STEP_ID)
     def _run_batch(self, contents: list[UploadStepContent]) -> None:
         upload_contents = [
-            UploadContent(path=Path(c["path"]), file_data=FileData.from_file(c["file_data_path"]))
+            UploadContent(path=Path(c["path"]), file_data=file_data_from_file(c["file_data_path"]))
             for c in contents
         ]
         self.process.run_batch(contents=upload_contents)
 
     async def _run_async(self, path: str, file_data_path: str, fn: Optional[Callable] = None):
         fn = fn or self.process.run_async
-        fn_kwargs = {"path": Path(path), "file_data": FileData.from_file(path=file_data_path)}
+        fn_kwargs = {"path": Path(path), "file_data": file_data_from_file(path=file_data_path)}
         if not asyncio.iscoroutinefunction(fn):
             fn(**fn_kwargs)
         elif semaphore := self.context.semaphore:

unstructured_ingest/v2/processes/connectors/__init__.py

@@ -40,6 +40,8 @@ from .milvus import CONNECTOR_TYPE as MILVUS_CONNECTOR_TYPE
 from .milvus import milvus_destination_entry
 from .mongodb import CONNECTOR_TYPE as MONGODB_CONNECTOR_TYPE
 from .mongodb import mongodb_destination_entry, mongodb_source_entry
+from .neo4j import CONNECTOR_TYPE as NEO4J_CONNECTOR_TYPE
+from .neo4j import neo4j_destination_entry
 from .onedrive import CONNECTOR_TYPE as ONEDRIVE_CONNECTOR_TYPE
 from .onedrive import onedrive_destination_entry, onedrive_source_entry
 from .outlook import CONNECTOR_TYPE as OUTLOOK_CONNECTOR_TYPE
@@ -74,6 +76,7 @@ add_destination_entry(destination_type=LOCAL_CONNECTOR_TYPE, entry=local_destina
 add_source_entry(source_type=ONEDRIVE_CONNECTOR_TYPE, entry=onedrive_source_entry)
 add_destination_entry(destination_type=ONEDRIVE_CONNECTOR_TYPE, entry=onedrive_destination_entry)
 
+add_destination_entry(destination_type=NEO4J_CONNECTOR_TYPE, entry=neo4j_destination_entry)
 
 add_source_entry(source_type=SALESFORCE_CONNECTOR_TYPE, entry=salesforce_source_entry)
 

unstructured_ingest/v2/processes/connectors/astradb.py

@@ -1,13 +1,11 @@
-import copy
 import csv
 import hashlib
-import sys
 from dataclasses import dataclass, field
 from pathlib import Path
 from time import time
 from typing import TYPE_CHECKING, Any, Generator, Optional
 
-from pydantic import Field, Secret
+from pydantic import BaseModel, Field, Secret
 
 from unstructured_ingest import __name__ as integration_name
 from unstructured_ingest.__version__ import __version__ as integration_version
@@ -22,6 +20,8 @@ from unstructured_ingest.utils.string_and_date_utils import truncate_string_byte
 from unstructured_ingest.v2.constants import RECORD_ID_LABEL
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
+    BatchFileData,
+    BatchItem,
     ConnectionConfig,
     Downloader,
     DownloaderConfig,
@@ -30,6 +30,7 @@ from unstructured_ingest.v2.interfaces import (
     FileDataSourceMetadata,
     Indexer,
     IndexerConfig,
+    SourceIdentifiers,
     Uploader,
     UploaderConfig,
     UploadStager,
@@ -53,6 +54,15 @@ CONNECTOR_TYPE = "astradb"
 MAX_CONTENT_PARAM_BYTE_SIZE = 8000
 
 
+class AstraDBAdditionalMetadata(BaseModel):
+    collection_name: str
+    keyspace: Optional[str] = None
+
+
+class AstraDBBatchFileData(BatchFileData):
+    additional_metadata: AstraDBAdditionalMetadata
+
+
 class AstraDBAccessConfig(AccessConfig):
     token: str = Field(description="Astra DB Token with access to the database.")
     api_endpoint: str = Field(description="The API endpoint for the Astra DB.")
@@ -179,9 +189,6 @@ class AstraDBIndexer(Indexer):
 
     def _get_doc_ids(self) -> set[str]:
         """Fetches all document ids in an index"""
-        # Initialize set of ids
-        ids = set()
-
         # Get the collection
         collection = self.get_collection()
 
@@ -194,31 +201,26 @@ class AstraDBIndexer(Indexer):
             astra_db_docs.append(result)
 
         # Create file data for each astra record
-        for astra_record in astra_db_docs:
-            ids.add(astra_record["_id"])
+        ids = sorted([astra_record["_id"] for astra_record in astra_db_docs])
 
-        return ids
+        return set(ids)
 
-    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+    def run(self, **kwargs: Any) -> Generator[AstraDBBatchFileData, None, None]:
         all_ids = self._get_doc_ids()
         ids = list(all_ids)
         id_batches = batch_generator(ids, self.index_config.batch_size)
 
         for batch in id_batches:
-            # Make sure the hash is always a positive number to create identified
-            identified = str(hash(batch) + sys.maxsize + 1)
-            fd = FileData(
-                identifier=identified,
+            fd = AstraDBBatchFileData(
                 connector_type=CONNECTOR_TYPE,
-                doc_type="batch",
                 metadata=FileDataSourceMetadata(
                     date_processed=str(time()),
                 ),
-                additional_metadata={
-                    "ids": list(batch),
-                    "collection_name": self.index_config.collection_name,
-                    "keyspace": self.index_config.keyspace,
-                },
+                additional_metadata=AstraDBAdditionalMetadata(
+                    collection_name=self.index_config.collection_name,
+                    keyspace=self.index_config.keyspace,
+                ),
+                batch_items=[BatchItem(identifier=b) for b in batch],
             )
             yield fd
 
@@ -247,7 +249,9 @@ class AstraDBDownloader(Downloader):
             writer.writerow(astra_result.keys())
             writer.writerow(astra_result.values())
 
-    def generate_download_response(self, result: dict, file_data: FileData) -> DownloadResponse:
+    def generate_download_response(
+        self, result: dict, file_data: AstraDBBatchFileData
+    ) -> DownloadResponse:
         record_id = result["_id"]
         filename_id = self.get_identifier(record_id=record_id)
         filename = f"{filename_id}.csv"  # csv to preserve column info
@@ -255,7 +259,7 @@ class AstraDBDownloader(Downloader):
         logger.debug(f"Downloading results from record {record_id} as csv to {download_path}")
         download_path.parent.mkdir(parents=True, exist_ok=True)
         try:
-            self.write_astra_result_to_csv(astra_result=result, download_path=download_path)
+            self.write_astra_result_to_csv(astra_result=result, download_path=str(download_path))
         except Exception as e:
             logger.error(
                 f"failed to download from record {record_id} to {download_path}: {e}",
@@ -264,14 +268,13 @@ class AstraDBDownloader(Downloader):
             raise SourceConnectionNetworkError(f"failed to download file {file_data.identifier}")
 
         # modify input file_data for download_response
-        copied_file_data = copy.deepcopy(file_data)
-        copied_file_data.identifier = filename
-        copied_file_data.doc_type = "file"
-        copied_file_data.metadata.date_processed = str(time())
-        copied_file_data.metadata.record_locator = {"document_id": record_id}
-        copied_file_data.additional_metadata.pop("ids", None)
+        file_data.source_identifiers = SourceIdentifiers(filename=filename, fullpath=filename)
+        cast_file_data = FileData.cast(file_data=file_data)
+        cast_file_data.identifier = filename
+        cast_file_data.metadata.date_processed = str(time())
+        cast_file_data.metadata.record_locator = {"document_id": record_id}
         return super().generate_download_response(
-            file_data=copied_file_data, download_path=download_path
+            file_data=cast_file_data, download_path=download_path
         )
 
     def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
@@ -279,9 +282,10 @@ class AstraDBDownloader(Downloader):
 
     async def run_async(self, file_data: FileData, **kwargs: Any) -> download_responses:
         # Get metadata from file_data
-        ids: list[str] = file_data.additional_metadata["ids"]
-        collection_name: str = file_data.additional_metadata["collection_name"]
-        keyspace: str = file_data.additional_metadata["keyspace"]
+        astra_file_data = AstraDBBatchFileData.cast(file_data=file_data)
+        ids: list[str] = [item.identifier for item in astra_file_data.batch_items]
+        collection_name: str = astra_file_data.additional_metadata.collection_name
+        keyspace: str = astra_file_data.additional_metadata.keyspace
 
         # Retrieve results from async collection
         download_responses = []
@@ -292,7 +296,7 @@ class AstraDBDownloader(Downloader):
         )
         async for result in async_astra_collection.find({"_id": {"$in": ids}}):
             download_responses.append(
-                self.generate_download_response(result=result, file_data=file_data)
+                self.generate_download_response(result=result, file_data=astra_file_data)
             )
         return download_responses
 

unstructured_ingest/v2/processes/connectors/couchbase.py

@@ -1,5 +1,4 @@
 import hashlib
-import sys
 import time
 from contextlib import contextmanager
 from dataclasses import dataclass, field
@@ -7,7 +6,7 @@ from datetime import timedelta
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Generator, List
 
-from pydantic import Field, Secret
+from pydantic import BaseModel, Field, Secret
 
 from unstructured_ingest.error import (
     DestinationConnectionError,
@@ -18,6 +17,8 @@ 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,
+    BatchFileData,
+    BatchItem,
     ConnectionConfig,
     Downloader,
     DownloaderConfig,
@@ -26,6 +27,7 @@ from unstructured_ingest.v2.interfaces import (
     FileDataSourceMetadata,
     Indexer,
     IndexerConfig,
+    SourceIdentifiers,
     Uploader,
     UploaderConfig,
     UploadStager,
@@ -40,11 +42,20 @@ from unstructured_ingest.v2.processes.connector_registry import (
 
 if TYPE_CHECKING:
     from couchbase.cluster import Cluster
+    from couchbase.collection import Collection
 
 CONNECTOR_TYPE = "couchbase"
 SERVER_API_VERSION = "1"
 
 
+class CouchbaseAdditionalMetadata(BaseModel):
+    bucket: str
+
+
+class CouchbaseBatchFileData(BatchFileData):
+    additional_metadata: CouchbaseAdditionalMetadata
+
+
 class CouchbaseAccessConfig(AccessConfig):
     password: str = Field(description="The password for the Couchbase server")
 
@@ -180,31 +191,21 @@ class CouchbaseIndexer(Indexer):
                 if attempts == max_attempts:
                     raise SourceConnectionError(f"failed to get document ids: {e}")
 
-    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+    def run(self, **kwargs: Any) -> Generator[CouchbaseBatchFileData, None, None]:
         ids = self._get_doc_ids()
-
-        id_batches = [
-            ids[i * self.index_config.batch_size : (i + 1) * self.index_config.batch_size]
-            for i in range(
-                (len(ids) + self.index_config.batch_size - 1) // self.index_config.batch_size
-            )
-        ]
-        for batch in id_batches:
+        for batch in batch_generator(ids, self.index_config.batch_size):
             # Make sure the hash is always a positive number to create identified
-            identified = str(hash(tuple(batch)) + sys.maxsize + 1)
-            yield FileData(
-                identifier=identified,
+            yield CouchbaseBatchFileData(
                 connector_type=CONNECTOR_TYPE,
-                doc_type="batch",
                 metadata=FileDataSourceMetadata(
                     url=f"{self.connection_config.connection_string}/"
                     f"{self.connection_config.bucket}",
                     date_processed=str(time.time()),
                 ),
-                additional_metadata={
-                    "ids": list(batch),
-                    "bucket": self.connection_config.bucket,
-                },
+                additional_metadata=CouchbaseAdditionalMetadata(
+                    bucket=self.connection_config.bucket
+                ),
+                batch_items=[BatchItem(identifier=b) for b in batch],
             )
 
 
@@ -241,7 +242,7 @@ class CouchbaseDownloader(Downloader):
         return concatenated_values
 
     def generate_download_response(
-        self, result: dict, bucket: str, file_data: FileData
+        self, result: dict, bucket: str, file_data: CouchbaseBatchFileData
     ) -> DownloadResponse:
         record_id = result[self.download_config.collection_id]
         filename_id = self.get_identifier(bucket=bucket, record_id=record_id)
@@ -261,28 +262,26 @@ class CouchbaseDownloader(Downloader):
                 exc_info=True,
             )
             raise SourceConnectionNetworkError(f"failed to download file {file_data.identifier}")
-        return DownloadResponse(
-            file_data=FileData(
-                identifier=filename_id,
-                connector_type=CONNECTOR_TYPE,
-                metadata=FileDataSourceMetadata(
-                    version=None,
-                    date_processed=str(time.time()),
-                    record_locator={
-                        "connection_string": self.connection_config.connection_string,
-                        "bucket": bucket,
-                        "scope": self.connection_config.scope,
-                        "collection": self.connection_config.collection,
-                        "document_id": record_id,
-                    },
-                ),
-            ),
-            path=download_path,
+        file_data.source_identifiers = SourceIdentifiers(filename=filename, fullpath=filename)
+        cast_file_data = FileData.cast(file_data=file_data)
+        cast_file_data.identifier = filename_id
+        cast_file_data.metadata.date_processed = str(time.time())
+        cast_file_data.metadata.record_locator = {
+            "connection_string": self.connection_config.connection_string,
+            "bucket": bucket,
+            "scope": self.connection_config.scope,
+            "collection": self.connection_config.collection,
+            "document_id": record_id,
+        }
+        return super().generate_download_response(
+            file_data=cast_file_data,
+            download_path=download_path,
         )
 
     def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
-        bucket_name: str = file_data.additional_metadata["bucket"]
-        ids: list[str] = file_data.additional_metadata["ids"]
+        couchbase_file_data = CouchbaseBatchFileData.cast(file_data=file_data)
+        bucket_name: str = couchbase_file_data.additional_metadata.bucket
+        ids: list[str] = [item.identifier for item in couchbase_file_data.batch_items]
 
         with self.connection_config.get_client() as client:
             bucket = client.bucket(bucket_name)
@@ -292,13 +291,25 @@ class CouchbaseDownloader(Downloader):
             download_resp = self.process_all_doc_ids(ids, collection, bucket_name, file_data)
             return list(download_resp)
 
-    def process_doc_id(self, doc_id, collection, bucket_name, file_data):
+    def process_doc_id(
+        self,
+        doc_id: str,
+        collection: "Collection",
+        bucket_name: str,
+        file_data: CouchbaseBatchFileData,
+    ):
         result = collection.get(doc_id)
         return self.generate_download_response(
             result=result.content_as[dict], bucket=bucket_name, file_data=file_data
         )
 
-    def process_all_doc_ids(self, ids, collection, bucket_name, file_data):
+    def process_all_doc_ids(
+        self,
+        ids: list[str],
+        collection: "Collection",
+        bucket_name: str,
+        file_data: CouchbaseBatchFileData,
+    ):
         for doc_id in ids:
             yield self.process_doc_id(doc_id, collection, bucket_name, file_data)
 

unstructured_ingest/v2/processes/connectors/databricks/volumes.py

@@ -14,6 +14,7 @@ from unstructured_ingest.error import (
 )
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
     ConnectionConfig,
     Downloader,
     DownloaderConfig,
@@ -52,6 +53,10 @@ class DatabricksPathMixin(BaseModel):
         return path
 
 
+class DatabricksVolumesAccessConfig(AccessConfig):
+    token: Optional[str] = Field(default=None, description="Databricks Personal Access Token")
+
+
 class DatabricksVolumesConnectionConfig(ConnectionConfig, ABC):
     host: Optional[str] = Field(
         default=None,

unstructured_ingest/v2/processes/connectors/databricks/volumes_aws.py

@@ -3,12 +3,12 @@ from typing import Optional
 
 from pydantic import Field, Secret
 
-from unstructured_ingest.v2.interfaces import AccessConfig
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
     SourceRegistryEntry,
 )
 from unstructured_ingest.v2.processes.connectors.databricks.volumes import (
+    DatabricksVolumesAccessConfig,
     DatabricksVolumesConnectionConfig,
     DatabricksVolumesDownloader,
     DatabricksVolumesDownloaderConfig,
@@ -21,7 +21,7 @@ from unstructured_ingest.v2.processes.connectors.databricks.volumes import (
 CONNECTOR_TYPE = "databricks_volumes_aws"
 
 
-class DatabricksAWSVolumesAccessConfig(AccessConfig):
+class DatabricksAWSVolumesAccessConfig(DatabricksVolumesAccessConfig):
     account_id: Optional[str] = Field(
         default=None,
         description="The Databricks account ID for the Databricks " "accounts endpoint",

unstructured_ingest/v2/processes/connectors/databricks/volumes_azure.py

@@ -3,12 +3,12 @@ from typing import Optional
 
 from pydantic import Field, Secret
 
-from unstructured_ingest.v2.interfaces import AccessConfig
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
     SourceRegistryEntry,
 )
 from unstructured_ingest.v2.processes.connectors.databricks.volumes import (
+    DatabricksVolumesAccessConfig,
     DatabricksVolumesConnectionConfig,
     DatabricksVolumesDownloader,
     DatabricksVolumesDownloaderConfig,
@@ -21,7 +21,7 @@ from unstructured_ingest.v2.processes.connectors.databricks.volumes import (
 CONNECTOR_TYPE = "databricks_volumes_azure"
 
 
-class DatabricksAzureVolumesAccessConfig(AccessConfig):
+class DatabricksAzureVolumesAccessConfig(DatabricksVolumesAccessConfig):
     account_id: Optional[str] = Field(
         default=None,
         description="The Databricks account ID for the Databricks " "accounts endpoint.",

unstructured_ingest/v2/processes/connectors/databricks/volumes_gcp.py

@@ -3,12 +3,12 @@ from typing import Optional
 
 from pydantic import Field, Secret
 
-from unstructured_ingest.v2.interfaces import AccessConfig
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
     SourceRegistryEntry,
 )
 from unstructured_ingest.v2.processes.connectors.databricks.volumes import (
+    DatabricksVolumesAccessConfig,
     DatabricksVolumesConnectionConfig,
     DatabricksVolumesDownloader,
     DatabricksVolumesDownloaderConfig,
@@ -21,7 +21,7 @@ from unstructured_ingest.v2.processes.connectors.databricks.volumes import (
 CONNECTOR_TYPE = "databricks_volumes_gcp"
 
 
-class DatabricksGoogleVolumesAccessConfig(AccessConfig):
+class DatabricksGoogleVolumesAccessConfig(DatabricksVolumesAccessConfig):
     account_id: Optional[str] = Field(
         default=None,
         description="The Databricks account ID for the Databricks " "accounts endpoint.",