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

unstructured_ingest/v2/processes/connectors/google_drive.py

@@ -1,18 +1,18 @@
 import io
-import os
+import json
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Generator, Optional, Union
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Generator, Optional
 
 from dateutil import parser
-from unstructured.file_utils.google_filetype import GOOGLE_DRIVE_EXPORT_TYPES
+from pydantic import Field, Secret
 
-from unstructured_ingest.enhanced_dataclass import enhanced_field
 from unstructured_ingest.error import (
     SourceConnectionError,
     SourceConnectionNetworkError,
 )
 from unstructured_ingest.utils.dep_check import requires_dependencies
-from unstructured_ingest.utils.string_and_date_utils import json_to_dict
+from unstructured_ingest.utils.google_filetype import GOOGLE_DRIVE_EXPORT_TYPES
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
     ConnectionConfig,
@@ -37,46 +37,54 @@ if TYPE_CHECKING:
     from googleapiclient.http import MediaIoBaseDownload
 
 
-@dataclass
 class GoogleDriveAccessConfig(AccessConfig):
-    service_account_key: Union[str, dict]
+    service_account_key: Optional[dict] = Field(
+        default=None, description="Credentials values to use for authentication"
+    )
+    service_account_key_path: Optional[Path] = Field(
+        default=None, description="File path to credentials values to use for authentication"
+    )
+
+    def model_post_init(self, __context: Any) -> None:
+        if self.service_account_key is None and self.service_account_key_path is None:
+            raise ValueError(
+                "either service_account_key or service_account_key_path must be provided"
+            )
+
+    def get_service_account_key(self) -> dict:
+        key_data = None
+        if self.service_account_key_path:
+            with self.service_account_key_path.open() as f:
+                key_data = json.load(f)
+        if key_data and self.service_account_key:
+            if key_data == self.service_account_key:
+                return key_data
+            else:
+                raise ValueError(
+                    "service_account_key and service_account_key_path "
+                    "both provided and have different values"
+                )
+        if key_data:
+            return key_data
+        return self.service_account_key
 
 
-@dataclass
 class GoogleDriveConnectionConfig(ConnectionConfig):
-    drive_id: str
-    access_config: GoogleDriveAccessConfig = enhanced_field(sensitive=True)
+    drive_id: str = Field(description="Google Drive File or Folder ID.")
+    access_config: Secret[GoogleDriveAccessConfig]
 
     @requires_dependencies(["googleapiclient"], extras="google-drive")
     def get_files_service(self) -> "GoogleAPIResource":
-        from google.auth import default, exceptions
+        from google.auth import exceptions
         from google.oauth2 import service_account
         from googleapiclient.discovery import build
         from googleapiclient.errors import HttpError
 
-        # Service account key can be a dict or a file path(str)
-        # But the dict may come in as a string
-        if isinstance(self.access_config.service_account_key, str):
-            key_path = json_to_dict(self.access_config.service_account_key)
-        elif isinstance(self.access_config.service_account_key, dict):
-            key_path = self.access_config.service_account_key
-        else:
-            raise TypeError(
-                f"access_config.service_account_key must be "
-                f"str or dict, got: {type(self.access_config.service_account_key)}"
-            )
+        access_config = self.access_config.get_secret_value()
+        key_data = access_config.get_service_account_key()
 
         try:
-            if isinstance(key_path, dict):
-                creds = service_account.Credentials.from_service_account_info(key_path)
-            elif isinstance(key_path, str):
-                os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = key_path
-                creds, _ = default()
-            else:
-                raise ValueError(
-                    f"key path not recognized as a dictionary or a file path: "
-                    f"[{type(key_path)}] {key_path}",
-                )
+            creds = service_account.Credentials.from_service_account_info(key_data)
             service = build("drive", "v3", credentials=creds)
             return service.files()
 
@@ -86,7 +94,6 @@ class GoogleDriveConnectionConfig(ConnectionConfig):
             raise ValueError("The provided API key is invalid.")
 
 
-@dataclass
 class GoogleDriveIndexerConfig(IndexerConfig):
     extensions: Optional[list[str]] = None
     recursive: bool = False
@@ -268,7 +275,6 @@ class GoogleDriveIndexer(Indexer):
             yield f
 
 
-@dataclass
 class GoogleDriveDownloaderConfig(DownloaderConfig):
     pass
 

