unstructured-ingest 1.0.1__tar.gz → 1.0.2__tar.gz

{unstructured_ingest-1.0.1 → unstructured_ingest-1.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unstructured_ingest
-Version: 1.0.1
+Version: 1.0.2
 Summary: Local ETL data pipeline to get data RAG ready
 Author-email: Unstructured Technologies <devops@unstructuredai.io>
 License-Expression: Apache-2.0
@@ -28,7 +28,7 @@ Provides-Extra: airtable
 Requires-Dist: pandas; extra == 'airtable'
 Requires-Dist: pyairtable; extra == 'airtable'
 Provides-Extra: astradb
-Requires-Dist: astrapy; extra == 'astradb'
+Requires-Dist: astrapy>2.0.0; extra == 'astradb'
 Provides-Extra: azure
 Requires-Dist: adlfs; extra == 'azure'
 Requires-Dist: fsspec; extra == 'azure'

{unstructured_ingest-1.0.1 → unstructured_ingest-1.0.2}/pyproject.toml

@@ -172,7 +172,6 @@ ci = [
     "lancedb<=0.15.0",
     # TODO: versions higher than this are missing the macos wheel
     "pykx==2.5.3",
-    "astrapy<2.0.0"
 ]
 
 [project.scripts]

unstructured_ingest-1.0.2/unstructured_ingest/__version__.py

@@ -0,0 +1 @@
+__version__ = "1.0.2"  # pragma: no cover

{unstructured_ingest-1.0.1 → unstructured_ingest-1.0.2}/unstructured_ingest/processes/connectors/astradb.py

@@ -1,5 +1,7 @@
+import asyncio
 import csv
 import hashlib
+import os
 import re
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -8,7 +10,6 @@ from typing import TYPE_CHECKING, Any, Generator, Optional
 
 from pydantic import BaseModel, Field, Secret
 
-from unstructured_ingest import __name__ as integration_name
 from unstructured_ingest.__version__ import __version__ as integration_version
 from unstructured_ingest.data_types.file_data import (
     BatchFileData,
@@ -83,10 +84,8 @@ class AstraDBConnectionConfig(ConnectionConfig):
 
         # Create a client object to interact with the Astra DB
         # caller_name/version for Astra DB tracking
-        return AstraDBClient(
-            caller_name=integration_name,
-            caller_version=integration_version,
-        )
+        user_agent = os.getenv("UNSTRUCTURED_USER_AGENT", "unstructuredio_oss")
+        return AstraDBClient(callers=[(user_agent, integration_version)])
 
 
 def get_astra_db(
@@ -141,7 +140,7 @@ async def get_async_astra_collection(
     )
 
     # Get async collection from AsyncDatabase
-    async_astra_db_collection = await async_astra_db.get_collection(name=collection_name)
+    async_astra_db_collection = async_astra_db.get_collection(name=collection_name)
     return async_astra_db_collection
 
 
@@ -360,13 +359,22 @@ class AstraDBUploader(Uploader):
     upload_config: AstraDBUploaderConfig
     connector_type: str = CONNECTOR_TYPE
 
+    def is_async(self) -> bool:
+        return True
+
     def init(self, **kwargs: Any) -> None:
         self.create_destination(**kwargs)
 
+    @requires_dependencies(["astrapy"], extras="astradb")
     def precheck(self) -> None:
         try:
             if self.upload_config.collection_name:
-                self.get_collection(collection_name=self.upload_config.collection_name).options()
+                collection = get_astra_collection(
+                    connection_config=self.connection_config,
+                    collection_name=self.upload_config.collection_name,
+                    keyspace=self.upload_config.keyspace,
+                )
+                collection.options()
             else:
                 # check for db connection only if collection name is not provided
                 get_astra_db(
@@ -377,17 +385,7 @@ class AstraDBUploader(Uploader):
             logger.error(f"Failed to validate connection {e}", exc_info=True)
             raise DestinationConnectionError(f"failed to validate connection: {e}")
 
-    @requires_dependencies(["astrapy"], extras="astradb")
-    def get_collection(self, collection_name: Optional[str] = None) -> "AstraDBCollection":
-        return get_astra_collection(
-            connection_config=self.connection_config,
-            collection_name=collection_name or self.upload_config.collection_name,
-            keyspace=self.upload_config.keyspace,
-        )
-
     def _collection_exists(self, collection_name: str):
-        from astrapy.exceptions import CollectionNotFoundException
-
         collection = get_astra_collection(
             connection_config=self.connection_config,
             collection_name=collection_name,
@@ -397,8 +395,10 @@ class AstraDBUploader(Uploader):
         try:
             collection.options()
             return True
-        except CollectionNotFoundException:
-            return False
+        except RuntimeError as e:
+            if "not found" in str(e):
+                return False
+            raise DestinationConnectionError(f"failed to check if astra collection exists : {e}")
         except Exception as e:
             logger.error(f"failed to check if astra collection exists : {e}")
             raise DestinationConnectionError(f"failed to check if astra collection exists : {e}")
@@ -422,6 +422,8 @@ class AstraDBUploader(Uploader):
         self.upload_config.collection_name = collection_name
 
         if not self._collection_exists(collection_name):
+            from astrapy.info import CollectionDefinition
+
             astra_db = get_astra_db(
                 connection_config=self.connection_config, keyspace=self.upload_config.keyspace
             )
@@ -429,44 +431,56 @@ class AstraDBUploader(Uploader):
                 f"creating default astra collection '{collection_name}' with dimension "
                 f"{vector_length} and metric {similarity_metric}"
             )
-            astra_db.create_collection(
-                collection_name,
-                dimension=vector_length,
-                metric=similarity_metric,
+            definition = (
+                CollectionDefinition.builder()
+                .set_vector_dimension(dimension=vector_length)
+                .set_vector_metric(similarity_metric)
+                .build()
             )
+            (astra_db.create_collection(collection_name, definition=definition),)
             return True
         logger.debug(f"collection with name '{collection_name}' already exists, skipping creation")
         return False
 
-    def delete_by_record_id(self, collection: "AstraDBCollection", file_data: FileData):
+    async def delete_by_record_id(self, collection: "AstraDBAsyncCollection", 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)
+        delete_resp = await collection.delete_many(filter=delete_filter)
         logger.debug(
             f"deleted {delete_resp.deleted_count} records from collection {collection.name}"
         )
 
-    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+    async def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
         logger.info(
             f"writing {len(data)} objects to destination "
             f"collection {self.upload_config.collection_name}"
         )
 
         astra_db_batch_size = self.upload_config.batch_size
-        collection = self.get_collection()
-
-        self.delete_by_record_id(collection=collection, file_data=file_data)
+        async_astra_collection = await get_async_astra_collection(
+            connection_config=self.connection_config,
+            collection_name=self.upload_config.collection_name,
+            keyspace=self.upload_config.keyspace,
+        )
 
-        for chunk in batch_generator(data, astra_db_batch_size):
-            collection.insert_many(chunk)
+        await self.delete_by_record_id(collection=async_astra_collection, file_data=file_data)
+        await asyncio.gather(
+            *[
+                async_astra_collection.insert_many(chunk)
+                for chunk in batch_generator(data, astra_db_batch_size)
+            ]
+        )
 
-    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+    async def run_async(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
         data = get_json_data(path=path)
-        self.run_data(data=data, file_data=file_data, **kwargs)
+        await self.run_data(data=data, file_data=file_data)
+
+    def run(self, **kwargs: Any) -> Any:
+        raise NotImplementedError("Use astradb run_async instead")
 
 
 astra_db_source_entry = SourceRegistryEntry(

unstructured_ingest-1.0.1/unstructured_ingest/__version__.py

@@ -1 +0,0 @@
-__version__ = "1.0.1"  # pragma: no cover