unstructured-ingest 0.0.3__py3-none-any.whl → 0.0.4__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/azure_cognitive_search.py

@@ -1,10 +1,11 @@
 import json
-import typing as t
 import uuid
 from dataclasses import dataclass, field
 from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from pydantic import Field, Secret
 
-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
@@ -20,27 +21,31 @@ from unstructured_ingest.v2.interfaces import (
 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:
+if 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")
+    azure_cognitive_search_key: str = Field(
+        alias="key", description="Credential that is used for authenticating to an Azure service"
+    )
 
 
-@dataclass
 class AzureCognitiveSearchConnectionConfig(ConnectionConfig):
-    endpoint: str
-    index: str
-    access_config: AzureCognitiveSearchAccessConfig = enhanced_field(sensitive=True)
+    endpoint: str = Field(
+        description="The URL endpoint of an Azure AI (Cognitive) search service. "
+        "In the form of https://{{service_name}}.search.windows.net"
+    )
+    index: str = Field(
+        description="The name of the Azure AI (Cognitive) Search index to connect to."
+    )
+    access_config: Secret[AzureCognitiveSearchAccessConfig]
 
     @requires_dependencies(["azure.search", "azure.core"], extras="azure-cognitive-search")
     def generate_client(self) -> "SearchClient":
@@ -50,18 +55,18 @@ class AzureCognitiveSearchConnectionConfig(ConnectionConfig):
         return SearchClient(
             endpoint=self.endpoint,
             index_name=self.index,
-            credential=AzureKeyCredential(self.access_config.key),
+            credential=AzureKeyCredential(
+                self.access_config.get_secret_value().azure_cognitive_search_key
+            ),
         )
 
 
-@dataclass
 class AzureCognitiveSearchUploadStagerConfig(UploadStagerConfig):
     pass
 
 
-@dataclass
 class AzureCognitiveSearchUploaderConfig(UploaderConfig):
-    batch_size: int = 100
+    batch_size: int = Field(default=100, description="Number of records per batch")
 
 
 @dataclass
@@ -122,7 +127,7 @@ class AzureCognitiveSearchUploadStager(UploadStager):
         elements_filepath: Path,
         output_dir: Path,
         output_filename: str,
-        **kwargs: t.Any,
+        **kwargs: Any,
     ) -> Path:
         with open(elements_filepath) as elements_file:
             elements_contents = json.load(elements_file)
@@ -143,7 +148,7 @@ class AzureCognitiveSearchUploader(Uploader):
 
     @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:
