unstructured-ingest 0.3.13__py3-none-any.whl

unstructured_ingest/connector/astradb.py

@@ -0,0 +1,267 @@
+import copy
+import typing as t
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from unstructured_ingest import __name__ as integration_name
+from unstructured_ingest.__version__ import __version__ as integration_version
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.enhanced_dataclass.core import _asdict
+from unstructured_ingest.error import DestinationConnectionError, SourceConnectionError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    BaseConnectorConfig,
+    BaseDestinationConnector,
+    BaseSingleIngestDoc,
+    BaseSourceConnector,
+    IngestDocCleanupMixin,
+    SourceConnectorCleanupMixin,
+    SourceMetadata,
+    WriteConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.data_prep import batch_generator, flatten_dict
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if t.TYPE_CHECKING:
+    from astrapy import Collection as AstraDBCollection
+    from astrapy import Database as AstraDB
+
+NON_INDEXED_FIELDS = ["metadata._node_content", "content"]
+
+
+@dataclass
+class AstraDBAccessConfig(AccessConfig):
+    token: str = enhanced_field(sensitive=True)
+    api_endpoint: str = enhanced_field(sensitive=True)
+
+
+@dataclass
+class SimpleAstraDBConfig(BaseConnectorConfig):
+    access_config: AstraDBAccessConfig
+    collection_name: str
+    keyspace: t.Optional[str] = None
+    namespace: t.Optional[str] = None
+
+
+@dataclass
+class AstraDBIngestDoc(IngestDocCleanupMixin, BaseSingleIngestDoc):
+    connector_config: SimpleAstraDBConfig
+    metadata: t.Dict[str, str] = field(default_factory=dict)
+    registry_name: str = "astradb"
+
+    @property
+    def filename(self):
+        return (
+            Path(self.read_config.download_dir)
+            / self.connector_config.collection_name
+            / f"{self.metadata['_id']}.txt"
+        ).resolve()
+
+    @property
+    def _output_filename(self):
+        return (
+            Path(self.processor_config.output_dir)
+            / self.connector_config.collection_name
+            / f"{self.metadata['_id']}.json"
+        ).resolve()
+
+    def update_source_metadata(self, **kwargs):
+        if not self.metadata:
+            self.source_metadata = SourceMetadata(
+                exists=False,
+            )
+            return
+        self.source_metadata = SourceMetadata(
+            exists=True,
+        )
+
+    @SourceConnectionError.wrap
+    @requires_dependencies(["astrapy"], extras="astradb")
+    @BaseSingleIngestDoc.skip_if_file_exists
+    def get_file(self):
+        self.filename.parent.mkdir(parents=True, exist_ok=True)
+
+        flattened_dict = flatten_dict(dictionary=self.metadata)
+        str_values = [str(value) for value in flattened_dict.values()]
+        concatenated_values = "\n".join(str_values)
+
+        with open(self.filename, "w") as f:
+            f.write(concatenated_values)
+
+
+@dataclass
+class AstraDBSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
+    connector_config: SimpleAstraDBConfig
+    _astra_db: t.Optional["AstraDB"] = field(init=False, default=None)
+    _astra_db_collection: t.Optional["AstraDBCollection"] = field(init=False, default=None)
+
+    @property
+    @requires_dependencies(["astrapy"], extras="astradb")
+    def astra_db_collection(self) -> "AstraDBCollection":
+        if self._astra_db_collection is None:
+            from astrapy import DataAPIClient as AstraDBClient
+
+            # Choose keyspace or deprecated namespace
+            keyspace_param = self.connector_config.keyspace or self.connector_config.namespace
+
+            # Create a client object to interact with the Astra DB
+            # caller_name/version for Astra DB tracking
+            my_client = AstraDBClient(
+                caller_name=integration_name,
+                caller_version=integration_version,
+            )
+
+            # Get the database object
+            self._astra_db = my_client.get_database(
+                api_endpoint=self.connector_config.access_config.api_endpoint,
+                token=self.connector_config.access_config.token,
+                keyspace=keyspace_param,
+            )
+
+            # Create and connect to the newly created collection
+            self._astra_db_collection = self._astra_db.get_collection(
+                name=self.connector_config.collection_name,
+            )
+
+        return self._astra_db_collection  # type: ignore
+
+    @requires_dependencies(["astrapy"], extras="astradb")
+    @SourceConnectionError.wrap  # type: ignore
+    def initialize(self):
+        _ = self.astra_db_collection
+
+    @requires_dependencies(["astrapy"], extras="astradb")
+    def check_connection(self):
+        try:
+            _ = self.astra_db_collection
+        except Exception as e:
+            logger.error(f"Failed to validate connection {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+    @requires_dependencies(["astrapy"], extras="astradb")
+    def get_ingest_docs(self):  # type: ignore
+        # Perform the find operation
+        astra_db_docs_cursor = self.astra_db_collection.find({})
+
+        # Iterate over the cursor
+        astra_db_docs = []
+        for result in astra_db_docs_cursor:
+            astra_db_docs.append(result)
+
+        # Create a list of AstraDBIngestDoc objects
+        doc_list = []
+        for record in astra_db_docs:
+            doc = AstraDBIngestDoc(
+                connector_config=self.connector_config,
+                processor_config=self.processor_config,
+                read_config=self.read_config,
+                metadata=record,
+            )
+
+            doc.update_source_metadata()
+
+            doc_list.append(doc)
+
+        return doc_list
+
+
+@dataclass
+class AstraDBWriteConfig(WriteConfig):
+    embedding_dimension: int
+    requested_indexing_policy: t.Optional[t.Dict[str, t.Any]] = None
+    batch_size: int = 20
+
+
+@dataclass
+class AstraDBDestinationConnector(BaseDestinationConnector):
+    write_config: AstraDBWriteConfig
+    connector_config: SimpleAstraDBConfig
+    _astra_db: t.Optional["AstraDB"] = field(init=False, default=None)
+    _astra_db_collection: t.Optional["AstraDBCollection"] = field(init=False, default=None)
+
+    def to_dict(self, **kwargs):
+        """
+        The _astra_db_collection variable in this dataclass breaks deepcopy due to:
+        TypeError: cannot pickle '_thread.lock' object
+        When serializing, remove it, meaning client data will need to be reinitialized
+        when deserialized
+        """
+        self_cp = copy.copy(self)
+
+        if hasattr(self_cp, "_astra_db_collection"):
+            setattr(self_cp, "_astra_db_collection", None)
+
+        return _asdict(self_cp, **kwargs)
+
+    @property
+    @requires_dependencies(["astrapy"], extras="astradb")
+    def astra_db_collection(self) -> "AstraDBCollection":
+        if self._astra_db_collection is None:
+            from astrapy import DataAPIClient as AstraDBClient
+            from astrapy.exceptions import CollectionAlreadyExistsException
+
+            # Choose keyspace or deprecated namespace
+            keyspace_param = self.connector_config.keyspace or self.connector_config.namespace
+
+            collection_name = self.connector_config.collection_name
+            embedding_dimension = self.write_config.embedding_dimension
+            requested_indexing_policy = self.write_config.requested_indexing_policy
+
+            # Create a client object to interact with the Astra DB
+            # caller_name/version for Astra DB tracking
+            my_client = AstraDBClient(
+                caller_name=integration_name,
+                caller_version=integration_version,
+            )
+
+            # Get the database object
+            self._astra_db = my_client.get_database(
+                api_endpoint=self.connector_config.access_config.api_endpoint,
+                token=self.connector_config.access_config.token,
+                keyspace=keyspace_param,
+            )
+
+            # Create and connect to the newly created collection
+            try:
+                self._astra_db_collection = self._astra_db.create_collection(
+                    name=collection_name,
+                    dimension=embedding_dimension,
+                    indexing=requested_indexing_policy,
+                )
+            except CollectionAlreadyExistsException as e:
+                logger.info(f"{e}", exc_info=True)
+                self._astra_db_collection = self._astra_db.get_collection(name=collection_name)
+
+        return self._astra_db_collection
+
+    @requires_dependencies(["astrapy"], extras="astradb")
+    @DestinationConnectionError.wrap
+    def initialize(self):
+        _ = self.astra_db_collection
+
+    @requires_dependencies(["astrapy"], extras="astradb")
+    def check_connection(self):
+        try:
+            _ = self.astra_db_collection
+        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(self, *args, elements_dict: t.List[t.Dict[str, t.Any]], **kwargs) -> None:
+        logger.info(f"inserting / updating {len(elements_dict)} documents to Astra DB.")
+
+        if self._astra_db_collection is None:
+            raise DestinationConnectionError("Astra DB collection not available for insertion.")
+
+        astra_db_batch_size = self.write_config.batch_size
+
+        for batch in batch_generator(elements_dict, astra_db_batch_size):
+            self._astra_db_collection.insert_many(batch)
+
+    def normalize_dict(self, element_dict: dict) -> dict:
+        return {
+            "$vector": element_dict.pop("embeddings", None),
+            "content": element_dict.pop("text", None),
+            "metadata": element_dict,
+        }

unstructured_ingest/connector/azure_ai_search.py

@@ -0,0 +1,144 @@
+import json
+import typing as t
+import uuid
+from dataclasses import dataclass, field
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.error import DestinationConnectionError, WriteError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    BaseConnectorConfig,
+    BaseDestinationConnector,
+    WriteConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if t.TYPE_CHECKING:
+    from azure.search.documents import SearchClient
+
+
+@dataclass
+class AzureAiSearchAccessConfig(AccessConfig):
+    key: str = enhanced_field(sensitive=True)
+
+
+@dataclass
+class SimpleAzureAISearchStorageConfig(BaseConnectorConfig):
+    endpoint: str
+    access_config: AzureAiSearchAccessConfig
+
+
+@dataclass
+class AzureAISearchWriteConfig(WriteConfig):
+    index: str
+
+
+@dataclass
+class AzureAISearchDestinationConnector(BaseDestinationConnector):
+    write_config: AzureAISearchWriteConfig
+    connector_config: SimpleAzureAISearchStorageConfig
+    _client: t.Optional["SearchClient"] = field(init=False, default=None)
+
+    @requires_dependencies(["azure.search"], extras="azure-ai-search")
+    def generate_client(self) -> "SearchClient":
+        from azure.core.credentials import AzureKeyCredential
+        from azure.search.documents import SearchClient
+
+        # Create a client
+        credential = AzureKeyCredential(self.connector_config.access_config.key)
+        return SearchClient(
+            endpoint=self.connector_config.endpoint,
+            index_name=self.write_config.index,
+            credential=credential,
+        )
+
+    @property
+    def client(self) -> "SearchClient":
+        if self._client is None:
+            self._client = self.generate_client()
+        return self._client
+
+    def check_connection(self):
+        try:
+            self.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 initialize(self):
+        _ = self.client
+
+    def conform_dict(self, data: dict) -> None:
+        """
+        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
+        Azure Cognitive Search index
+        """
+        from dateutil import parser  # type: ignore
+
+        data["id"] = str(uuid.uuid4())
+
+        if points := data.get("metadata", {}).get("coordinates", {}).get("points"):
+            data["metadata"]["coordinates"]["points"] = json.dumps(points)
+        if version := data.get("metadata", {}).get("data_source", {}).get("version"):
+            data["metadata"]["data_source"]["version"] = str(version)
+        if record_locator := data.get("metadata", {}).get("data_source", {}).get("record_locator"):
+            data["metadata"]["data_source"]["record_locator"] = json.dumps(record_locator)
+        if permissions_data := (
+            data.get("metadata", {}).get("data_source", {}).get("permissions_data")
+        ):
+            data["metadata"]["data_source"]["permissions_data"] = json.dumps(permissions_data)
+        if links := data.get("metadata", {}).get("links"):
+            data["metadata"]["links"] = [json.dumps(link) for link in links]
+        if last_modified := data.get("metadata", {}).get("last_modified"):
+            data["metadata"]["last_modified"] = parser.parse(last_modified).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            )
+        if date_created := data.get("metadata", {}).get("data_source", {}).get("date_created"):
+            data["metadata"]["data_source"]["date_created"] = parser.parse(date_created).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            )
+        if date_modified := data.get("metadata", {}).get("data_source", {}).get("date_modified"):
+            data["metadata"]["data_source"]["date_modified"] = parser.parse(date_modified).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            )
+        if date_processed := data.get("metadata", {}).get("data_source", {}).get("date_processed"):
+            data["metadata"]["data_source"]["date_processed"] = parser.parse(
+                date_processed,
+            ).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+        if regex_metadata := data.get("metadata", {}).get("regex_metadata"):
+            data["metadata"]["regex_metadata"] = json.dumps(regex_metadata)
+        if page_number := data.get("metadata", {}).get("page_number"):
+            data["metadata"]["page_number"] = str(page_number)
+
+    @requires_dependencies(["azure"], extras="azure-ai-search")
+    def write_dict(self, *args, elements_dict: t.List[t.Dict[str, t.Any]], **kwargs) -> None:
+        import azure.core.exceptions
+
+        logger.info(
+            f"writing {len(elements_dict)} documents to destination "
+            f"index at {self.write_config.index}",
+        )
+        try:
+            results = self.client.upload_documents(documents=elements_dict)
+
+        except azure.core.exceptions.HttpResponseError as http_error:
+            raise WriteError(f"http error: {http_error}") from http_error
+        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.key}: [{error.status_code}] {error.error_message}"
+                        for error in errors
+                    ],
+                ),
+            )