unstructured-ingest 0.0.0__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/astra.py

@@ -0,0 +1,152 @@
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional
+
+from unstructured import __name__ as integration_name
+from unstructured.__version__ import __version__ as integration_version
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.utils.data_prep import batch_generator
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    FileData,
+    UploadContent,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+)
+
+if TYPE_CHECKING:
+    from astrapy.db import AstraDBCollection
+
+CONNECTOR_TYPE = "astra"
+
+
+@dataclass
+class AstraAccessConfig(AccessConfig):
+    token: str
+    api_endpoint: str
+
+
+@dataclass
+class AstraConnectionConfig(ConnectionConfig):
+    connection_type: str = CONNECTOR_TYPE
+    access_config: AstraAccessConfig = enhanced_field(sensitive=True)
+
+
+@dataclass
+class AstraUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class AstraUploadStager(UploadStager):
+    upload_stager_config: AstraUploadStagerConfig = field(
+        default_factory=lambda: AstraUploadStagerConfig()
+    )
+
+    def conform_dict(self, element_dict: dict) -> dict:
+        return {
+            "$vector": element_dict.pop("embeddings", None),
+            "content": element_dict.pop("text", None),
+            "metadata": element_dict,
+        }
+
+    def run(
+        self,
+        elements_filepath: Path,
+        file_data: FileData,
+        output_dir: Path,
+        output_filename: str,
+        **kwargs: Any,
+    ) -> Path:
+        with open(elements_filepath) as elements_file:
+            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")
+        with open(output_path, "w") as output_file:
+            json.dump(conformed_elements, output_file)
+        return output_path
+
+
+@dataclass
+class AstraUploaderConfig(UploaderConfig):
+    collection_name: str
+    embedding_dimension: int
+    namespace: Optional[str] = None
+    requested_indexing_policy: Optional[dict[str, Any]] = None
+    batch_size: int = 20
+
+
+@dataclass
+class AstraUploader(Uploader):
+    connection_config: AstraConnectionConfig
+    upload_config: AstraUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @requires_dependencies(["astrapy"], extras="astra")
+    def get_collection(self) -> "AstraDBCollection":
+        from astrapy.db import AstraDB
+
+        # Get the collection_name and embedding dimension
+        collection_name = self.upload_config.collection_name
+        embedding_dimension = self.upload_config.embedding_dimension
+        requested_indexing_policy = self.upload_config.requested_indexing_policy
+
+        # If the user has requested an indexing policy, pass it to the AstraDB
+        options = {"indexing": requested_indexing_policy} if requested_indexing_policy else None
+
+        # Build the Astra DB object.
+        # caller_name/version for AstraDB tracking
+        astra_db = AstraDB(
+            api_endpoint=self.connection_config.access_config.api_endpoint,
+            token=self.connection_config.access_config.token,
+            namespace=self.upload_config.namespace,
+            caller_name=integration_name,
+            caller_version=integration_version,
+        )
+
+        # Create and connect to the newly created collection
+        astra_db_collection = astra_db.create_collection(
+            collection_name=collection_name,
+            dimension=embedding_dimension,
+            options=options,
+        )
+        return astra_db_collection
+
+    def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
+        elements_dict = []
+        for content in contents:
+            with open(content.path) as elements_file:
+                elements = json.load(elements_file)
+                elements_dict.extend(elements)
+
+        logger.info(
+            f"writing {len(elements_dict)} objects to destination "
+            f"collection {self.upload_config.collection_name}"
+        )
+
+        astra_batch_size = self.upload_config.batch_size
+        collection = self.get_collection()
+
+        for chunk in batch_generator(elements_dict, astra_batch_size):
+            collection.insert_many(chunk)
+
+
+astra_destination_entry = DestinationRegistryEntry(
+    connection_config=AstraConnectionConfig,
+    upload_stager_config=AstraUploadStagerConfig,
+    upload_stager=AstraUploadStager,
+    uploader_config=AstraUploaderConfig,
+    uploader=AstraUploader,
+)

unstructured_ingest/v2/processes/connectors/azure_cognitive_search.py