+    def write_dict(self, *args, elements_dict: list[dict[str, Any]], **kwargs) -> None:
         import azure.core.exceptions
 
         logger.info(
@@ -169,7 +174,8 @@ class AzureCognitiveSearchUploader(Uploader):
             raise WriteError(
                 ", ".join(
                     [
-                        f"{error.key}: [{error.status_code}] {error.error_message}"
+                        f"{error.azure_cognitive_search_key}: "
+                        f"[{error.status_code}] {error.error_message}"
                         for error in errors
                     ],
                 ),
@@ -186,7 +192,7 @@ class AzureCognitiveSearchUploader(Uploader):
     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:
+    def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
 
         elements_dict = []
         for content in contents:
@@ -207,13 +213,10 @@ class AzureCognitiveSearchUploader(Uploader):
             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,
-    ),
+azure_cognitive_search_destination_entry = DestinationRegistryEntry(
+    connection_config=AzureCognitiveSearchConnectionConfig,
+    uploader=AzureCognitiveSearchUploader,
+    uploader_config=AzureCognitiveSearchUploaderConfig,
+    upload_stager=AzureCognitiveSearchUploadStager,
+    upload_stager_config=AzureCognitiveSearchUploadStagerConfig,
 )

unstructured_ingest/v2/processes/connectors/chroma.py

@@ -3,11 +3,11 @@ 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 typing import TYPE_CHECKING, Any, Optional
 
 from dateutil import parser
+from pydantic import Field, Secret
 
-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
@@ -32,26 +32,35 @@ if TYPE_CHECKING:
 CONNECTOR_TYPE = "chroma"
 
 
-@dataclass
 class ChromaAccessConfig(AccessConfig):
-    settings: Optional[Dict[str, str]] = None
-    headers: Optional[Dict[str, str]] = None
+    settings: Optional[dict[str, str]] = Field(
+        default=None, description="A dictionary of settings to communicate with the chroma server."
+    )
+    headers: Optional[dict[str, str]] = Field(
+        default=None, description="A dictionary of headers to send to the Chroma server."
+    )
 
 
-@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
+    collection_name: str = Field(description="The name of the Chroma collection to write into.")
+    access_config: Secret[ChromaAccessConfig]
+    path: Optional[str] = Field(
+        default=None, description="Location where Chroma is persisted, if not connecting via http."
+    )
+    tenant: Optional[str] = Field(
+        default="default_tenant", description="The tenant to use for this client."
+    )
+    database: Optional[str] = Field(
+        default="default_database", description="The database to use for this client."
+    )
+    host: Optional[str] = Field(default=None, description="The hostname of the Chroma server.")
+    port: Optional[int] = Field(default=None, description="The port of the Chroma server.")
+    ssl: bool = Field(
+        default=False, description="Whether to use SSL to connect to the Chroma server."
+    )
+    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
 
 
-@dataclass
 class ChromaUploadStagerConfig(UploadStagerConfig):
     pass
 
@@ -101,9 +110,8 @@ class ChromaUploadStager(UploadStager):
         return output_path
 
 
-@dataclass
 class ChromaUploaderConfig(UploaderConfig):
-    batch_size: int = 100
+    batch_size: int = Field(default=100, description="Number of records per batch")
 
 
 @dataclass
@@ -123,10 +131,11 @@ class ChromaUploader(Uploader):
     def create_client(self) -> "Client":
         import chromadb
 
+        access_config = self.connection_config.access_config.get_secret_value()
         if self.connection_config.path:
             return chromadb.PersistentClient(
                 path=self.connection_config.path,
-                settings=self.connection_config.access_config.settings,
+                settings=access_config.settings,
                 tenant=self.connection_config.tenant,
                 database=self.connection_config.database,
             )
@@ -136,8 +145,8 @@ class ChromaUploader(Uploader):
                 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,
+                headers=access_config.headers,
+                settings=access_config.settings,
                 tenant=self.connection_config.tenant,
                 database=self.connection_config.database,
             )

unstructured_ingest/v2/processes/connectors/couchbase.py

