unstructured-ingest 0.6.4__py3-none-any.whl → 0.7.1__py3-none-any.whl

unstructured_ingest/connector/astradb.py

@@ -1,267 +0,0 @@
-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

@@ -1,144 +0,0 @@
-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
-                    ],
-                ),
-            )