@@ -0,0 +1,211 @@
+import json
+import typing as t
+import uuid
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+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.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    UploadContent,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    add_destination_entry,
+)
+from unstructured_ingest.v2.processes.connectors.utils import parse_datetime
+
+if t.TYPE_CHECKING:
+    from azure.search.documents import SearchClient
+
+
+CONNECTOR_TYPE = "azure_cognitive_search"
+
+
+@dataclass
+class AzureCognitiveSearchAccessConfig(AccessConfig):
+    key: t.Optional[str] = enhanced_field(default=None, overload_name="azure_cognitive_search_key")
+
+
+@dataclass
+class AzureCognitiveSearchConnectionConfig(ConnectionConfig):
+    endpoint: str
+    index: str
+    access_config: AzureCognitiveSearchAccessConfig = enhanced_field(sensitive=True)
+
+    @requires_dependencies(["azure.search", "azure.core"], extras="azure-cognitive-search")
+    def generate_client(self) -> "SearchClient":
+        from azure.core.credentials import AzureKeyCredential
+        from azure.search.documents import SearchClient
+
+        return SearchClient(
+            endpoint=self.endpoint,
+            index_name=self.index,
+            credential=AzureKeyCredential(self.access_config.key),
+        )
+
+
+@dataclass
+class AzureCognitiveSearchUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class AzureCognitiveSearchUploaderConfig(UploaderConfig):
+    batch_size: int = 100
+
+
+@dataclass
+class AzureCognitiveSearchUploadStager(UploadStager):
+    upload_stager_config: AzureCognitiveSearchUploadStagerConfig = field(
+        default_factory=lambda: AzureCognitiveSearchUploadStagerConfig()
+    )
+
+    @staticmethod
+    def conform_dict(data: dict) -> 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
+        Azure Cognitive Search index
+        """
+
+        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"] = parse_datetime(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"] = parse_datetime(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"] = parse_datetime(
+                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"] = parse_datetime(
+                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)
+        return data
+
+    def run(
+        self,
+        elements_filepath: Path,
+        output_dir: Path,
+        output_filename: str,
+        **kwargs: t.Any,
+    ) -> Path:
+        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]
+
+        output_path = Path(output_dir) / Path(f"{output_filename}.json")
+        with open(output_path, "w") as output_file:
+            json.dump(conformed_elements, output_file)
+        return output_path
+
+
+@dataclass
+class AzureCognitiveSearchUploader(Uploader):
+    upload_config: AzureCognitiveSearchUploaderConfig
+    connection_config: AzureCognitiveSearchConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @DestinationConnectionError.wrap
+    @requires_dependencies(["azure"], extras="azure-cognitive-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.connection_config.index}",
+        )
+        try:
+            results = self.connection_config.generate_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
+                    ],
+                ),
+            )
+
+    def write_dict_wrapper(self, elements_dict):
+        return self.write_dict(elements_dict=elements_dict)
+
+    def run(self, contents: list[UploadContent], **kwargs: t.Any) -> None:
+
+        elements_dict = []
+        for content in contents:
+            with open(content.path) as elements_file:
+                elements = json.load(elements_file)
+                elements_dict.extend(elements)
+
+        logger.info(
+            f"writing document batches to destination"
+            f" endpoint at {str(self.connection_config.endpoint)}"
+            f" index at {str(self.connection_config.index)}"
+            f" with batch size {str(self.upload_config.batch_size)}"
+        )
+
+        batch_size = self.upload_config.batch_size
+
+        for chunk in batch_generator(elements_dict, batch_size):
+            self.write_dict(elements_dict=chunk)  # noqa: E203
+
+
+add_destination_entry(
+    destination_type=CONNECTOR_TYPE,
+    entry=DestinationRegistryEntry(
+        connection_config=AzureCognitiveSearchConnectionConfig,
+        uploader=AzureCognitiveSearchUploader,
+        uploader_config=AzureCognitiveSearchUploaderConfig,
+        upload_stager=AzureCognitiveSearchUploadStager,
+        upload_stager_config=AzureCognitiveSearchUploadStagerConfig,
+    ),
+)

unstructured_ingest/v2/processes/connectors/chroma.py

@@ -0,0 +1,204 @@
+import json
+import uuid
+from dataclasses import dataclass, field
+from datetime import date, datetime
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+from dateutil import parser
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.error import DestinationConnectionError
+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,
+    ConnectionConfig,
+    FileData,
+    UploadContent,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+)
+
+if TYPE_CHECKING:
+    from chromadb import Client
+
+CONNECTOR_TYPE = "chroma"
+
+
+@dataclass
+class ChromaAccessConfig(AccessConfig):
+    settings: Optional[Dict[str, str]] = None
+    headers: Optional[Dict[str, str]] = None
+
+
+@dataclass
+class ChromaConnectionConfig(ConnectionConfig):
+    collection_name: str
+    access_config: ChromaAccessConfig = enhanced_field(sensitive=True)
+    path: Optional[str] = None
+    tenant: Optional[str] = "default_tenant"
+    database: Optional[str] = "default_database"
+    host: Optional[str] = None
+    port: Optional[int] = None
+    ssl: bool = False
+    connector_type: str = CONNECTOR_TYPE
+
+
+@dataclass
+class ChromaUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class ChromaUploadStager(UploadStager):
+    upload_stager_config: ChromaUploadStagerConfig = field(
+        default_factory=lambda: ChromaUploadStagerConfig()
+    )
+
+    @staticmethod
+    def parse_date_string(date_string: str) -> date:
+        try:
+            timestamp = float(date_string)
+            return datetime.fromtimestamp(timestamp)
+        except Exception as e:
+            logger.debug(f"date {date_string} string not a timestamp: {e}")
+        return parser.parse(date_string)
+
+    @staticmethod
+    def conform_dict(data: dict) -> dict:
+        """
+        Prepares dictionary in the format that Chroma requires
+        """
+        element_id = data.get("element_id", str(uuid.uuid4()))
+        return {
+            "id": element_id,
+            "embedding": data.pop("embeddings", None),
+            "document": data.pop("text", None),
+            "metadata": flatten_dict(data, separator="-", flatten_lists=True, remove_none=True),
+        }
+
+    def run(
+        self,
+        elements_filepath: Path,
+        file_data: FileData,
+        output_dir: Path,
+        output_filename: str,
+        **kwargs: Any,
+    ) -> Path:
+        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]
+        output_path = Path(output_dir) / Path(f"{output_filename}.json")
+        with open(output_path, "w") as output_file:
+            json.dump(conformed_elements, output_file)
+        return output_path
+
+
+@dataclass
+class ChromaUploaderConfig(UploaderConfig):
+    batch_size: int = 100
+
+
+@dataclass
+class ChromaUploader(Uploader):
+    connector_type: str = CONNECTOR_TYPE
+    upload_config: ChromaUploaderConfig
+    connection_config: ChromaConnectionConfig
+    client: Optional["Client"] = field(init=False)
+
+    def __post_init__(self):
+        self.client = self.create_client()
+
+    @requires_dependencies(["chromadb"], extras="chroma")
+    def create_client(self) -> "Client":
+        import chromadb
+
+        if self.connection_config.path:
+            return chromadb.PersistentClient(
+                path=self.connection_config.path,
+                settings=self.connection_config.access_config.settings,
+                tenant=self.connection_config.tenant,
+                database=self.connection_config.database,
+            )
+
+        elif self.connection_config.host and self.connection_config.port:
+            return chromadb.HttpClient(
+                host=self.connection_config.host,
+                port=self.connection_config.port,
+                ssl=self.connection_config.ssl,
+                headers=self.connection_config.access_config.headers,
+                settings=self.connection_config.access_config.settings,
+                tenant=self.connection_config.tenant,
+                database=self.connection_config.database,
+            )
+        else:
+            raise ValueError("Chroma connector requires either path or host and port to be set.")
+
+    @DestinationConnectionError.wrap
+    def upsert_batch(self, collection, batch):
+
+        try:
+            # Chroma wants lists even if there is only one element
+            # Upserting to prevent duplicates
+            collection.upsert(
+                ids=batch["ids"],
+                documents=batch["documents"],
+                embeddings=batch["embeddings"],
+                metadatas=batch["metadatas"],
+            )
+        except Exception as e:
+            raise ValueError(f"chroma error: {e}") from e
+
+    @staticmethod
+    def prepare_chroma_list(chunk: tuple[dict[str, Any]]) -> dict[str, list[Any]]:
+        """Helper function to break a tuple of dicts into list of parallel lists for ChromaDb.
+        ({'id':1}, {'id':2}, {'id':3}) -> {'ids':[1,2,3]}"""
+        chroma_dict = {}
+        chroma_dict["ids"] = [x.get("id") for x in chunk]
+        chroma_dict["documents"] = [x.get("document") for x in chunk]
+        chroma_dict["embeddings"] = [x.get("embedding") for x in chunk]
+        chroma_dict["metadatas"] = [x.get("metadata") for x in chunk]
+        # Make sure all lists are of the same length
+        assert (
+            len(chroma_dict["ids"])
+            == len(chroma_dict["documents"])
+            == len(chroma_dict["embeddings"])
+            == len(chroma_dict["metadatas"])
+        )
+        return chroma_dict
+
+    def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
+
+        elements_dict = []
+        for content in contents:
+            with open(content.path) as elements_file:
+                elements = json.load(elements_file)
+                elements_dict.extend(elements)
+
+        logger.info(
+            f"writing {len(elements_dict)} objects to destination "
+            f"collection {self.connection_config.collection_name} "
+            f"at {self.connection_config.host}",
+        )
+
+        collection = self.client.get_or_create_collection(
+            name=self.connection_config.collection_name
+        )
+        for chunk in batch_generator(elements_dict, self.upload_config.batch_size):
+            self.upsert_batch(collection, self.prepare_chroma_list(chunk))
+
+
+chroma_destination_entry = DestinationRegistryEntry(
+    connection_config=ChromaConnectionConfig,
+    uploader=ChromaUploader,
+    uploader_config=ChromaUploaderConfig,
+    upload_stager=ChromaUploadStager,
+    upload_stager_config=ChromaUploadStagerConfig,
+)

unstructured_ingest/v2/processes/connectors/databricks_volumes.py

@@ -0,0 +1,96 @@
+import os
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Optional
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    UploadContent,
+    Uploader,
+    UploaderConfig,
+)
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+
+if TYPE_CHECKING:
+    from databricks.sdk import WorkspaceClient
+
+CONNECTOR_TYPE = "databricks_volumes"
+
+
+@dataclass
+class DatabricksVolumesAccessConfig(AccessConfig):
+    account_id: Optional[str] = None
+    username: Optional[str] = None
+    password: Optional[str] = None
+    client_id: Optional[str] = None
+    client_secret: Optional[str] = None
+    token: Optional[str] = None
+    profile: Optional[str] = None
+    azure_workspace_resource_id: Optional[str] = None
+    azure_client_secret: Optional[str] = None
+    azure_client_id: Optional[str] = None
+    azure_tenant_id: Optional[str] = None
+    azure_environment: Optional[str] = None
+    auth_type: Optional[str] = None
+    cluster_id: Optional[str] = None
+    google_credentials: Optional[str] = None
+    google_service_account: Optional[str] = None
+
+
+@dataclass
+class DatabricksVolumesConnectionConfig(ConnectionConfig):
+    access_config: DatabricksVolumesAccessConfig = enhanced_field(
+        default_factory=DatabricksVolumesAccessConfig, sensitive=True
+    )
+    host: Optional[str] = None
+
+
+@dataclass
+class DatabricksVolumesUploaderConfig(UploaderConfig):
+    volume: str
+    catalog: str
+    volume_path: Optional[str] = None
+    overwrite: bool = False
+    schema: str = "default"
+
+    @property
+    def path(self) -> str:
+        path = f"/Volumes/{self.catalog}/{self.schema}/{self.volume}"
+        if self.volume_path:
+            path = f"{path}/{self.volume_path}"
+        return path
+
+
+@dataclass
+class DatabricksVolumesUploader(Uploader):
+    connector_type: str = CONNECTOR_TYPE
+    upload_config: DatabricksVolumesUploaderConfig
+    connection_config: DatabricksVolumesConnectionConfig
+    client: Optional["WorkspaceClient"] = field(init=False, default=None)
+
+    @requires_dependencies(dependencies=["databricks.sdk"], extras="databricks-volumes")
+    def __post_init__(self) -> "WorkspaceClient":
+        from databricks.sdk import WorkspaceClient
+
+        self.client = WorkspaceClient(
+            host=self.connection_config.host, **self.connection_config.access_config.to_dict()
+        )
+
+    def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
+        for content in contents:
+            with open(content.path, "rb") as elements_file:
+                output_path = os.path.join(self.upload_config.path, content.path.name)
+                self.client.files.upload(
+                    file_path=output_path,
+                    contents=elements_file,
+                    overwrite=self.upload_config.overwrite,
+                )
+
+
+databricks_volumes_destination_entry = DestinationRegistryEntry(
+    connection_config=DatabricksVolumesConnectionConfig,
+    uploader=DatabricksVolumesUploader,
+    uploader_config=DatabricksVolumesUploaderConfig,
+)