@@ -0,0 +1,151 @@
+import json
+from dataclasses import dataclass, field
+from datetime import timedelta
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.error import DestinationConnectionError
+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,
+)
+
+if TYPE_CHECKING:
+    from couchbase.cluster import Cluster
+
+CONNECTOR_TYPE = "couchbase"
+SERVER_API_VERSION = "1"
+
+
+class CouchbaseAccessConfig(AccessConfig):
+    password: str = Field(description="The password for the Couchbase server")
+
+
+class CouchbaseConnectionConfig(ConnectionConfig):
+    username: str = Field(description="The username for the Couchbase server")
+    bucket: str = Field(description="The bucket to connect to on the Couchbase server")
+    connection_string: str = Field(
+        default="couchbase://localhost", description="The connection string of the Couchbase server"
+    )
+    scope: str = Field(
+        default="_default", description="The scope to connect to on the Couchbase server"
+    )
+    collection: str = Field(
+        default="_default", description="The collection to connect to on the Couchbase server"
+    )
+    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
+    access_config: Secret[CouchbaseAccessConfig]
+
+
+class CouchbaseUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class CouchbaseUploadStager(UploadStager):
+    upload_stager_config: CouchbaseUploadStagerConfig = field(
+        default_factory=lambda: CouchbaseUploadStagerConfig()
+    )
+
+    def run(
+        self,
+        elements_filepath: Path,
+        output_dir: Path,
+        output_filename: str,
+        **kwargs: Any,
+    ) -> Path:
+        with open(elements_filepath) as elements_file:
+            elements_contents = json.load(elements_file)
+
+        output_elements = []
+        for element in elements_contents:
+            new_doc = {
+                element["element_id"]: {
+                    "embedding": element.get("embeddings", None),
+                    "text": element.get("text", None),
+                    "metadata": element.get("metadata", None),
+                    "type": element.get("type", None),
+                }
+            }
+            output_elements.append(new_doc)
+
+        output_path = Path(output_dir) / Path(f"{output_filename}.json")
+        with open(output_path, "w") as output_file:
+            json.dump(output_elements, output_file)
+        return output_path
+
+
+class CouchbaseUploaderConfig(UploaderConfig):
+    batch_size: int = Field(default=50, description="Number of documents to upload per batch")
+
+
+@dataclass
+class CouchbaseUploader(Uploader):
+    connection_config: CouchbaseConnectionConfig
+    upload_config: CouchbaseUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @requires_dependencies(["couchbase"], extras="couchbase")
+    def connect_to_couchbase(self) -> "Cluster":
+        from couchbase.auth import PasswordAuthenticator
+        from couchbase.cluster import Cluster
+        from couchbase.options import ClusterOptions
+
+        connection_string = self.connection_config.connection_string
+        username = self.connection_config.username
+        password = self.connection_config.access_config.get_secret_value().password
+
+        auth = PasswordAuthenticator(username, password)
+        options = ClusterOptions(auth)
+        options.apply_profile("wan_development")
+        cluster = Cluster(connection_string, options)
+        cluster.wait_until_ready(timedelta(seconds=5))
+        return cluster
+
+    def precheck(self) -> None:
+        try:
+            self.connect_to_couchbase()
+        except Exception as e:
+            logger.error(f"Failed to validate connection {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
+        elements = []
+        for content in contents:
+            with open(content.path) as elements_file:
+                elements.extend(json.load(elements_file))
+
+        logger.info(
+            f"writing {len(elements)} objects to destination "
+            f"bucket, {self.connection_config.bucket} "
+            f"at {self.connection_config.connection_string}",
+        )
+        cluster = self.connect_to_couchbase()
+        bucket = cluster.bucket(self.connection_config.bucket)
+        scope = bucket.scope(self.connection_config.scope)
+        collection = scope.collection(self.connection_config.collection)
+
+        for chunk in batch_generator(elements, self.upload_config.batch_size):
+            collection.upsert_multi({doc_id: doc for doc in chunk for doc_id, doc in doc.items()})
+
+
+couchbase_destination_entry = DestinationRegistryEntry(
+    connection_config=CouchbaseConnectionConfig,
+    uploader=CouchbaseUploader,
+    uploader_config=CouchbaseUploaderConfig,
+    upload_stager=CouchbaseUploadStager,
+    upload_stager_config=CouchbaseUploadStagerConfig,
+)

unstructured_ingest/v2/processes/connectors/databricks_volumes.py

@@ -1,8 +1,9 @@
 import os
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Optional
 
-from unstructured_ingest.enhanced_dataclass import enhanced_field
+from pydantic import Field, Secret
+
 from unstructured_ingest.error import DestinationConnectionError
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import (
@@ -21,45 +22,99 @@ if TYPE_CHECKING:
 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
+    account_id: Optional[str] = Field(
+        default=None,
+        description="The Databricks account ID for the Databricks "
+        "accounts endpoint. Only has effect when Host is "
+        "either https://accounts.cloud.databricks.com/ (AWS), "
+        "https://accounts.azuredatabricks.net/ (Azure), "
+        "or https://accounts.gcp.databricks.com/ (GCP).",
+    )
+    username: Optional[str] = Field(
+        default=None,
+        description="The Databricks username part of basic authentication. "
+        "Only possible when Host is *.cloud.databricks.com (AWS).",
+    )
+    password: Optional[str] = Field(
+        default=None,
+        description="The Databricks password part of basic authentication. "
+        "Only possible when Host is *.cloud.databricks.com (AWS).",
+    )
+    client_id: Optional[str] = Field(default=None)
+    client_secret: Optional[str] = Field(default=None)
+    token: Optional[str] = Field(
+        default=None,
+        description="The Databricks personal access token (PAT) (AWS, Azure, and GCP) or "
+        "Azure Active Directory (Azure AD) token (Azure).",
+    )
     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
+    azure_workspace_resource_id: Optional[str] = Field(
+        default=None,
+        description="The Azure Resource Manager ID for the Azure Databricks workspace, "
+        "which is exchanged for a Databricks host URL.",
+    )
+    azure_client_secret: Optional[str] = Field(
+        default=None, description="The Azure AD service principal’s client secret."
+    )
+    azure_client_id: Optional[str] = Field(
+        default=None, description="The Azure AD service principal’s application ID."
+    )
+    azure_tenant_id: Optional[str] = Field(
+        default=None, description="The Azure AD service principal’s tenant ID."
+    )
+    azure_environment: Optional[str] = Field(
+        default=None,
+        description="The Azure environment type for a " "specific set of API endpoints",
+        examples=["Public", "UsGov", "China", "Germany"],
+    )
+    auth_type: Optional[str] = Field(
+        default=None,
+        description="When multiple auth attributes are available in the "
+        "environment, use the auth type specified by this "
+        "argument. This argument also holds the currently "
+        "selected auth.",
+    )
     cluster_id: Optional[str] = None
     google_credentials: Optional[str] = None
     google_service_account: Optional[str] = None
 
 
-@dataclass
+SecretDatabricksVolumesAccessConfig = Secret[DatabricksVolumesAccessConfig]
+
+
 class DatabricksVolumesConnectionConfig(ConnectionConfig):
-    access_config: DatabricksVolumesAccessConfig = enhanced_field(
-        default_factory=DatabricksVolumesAccessConfig, sensitive=True
+    access_config: SecretDatabricksVolumesAccessConfig = Field(
+        default_factory=lambda: SecretDatabricksVolumesAccessConfig(
+            secret_value=DatabricksVolumesAccessConfig()
+        )
+    )
+    host: Optional[str] = Field(
+        default=None,
+        description="The Databricks host URL for either the "
+        "Databricks workspace endpoint or the "
+        "Databricks accounts endpoint.",
     )
-    host: Optional[str] = None
 
 
-@dataclass
 class DatabricksVolumesUploaderConfig(UploaderConfig):
-    volume: str
-    catalog: str
-    volume_path: Optional[str] = None
-    overwrite: bool = False
-    schema: str = "default"
+    volume: str = Field(description="Name of volume in the Unity Catalog")
+    catalog: str = Field(description="Name of the catalog in the Databricks Unity Catalog service")
+    volume_path: Optional[str] = Field(
+        default=None, description="Optional path within the volume to write to"
+    )
+    overwrite: bool = Field(
+        default=False, description="If true, an existing file will be overwritten."
+    )
+    databricks_schema: str = Field(
+        default="default",
+        alias="schema",
+        description="Schema associated with the volume to write to in the Unity Catalog service",
+    )
 
     @property
     def path(self) -> str:
-        path = f"/Volumes/{self.catalog}/{self.schema}/{self.volume}"
+        path = f"/Volumes/{self.catalog}/{self.databricks_schema}/{self.volume}"
         if self.volume_path:
             path = f"{path}/{self.volume_path}"
         return path
@@ -70,19 +125,19 @@ 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":
+    def get_client(self) -> "WorkspaceClient":
         from databricks.sdk import WorkspaceClient
 
-        self.client = WorkspaceClient(
-            host=self.connection_config.host, **self.connection_config.access_config.to_dict()
+        return WorkspaceClient(
+            host=self.connection_config.host,
+            **self.connection_config.access_config.get_secret_value().dict(),
         )
 
     def precheck(self) -> None:
         try:
-            assert self.client.current_user.me().active
+            assert self.get_client().current_user.me().active
         except Exception as e:
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise DestinationConnectionError(f"failed to validate connection: {e}")
@@ -91,7 +146,7 @@ class DatabricksVolumesUploader(Uploader):
         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(
+                self.get_client().files.upload(
                     file_path=output_path,
                     contents=elements_file,
                     overwrite=self.upload_config.overwrite,