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

unstructured_ingest/v2/processes/connectors/{astra.py → astradb.py}

@@ -3,10 +3,10 @@ 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 pydantic import Field, Secret
 
-from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest import __name__ as integration_name
+from unstructured_ingest.__version__ import __version__ as integration_version
 from unstructured_ingest.error import DestinationConnectionError
 from unstructured_ingest.utils.data_prep import batch_generator
 from unstructured_ingest.utils.dep_check import requires_dependencies
@@ -28,30 +28,27 @@ from unstructured_ingest.v2.processes.connector_registry import (
 if TYPE_CHECKING:
     from astrapy.db import AstraDBCollection
 
-CONNECTOR_TYPE = "astra"
+CONNECTOR_TYPE = "astradb"
 
 
-@dataclass
-class AstraAccessConfig(AccessConfig):
-    token: str
-    api_endpoint: str
+class AstraDBAccessConfig(AccessConfig):
+    token: str = Field(description="Astra DB Token with access to the database.")
+    api_endpoint: str = Field(description="The API endpoint for the Astra DB.")
 
 
-@dataclass
-class AstraConnectionConfig(ConnectionConfig):
-    connection_type: str = CONNECTOR_TYPE
-    access_config: AstraAccessConfig = enhanced_field(sensitive=True)
+class AstraDBConnectionConfig(ConnectionConfig):
+    connection_type: str = Field(default=CONNECTOR_TYPE, init=False)
+    access_config: Secret[AstraDBAccessConfig]
 
 
-@dataclass
-class AstraUploadStagerConfig(UploadStagerConfig):
+class AstraDBUploadStagerConfig(UploadStagerConfig):
     pass
 
 
 @dataclass
-class AstraUploadStager(UploadStager):
-    upload_stager_config: AstraUploadStagerConfig = field(
-        default_factory=lambda: AstraUploadStagerConfig()
+class AstraDBUploadStager(UploadStager):
+    upload_stager_config: AstraDBUploadStagerConfig = field(
+        default_factory=lambda: AstraDBUploadStagerConfig()
     )
 
     def conform_dict(self, element_dict: dict) -> dict:
@@ -80,19 +77,28 @@ class AstraUploadStager(UploadStager):
         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
+class AstraDBUploaderConfig(UploaderConfig):
+    collection_name: str = Field(
+        description="The name of the Astra DB collection. "
+        "Note that the collection name must only include letters, "
+        "numbers, and underscores."
+    )
+    embedding_dimension: int = Field(
+        default=384, description="The dimensionality of the embeddings"
+    )
+    namespace: Optional[str] = Field(default=None, description="The Astra DB connection namespace.")
+    requested_indexing_policy: Optional[dict[str, Any]] = Field(
+        default=None,
+        description="The indexing policy to use for the collection.",
+        examples=['{"deny": ["metadata"]}'],
+    )
+    batch_size: int = Field(default=20, description="Number of records per batch")
 
 
 @dataclass
-class AstraUploader(Uploader):
-    connection_config: AstraConnectionConfig
-    upload_config: AstraUploaderConfig
+class AstraDBUploader(Uploader):
+    connection_config: AstraDBConnectionConfig
+    upload_config: AstraDBUploaderConfig
     connector_type: str = CONNECTOR_TYPE
 
     def precheck(self) -> None:
@@ -102,7 +108,7 @@ class AstraUploader(Uploader):
             logger.error(f"Failed to validate connection {e}", exc_info=True)
             raise DestinationConnectionError(f"failed to validate connection: {e}")
 
-    @requires_dependencies(["astrapy"], extras="astra")
+    @requires_dependencies(["astrapy"], extras="astradb")
     def get_collection(self) -> "AstraDBCollection":
         from astrapy.db import AstraDB
 
@@ -111,14 +117,15 @@ class AstraUploader(Uploader):
         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
+        # If the user has requested an indexing policy, pass it to the Astra DB
         options = {"indexing": requested_indexing_policy} if requested_indexing_policy else None
 
         # Build the Astra DB object.
         # caller_name/version for AstraDB tracking
+        access_configs = self.connection_config.access_config.get_secret_value()
         astra_db = AstraDB(
-            api_endpoint=self.connection_config.access_config.api_endpoint,
-            token=self.connection_config.access_config.token,
+            api_endpoint=access_configs.api_endpoint,
+            token=access_configs.token,
             namespace=self.upload_config.namespace,
             caller_name=integration_name,
             caller_version=integration_version,
@@ -144,17 +151,17 @@ class AstraUploader(Uploader):
             f"collection {self.upload_config.collection_name}"
         )
 
-        astra_batch_size = self.upload_config.batch_size
+        astra_db_batch_size = self.upload_config.batch_size
         collection = self.get_collection()
 
-        for chunk in batch_generator(elements_dict, astra_batch_size):
+        for chunk in batch_generator(elements_dict, astra_db_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,
+astra_db_destination_entry = DestinationRegistryEntry(
+    connection_config=AstraDBConnectionConfig,
+    upload_stager_config=AstraDBUploadStagerConfig,
+    upload_stager=AstraDBUploadStager,
+    uploader_config=AstraDBUploaderConfig,
+    uploader=AstraDBUploader,
 )

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,
             )