unstructured_ingest/v2/processes/connectors/kdbai.py

@@ -0,0 +1,170 @@
+import json
+import uuid
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional
+
+import numpy as np
+import pandas as pd
+from pydantic import Field, Secret
+
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.utils.data_prep import 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 kdbai_client import Session, Table
+
+CONNECTOR_TYPE = "kdbai"
+
+
+class KdbaiAccessConfig(AccessConfig):
+    api_key: Optional[str] = Field(
+        default=None,
+        description="A string for the api-key, can be left empty "
+        "when connecting to local KDBAI instance.",
+    )
+
+
+SecretKdbaiAccessConfig = Secret[KdbaiAccessConfig]
+
+
+class KdbaiConnectionConfig(ConnectionConfig):
+    access_config: SecretKdbaiAccessConfig = Field(
+        default=SecretKdbaiAccessConfig(secret_value=KdbaiAccessConfig())
+    )
+    endpoint: str = Field(
+        default="http://localhost:8082", description="Endpoint url where KDBAI is hosted."
+    )
+
+    @requires_dependencies(["kdbai_client"], extras="kdbai")
+    def get_session(self) -> "Session":
+        from kdbai_client import Session
+
+        return Session(
+            api_key=self.access_config.get_secret_value().api_key, endpoint=self.endpoint
+        )
+
+
+class KdbaiUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class KdbaiUploadStager(UploadStager):
+    upload_stager_config: KdbaiUploadStagerConfig = field(default_factory=KdbaiUploadStagerConfig)
+
+    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)
+        output_path = Path(output_dir) / Path(f"{output_filename}.json")
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        data = []
+        for element in elements_contents:
+            data.append(
+                {
+                    "id": str(uuid.uuid4()),
+                    "element_id": element.get("element_id"),
+                    "document": element.pop("text", None),
+                    "embeddings": element.get("embeddings"),
+                    "metadata": flatten_dict(
+                        dictionary=element.get("metadata"),
+                        flatten_lists=True,
+                        remove_none=True,
+                    ),
+                }
+            )
+        logger.debug(f"writing {len(data)} elements to {output_path}")
+        with output_path.open("w") as output_file:
+            json.dump(data, output_file, indent=2)
+        return output_path
+
+
+class KdbaiUploaderConfig(UploaderConfig):
+    table_name: str = Field(description="The name of the KDBAI table to write into.")
+    batch_size: int = Field(default=100, description="Number of records per batch")
+
+
+@dataclass
+class KdbaiUploader(Uploader):
+    connection_config: KdbaiConnectionConfig
+    upload_config: KdbaiUploaderConfig
+    connector_type: str = field(default=CONNECTOR_TYPE, init=False)
+
+    def precheck(self) -> None:
+        try:
+            self.get_table()
+        except Exception as e:
+            logger.error(f"Failed to validate connection {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def get_table(self) -> "Table":
+        session: Session = self.connection_config.get_session()
+        table = session.table(self.upload_config.table_name)
+        return table
+
+    def upsert_batch(self, batch: pd.DataFrame):
+        table = self.get_table()
+        table.insert(data=batch)
+
+    def process_dataframe(self, df: pd.DataFrame):
+        logger.debug(
+            f"uploading {len(df)} entries to {self.connection_config.endpoint} "
+            f"db in table {self.upload_config.table_name}"
+        )
+        for _, batch_df in df.groupby(np.arange(len(df)) // self.upload_config.batch_size):
+            self.upsert_batch(batch=batch_df)
+
+    def process_csv(self, csv_paths: list[Path]):
+        logger.debug(f"uploading content from {len(csv_paths)} csv files")
+        df = pd.concat((pd.read_csv(path) for path in csv_paths), ignore_index=True)
+        self.process_dataframe(df=df)
+
+    def process_json(self, json_paths: list[Path]):
+        logger.debug(f"uploading content from {len(json_paths)} json files")
+        all_records = []
+        for p in json_paths:
+            with open(p) as json_file:
+                all_records.extend(json.load(json_file))
+
+        df = pd.DataFrame(data=all_records)
+        self.process_dataframe(df=df)
+
+    def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
+        csv_paths = [c.path for c in contents if c.path.suffix == ".csv"]
+        if csv_paths:
+            self.process_csv(csv_paths=csv_paths)
+        json_paths = [c.path for c in contents if c.path.suffix == ".json"]
+        if json_paths:
+            self.process_json(json_paths=json_paths)
+
+
+kdbai_destination_entry = DestinationRegistryEntry(
+    connection_config=KdbaiConnectionConfig,
+    uploader=KdbaiUploader,
+    uploader_config=KdbaiUploaderConfig,
+    upload_stager=KdbaiUploadStager,
+    upload_stager_config=KdbaiUploadStagerConfig,
+)

unstructured_ingest/v2/processes/connectors/local.py

@@ -5,6 +5,8 @@ from pathlib import Path
 from time import time
 from typing import Any, Generator
 
+from pydantic import Field, Secret
+
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
     ConnectionConfig,
@@ -29,20 +31,28 @@ from unstructured_ingest.v2.processes.connector_registry import (
 CONNECTOR_TYPE = "local"
 
 
-@dataclass
 class LocalAccessConfig(AccessConfig):
     pass
 
 
-@dataclass
+SecretLocalAccessConfig = Secret[LocalAccessConfig]
+
+
 class LocalConnectionConfig(ConnectionConfig):
-    access_config: LocalAccessConfig = field(default_factory=lambda: LocalAccessConfig())
+    access_config: SecretLocalAccessConfig = Field(
+        default_factory=lambda: SecretLocalAccessConfig(secret_value=LocalAccessConfig())
+    )
 
 
-@dataclass
 class LocalIndexerConfig(IndexerConfig):
-    input_path: str
-    recursive: bool = False
+    input_path: Path = Field(
+        description="Path to the location in the local file system that will be processed."
+    )
+    recursive: bool = Field(
+        default=False,
+        description="Recursively download files in their respective folders "
+        "otherwise stop at the files in provided folder level.",
+    )
 
     @property
     def path(self) -> Path:
@@ -61,9 +71,12 @@ class LocalIndexer(Indexer):
         input_path = self.index_config.path
         if input_path.is_file():
             return [Path(s) for s in glob.glob(f"{self.index_config.path}")]
+        files = []
         if self.index_config.recursive:
-            return list(input_path.rglob("*"))
-        return list(input_path.glob("*"))
+            files.extend(list(input_path.rglob("*")))
+        else:
+            files.extend(list(input_path.glob("*")))
+        return [f for f in files if f.is_file()]
 
     def get_file_metadata(self, path: Path) -> FileDataSourceMetadata:
         stats = path.stat()
@@ -122,7 +135,6 @@ class LocalIndexer(Indexer):
             yield file_data
 
 
-@dataclass
 class LocalDownloaderConfig(DownloaderConfig):
     pass
 
@@ -130,10 +142,8 @@ class LocalDownloaderConfig(DownloaderConfig):
 @dataclass
 class LocalDownloader(Downloader):
     connector_type: str = CONNECTOR_TYPE
-    connection_config: LocalConnectionConfig = field(
-        default_factory=lambda: LocalConnectionConfig()
-    )
-    download_config: LocalDownloaderConfig = field(default_factory=lambda: LocalDownloaderConfig())
+    connection_config: LocalConnectionConfig = field(default_factory=LocalConnectionConfig)
+    download_config: LocalDownloaderConfig = field(default_factory=LocalDownloaderConfig)
 
     def get_download_path(self, file_data: FileData) -> Path:
         return Path(file_data.source_identifiers.fullpath)
@@ -144,9 +154,10 @@ class LocalDownloader(Downloader):
         )
 
 
-@dataclass
 class LocalUploaderConfig(UploaderConfig):
-    output_dir: str = field(default="structured-output")
+    output_dir: str = Field(
+        default="structured-output", description="Local path to write partitioned output to"
+    )
 
     @property
     def output_path(self) -> Path:
@@ -160,7 +171,7 @@ class LocalUploaderConfig(UploaderConfig):
 @dataclass
 class LocalUploader(Uploader):
     connector_type: str = CONNECTOR_TYPE
-    upload_config: LocalUploaderConfig = field(default_factory=lambda: LocalUploaderConfig())
+    upload_config: LocalUploaderConfig = field(default_factory=LocalUploaderConfig)
     connection_config: LocalConnectionConfig = field(
         default_factory=lambda: LocalConnectionConfig()
     )

unstructured_ingest/v2/processes/connectors/milvus.py

@@ -6,8 +6,8 @@ from typing import TYPE_CHECKING, Any, Optional, Union
 
 import pandas as pd
 from dateutil import parser
+from pydantic import Field, Secret
 
-from unstructured_ingest.enhanced_dataclass import enhanced_field
 from unstructured_ingest.error import WriteError
 from unstructured_ingest.utils.data_prep import flatten_dict
 from unstructured_ingest.utils.dep_check import requires_dependencies
@@ -32,24 +32,28 @@ if TYPE_CHECKING:
 CONNECTOR_TYPE = "milvus"
 
 
-@dataclass
 class MilvusAccessConfig(AccessConfig):
-    password: Optional[str] = None
-    token: Optional[str] = None
+    password: Optional[str] = Field(default=None, description="Milvus password")
+    token: Optional[str] = Field(default=None, description="Milvus access token")
+
+
+SecretMilvusAccessConfig = Secret[MilvusAccessConfig]
 
 
-@dataclass
 class MilvusConnectionConfig(ConnectionConfig):
-    access_config: MilvusAccessConfig = enhanced_field(
-        sensitive=True, default_factory=lambda: MilvusAccessConfig()
+    access_config: SecretMilvusAccessConfig = Field(
+        default_factory=lambda: SecretMilvusAccessConfig(secret_value=MilvusAccessConfig())
     )
-    uri: Optional[str] = None
-    user: Optional[str] = None
-    db_name: Optional[str] = None
+    uri: Optional[str] = Field(
+        default=None, description="Milvus uri", examples=["http://localhost:19530"]
+    )
+    user: Optional[str] = Field(default=None, description="Milvus user")
+    db_name: Optional[str] = Field(default=None, description="Milvus database name")
 
     def get_connection_kwargs(self) -> dict[str, Any]:
-        access_config_dict = self.access_config.to_dict()
-        connection_config_dict = self.to_dict()
+        access_config = self.access_config.get_secret_value()
+        access_config_dict = access_config.dict()
+        connection_config_dict = self.dict()
         connection_config_dict.pop("access_config", None)
         connection_config_dict.update(access_config_dict)
         # Drop any that were not set explicitly
@@ -63,7 +67,6 @@ class MilvusConnectionConfig(ConnectionConfig):
         return MilvusClient(**self.get_connection_kwargs())
 
 
-@dataclass
 class MilvusUploadStagerConfig(UploadStagerConfig):
     pass
 
@@ -130,10 +133,11 @@ class MilvusUploadStager(UploadStager):
         return output_path
 
 
-@dataclass
 class MilvusUploaderConfig(UploaderConfig):
-    collection_name: str
-    num_of_processes: int = 4
+    collection_name: str = Field(description="Milvus collections to write to")
+    num_processes: int = Field(
+        default=4, description="number of processes to use when writing to support parallel writes"
+    )
 
 
 @dataclass
@@ -180,13 +184,13 @@ class MilvusUploader(Uploader):
         self.insert_results(data=data)
 
     def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
-        if self.upload_config.num_of_processes == 1:
+        if self.upload_config.num_processes == 1:
             for content in contents:
                 self.upload(content=content)
 
         else:
             with mp.Pool(
-                processes=self.upload_config.num_of_processes,
+                processes=self.upload_config.num_processes,
             ) as pool:
                 pool.map(self.upload, contents)
 

unstructured_ingest/v2/processes/connectors/mongodb.py

@@ -3,9 +3,9 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Optional
 
-from unstructured.__version__ import __version__ as unstructured_version
+from pydantic import Field, Secret
 
-from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.__version__ import __version__ as unstructured_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
@@ -31,25 +31,28 @@ CONNECTOR_TYPE = "mongodb"
 SERVER_API_VERSION = "1"
 
 
-@dataclass
 class MongoDBAccessConfig(AccessConfig):
-    uri: Optional[str] = None
+    uri: Optional[str] = Field(default=None, description="URI to user when connecting")
+
+
+SecretMongoDBAccessConfig = Secret[MongoDBAccessConfig]
 
 
-@dataclass
 class MongoDBConnectionConfig(ConnectionConfig):
-    access_config: MongoDBAccessConfig = enhanced_field(
-        sensitive=True, default_factory=MongoDBAccessConfig
+    access_config: SecretMongoDBAccessConfig = Field(
+        default_factory=lambda: SecretMongoDBAccessConfig(secret_value=MongoDBAccessConfig())
     )
-    host: Optional[str] = None
-    database: Optional[str] = None
-    collection: Optional[str] = None
-    port: int = 27017
-    batch_size: int = 100
-    connector_type: str = CONNECTOR_TYPE
+    host: Optional[str] = Field(
+        default=None,
+        description="hostname or IP address or Unix domain socket path of a single mongod or "
+        "mongos instance to connect to, or a list of hostnames",
+    )
+    database: Optional[str] = Field(default=None, description="database name to connect to")
+    collection: Optional[str] = Field(default=None, description="collection name to connect to")
+    port: int = Field(default=27017)
+    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
 
 
-@dataclass
 class MongoDBUploadStagerConfig(UploadStagerConfig):
     pass
 
@@ -77,9 +80,8 @@ class MongoDBUploadStager(UploadStager):
         return output_path
 
 
-@dataclass
 class MongoDBUploaderConfig(UploaderConfig):
-    batch_size: int = 100
+    batch_size: int = Field(default=100, description="Number of records per batch")
 
 
 @dataclass
@@ -102,9 +104,11 @@ class MongoDBUploader(Uploader):
         from pymongo.driver_info import DriverInfo
         from pymongo.server_api import ServerApi
 
-        if self.connection_config.access_config.uri:
+        access_config = self.connection_config.access_config.get_secret_value()
+
+        if access_config.uri:
             return MongoClient(
-                self.connection_config.access_config.uri,
+                access_config.uri,
                 server_api=ServerApi(version=SERVER_API_VERSION),
                 driver=DriverInfo(name="unstructured", version=unstructured_version),
             )

unstructured_ingest/v2/processes/connectors/onedrive.py

@@ -1,12 +1,12 @@
 import json
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from pathlib import Path
 from time import time
 from typing import TYPE_CHECKING, Any, Generator, Optional
 
 from dateutil import parser
+from pydantic import Field, Secret
 
-from unstructured_ingest.enhanced_dataclass import enhanced_field
 from unstructured_ingest.error import SourceConnectionError, SourceConnectionNetworkError
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import (
@@ -35,18 +35,23 @@ CONNECTOR_TYPE = "onedrive"
 MAX_MB_SIZE = 512_000_000
 
 
-@dataclass
 class OnedriveAccessConfig(AccessConfig):
-    client_cred: str
+    client_cred: str = Field(description="Microsoft App client secret")
 
 
-@dataclass
 class OnedriveConnectionConfig(ConnectionConfig):
-    client_id: str
-    user_pname: str
-    tenant: str = field(repr=False)
-    authority_url: Optional[str] = field(repr=False, default="https://login.microsoftonline.com")
-    access_config: OnedriveAccessConfig = enhanced_field(sensitive=True)
+    client_id: str = Field(description="Microsoft app client ID")
+    user_pname: str = Field(description="User principal name, usually is your Azure AD email.")
+    tenant: str = Field(
+        repr=False, description="ID or domain name associated with your Azure AD instance"
+    )
+    authority_url: Optional[str] = Field(
+        repr=False,
+        default="https://login.microsoftonline.com",
+        examples=["https://login.microsoftonline.com"],
+        description="Authentication token provider for Microsoft apps",
+    )
+    access_config: Secret[OnedriveAccessConfig]
 
     @requires_dependencies(["msal"], extras="onedrive")
     def get_token(self):
@@ -56,7 +61,7 @@ class OnedriveConnectionConfig(ConnectionConfig):
             app = ConfidentialClientApplication(
                 authority=f"{self.authority_url}/{self.tenant}",
                 client_id=self.client_id,
-                client_credential=self.access_config.client_cred,
+                client_credential=self.access_config.get_secret_value().client_cred,
             )
             token = app.acquire_token_for_client(scopes=["https://graph.microsoft.com/.default"])
         except ValueError as exc:
@@ -76,9 +81,8 @@ class OnedriveConnectionConfig(ConnectionConfig):
         return client
 
 
-@dataclass
 class OnedriveIndexerConfig(IndexerConfig):
-    path: Optional[str] = field(default="")
+    path: Optional[str] = Field(default="")
     recursive: bool = False
 
 
@@ -171,7 +175,6 @@ class OnedriveIndexer(Indexer):
             yield file_data
 
 
-@dataclass
 class OnedriveDownloaderConfig(DownloaderConfig):
     pass