unstructured-ingest 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/lancedb/aws.py

@@ -0,0 +1,43 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.interfaces.connector import AccessConfig
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connectors.lancedb.lancedb import (
+    LanceDBRemoteConnectionConfig,
+    LanceDBUploader,
+    LanceDBUploaderConfig,
+    LanceDBUploadStager,
+    LanceDBUploadStagerConfig,
+)
+
+CONNECTOR_TYPE = "lancedb_aws"
+
+
+class LanceDBS3AccessConfig(AccessConfig):
+    aws_access_key_id: str = Field(description="The AWS access key ID to use.")
+    aws_secret_access_key: str = Field(description="The AWS secret access key to use.")
+
+
+class LanceDBS3ConnectionConfig(LanceDBRemoteConnectionConfig):
+    access_config: Secret[LanceDBS3AccessConfig]
+
+    def get_storage_options(self) -> dict:
+        return {**self.access_config.get_secret_value().model_dump(), "timeout": self.timeout}
+
+
+@dataclass
+class LanceDBS3Uploader(LanceDBUploader):
+    upload_config: LanceDBUploaderConfig
+    connection_config: LanceDBS3ConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+lancedb_aws_destination_entry = DestinationRegistryEntry(
+    connection_config=LanceDBS3ConnectionConfig,
+    uploader=LanceDBS3Uploader,
+    uploader_config=LanceDBUploaderConfig,
+    upload_stager_config=LanceDBUploadStagerConfig,
+    upload_stager=LanceDBUploadStager,
+)

unstructured_ingest/v2/processes/connectors/lancedb/azure.py

@@ -0,0 +1,43 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.interfaces.connector import AccessConfig
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connectors.lancedb.lancedb import (
+    LanceDBRemoteConnectionConfig,
+    LanceDBUploader,
+    LanceDBUploaderConfig,
+    LanceDBUploadStager,
+    LanceDBUploadStagerConfig,
+)
+
+CONNECTOR_TYPE = "lancedb_azure"
+
+
+class LanceDBAzureAccessConfig(AccessConfig):
+    azure_storage_account_name: str = Field(description="The name of the azure storage account.")
+    azure_storage_account_key: str = Field(description="The serialized azure service account key.")
+
+
+class LanceDBAzureConnectionConfig(LanceDBRemoteConnectionConfig):
+    access_config: Secret[LanceDBAzureAccessConfig]
+
+    def get_storage_options(self) -> dict:
+        return {**self.access_config.get_secret_value().model_dump(), "timeout": self.timeout}
+
+
+@dataclass
+class LanceDBAzureUploader(LanceDBUploader):
+    upload_config: LanceDBUploaderConfig
+    connection_config: LanceDBAzureConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+lancedb_azure_destination_entry = DestinationRegistryEntry(
+    connection_config=LanceDBAzureConnectionConfig,
+    uploader=LanceDBAzureUploader,
+    uploader_config=LanceDBUploaderConfig,
+    upload_stager_config=LanceDBUploadStagerConfig,
+    upload_stager=LanceDBUploadStager,
+)

unstructured_ingest/v2/processes/connectors/lancedb/gcp.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.interfaces.connector import AccessConfig
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connectors.lancedb.lancedb import (
+    LanceDBRemoteConnectionConfig,
+    LanceDBUploader,
+    LanceDBUploaderConfig,
+    LanceDBUploadStager,
+    LanceDBUploadStagerConfig,
+)
+
+CONNECTOR_TYPE = "lancedb_gcs"
+
+
+class LanceDBGCSAccessConfig(AccessConfig):
+    google_service_account_key: str = Field(
+        description="The serialized google service account key."
+    )
+
+
+class LanceDBGCSConnectionConfig(LanceDBRemoteConnectionConfig):
+    access_config: Secret[LanceDBGCSAccessConfig]
+
+    def get_storage_options(self) -> dict:
+        return {**self.access_config.get_secret_value().model_dump(), "timeout": self.timeout}
+
+
+@dataclass
+class LanceDBGSPUploader(LanceDBUploader):
+    upload_config: LanceDBUploaderConfig
+    connection_config: LanceDBGCSConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+lancedb_gcp_destination_entry = DestinationRegistryEntry(
+    connection_config=LanceDBGCSConnectionConfig,
+    uploader=LanceDBGSPUploader,
+    uploader_config=LanceDBUploaderConfig,
+    upload_stager_config=LanceDBUploadStagerConfig,
+    upload_stager=LanceDBUploadStager,
+)

unstructured_ingest/v2/processes/connectors/lancedb/lancedb.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+import asyncio
+import json
+from abc import ABC, abstractmethod
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional
+
+import pandas as pd
+from pydantic import Field
+
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.data_prep import flatten_dict
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces.connector import ConnectionConfig
+from unstructured_ingest.v2.interfaces.file_data import FileData
+from unstructured_ingest.v2.interfaces.upload_stager import UploadStager, UploadStagerConfig
+from unstructured_ingest.v2.interfaces.uploader import Uploader, UploaderConfig
+
+CONNECTOR_TYPE = "lancedb"
+
+if TYPE_CHECKING:
+    from lancedb import AsyncConnection
+    from lancedb.table import AsyncTable
+
+
+class LanceDBConnectionConfig(ConnectionConfig, ABC):
+    uri: str = Field(description="The uri of the database.")
+
+    @abstractmethod
+    def get_storage_options(self) -> Optional[dict[str, str]]:
+        raise NotImplementedError
+
+    @asynccontextmanager
+    @requires_dependencies(["lancedb"], extras="lancedb")
+    @DestinationConnectionError.wrap
+    async def get_async_connection(self) -> AsyncGenerator["AsyncConnection", None]:
+        import lancedb
+
+        connection = await lancedb.connect_async(
+            self.uri,
+            storage_options=self.get_storage_options(),
+        )
+        try:
+            yield connection
+        finally:
+            connection.close()
+
+
+class LanceDBRemoteConnectionConfig(LanceDBConnectionConfig):
+    timeout: str = Field(
+        default="30s",
+        description=(
+            "Timeout for the entire request, from connection until the response body has finished"
+            "in a [0-9]+(ns|us|ms|[smhdwy]) format."
+        ),
+        pattern=r"[0-9]+(ns|us|ms|[smhdwy])",
+    )
+
+
+class LanceDBUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class LanceDBUploadStager(UploadStager):
+    upload_stager_config: LanceDBUploadStagerConfig = field(
+        default_factory=LanceDBUploadStagerConfig
+    )
+
+    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: list[dict] = json.load(elements_file)
+
+        df = pd.DataFrame(
+            [
+                self._conform_element_contents(element_contents)
+                for element_contents in elements_contents
+            ]
+        )
+
+        output_path = (output_dir / output_filename).with_suffix(".feather")
+        df.to_feather(output_path)
+
+        return output_path
+
+    def _conform_element_contents(self, element: dict) -> dict:
+        return {
+            "vector": element.pop("embeddings", None),
+            **flatten_dict(element, separator="-"),
+        }
+
+
+class LanceDBUploaderConfig(UploaderConfig):
+    table_name: str = Field(description="The name of the table.")
+
+
+@dataclass
+class LanceDBUploader(Uploader):
+    upload_config: LanceDBUploaderConfig
+    connection_config: LanceDBConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @DestinationConnectionError.wrap
+    def precheck(self):
+        async def _precheck() -> None:
+            async with self.connection_config.get_async_connection() as conn:
+                table = await conn.open_table(self.upload_config.table_name)
+                table.close()
+
+        asyncio.run(_precheck())
+
+    @asynccontextmanager
+    async def get_table(self) -> AsyncGenerator["AsyncTable", None]:
+        async with self.connection_config.get_async_connection() as conn:
+            table = await conn.open_table(self.upload_config.table_name)
+            try:
+                yield table
+            finally:
+                table.close()
+
+    async def run_async(self, path, file_data, **kwargs):
+        df = pd.read_feather(path)
+        async with self.get_table() as table:
+            schema = await table.schema()
+            df = self._fit_to_schema(df, schema)
+            await table.add(data=df)
+
+    def _fit_to_schema(self, df: pd.DataFrame, schema) -> pd.DataFrame:
+        columns = set(df.columns)
+        schema_fields = set(schema.names)
+        columns_to_drop = columns - schema_fields
+        missing_columns = schema_fields - columns
+
+        if columns_to_drop:
+            logger.info(
+                "Following columns will be dropped to match the table's schema: "
+                f"{', '.join(columns_to_drop)}"
+            )
+        if missing_columns:
+            logger.info(
+                "Following null filled columns will be added to match the table's schema:"
+                f" {', '.join(missing_columns)} "
+            )
+
+        df = df.drop(columns=columns_to_drop)
+
+        for column in missing_columns:
+            df[column] = pd.Series()
+
+        return df

unstructured_ingest/v2/processes/connectors/lancedb/local.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.interfaces.connector import AccessConfig
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connectors.lancedb.lancedb import (
+    LanceDBConnectionConfig,
+    LanceDBUploader,
+    LanceDBUploaderConfig,
+    LanceDBUploadStager,
+    LanceDBUploadStagerConfig,
+)
+
+CONNECTOR_TYPE = "lancedb_local"
+
+
+class LanceDBLocalAccessConfig(AccessConfig):
+    pass
+
+
+class LanceDBLocalConnectionConfig(LanceDBConnectionConfig):
+    access_config: Secret[LanceDBLocalAccessConfig] = Field(
+        default_factory=LanceDBLocalAccessConfig, validate_default=True
+    )
+
+    def get_storage_options(self) -> None:
+        return None
+
+
+@dataclass
+class LanceDBLocalUploader(LanceDBUploader):
+    upload_config: LanceDBUploaderConfig
+    connection_config: LanceDBLocalConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+lancedb_local_destination_entry = DestinationRegistryEntry(
+    connection_config=LanceDBLocalConnectionConfig,
+    uploader=LanceDBLocalUploader,
+    uploader_config=LanceDBUploaderConfig,
+    upload_stager_config=LanceDBUploadStagerConfig,
+    upload_stager=LanceDBUploadStager,
+)

unstructured_ingest/v2/processes/connectors/milvus.py

@@ -1,7 +1,8 @@
 import json
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Generator, Optional, Union
 
 import pandas as pd
 from dateutil import parser
@@ -10,6 +11,7 @@ from pydantic import Field, Secret
 from unstructured_ingest.error import WriteError
 from unstructured_ingest.utils.data_prep import flatten_dict
 from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.constants import RECORD_ID_LABEL
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
     ConnectionConfig,
@@ -90,24 +92,27 @@ class MilvusUploadStager(UploadStager):
             pass
         return parser.parse(date_string).timestamp()
 
-    def conform_dict(self, data: dict) -> None:
-        if self.upload_stager_config.flatten_metadata and (metadata := data.pop("metadata", None)):
-            data.update(flatten_dict(metadata, keys_to_omit=["data_source_record_locator"]))
+    def conform_dict(self, data: dict, file_data: FileData) -> dict:
+        working_data = data.copy()
+        if self.upload_stager_config.flatten_metadata and (
+            metadata := working_data.pop("metadata", None)
+        ):
+            working_data.update(flatten_dict(metadata, keys_to_omit=["data_source_record_locator"]))
 
         # TODO: milvus sdk doesn't seem to support defaults via the schema yet,
         #  remove once that gets updated
         defaults = {"is_continuation": False}
         for default in defaults:
-            if default not in data:
-                data[default] = defaults[default]
+            if default not in working_data:
+                working_data[default] = defaults[default]
 
         if self.upload_stager_config.fields_to_include:
-            data_keys = set(data.keys())
+            data_keys = set(working_data.keys())
             for data_key in data_keys:
                 if data_key not in self.upload_stager_config.fields_to_include:
-                    data.pop(data_key)
+                    working_data.pop(data_key)
             for field_include_key in self.upload_stager_config.fields_to_include:
-                if field_include_key not in data:
+                if field_include_key not in working_data:
                     raise KeyError(f"Field '{field_include_key}' is missing in data!")
 
         datetime_columns = [
@@ -120,11 +125,15 @@ class MilvusUploadStager(UploadStager):
         json_dumps_fields = ["languages", "data_source_permissions_data"]
 
         for datetime_column in datetime_columns:
-            if datetime_column in data:
-                data[datetime_column] = self.parse_date_string(data[datetime_column])
+            if datetime_column in working_data:
+                working_data[datetime_column] = self.parse_date_string(
+                    working_data[datetime_column]
+                )
         for json_dumps_field in json_dumps_fields:
-            if json_dumps_field in data:
-                data[json_dumps_field] = json.dumps(data[json_dumps_field])
+            if json_dumps_field in working_data:
+                working_data[json_dumps_field] = json.dumps(working_data[json_dumps_field])
+        working_data[RECORD_ID_LABEL] = file_data.identifier
+        return working_data
 
     def run(
         self,
@@ -136,18 +145,27 @@ class MilvusUploadStager(UploadStager):
     ) -> Path:
         with open(elements_filepath) as elements_file:
             elements_contents: list[dict[str, Any]] = json.load(elements_file)
-        for element in elements_contents:
-            self.conform_dict(data=element)
-
-        output_path = Path(output_dir) / Path(f"{output_filename}.json")
+        new_content = [
+            self.conform_dict(data=element, file_data=file_data) for element in elements_contents
+        ]
+        output_filename_path = Path(output_filename)
+        if output_filename_path.suffix == ".json":
+            output_path = Path(output_dir) / output_filename_path
+        else:
+            output_path = Path(output_dir) / output_filename_path.with_suffix(".json")
         output_path.parent.mkdir(parents=True, exist_ok=True)
         with output_path.open("w") as output_file:
-            json.dump(elements_contents, output_file, indent=2)
+            json.dump(new_content, output_file, indent=2)
         return output_path
 
 
 class MilvusUploaderConfig(UploaderConfig):
+    db_name: Optional[str] = Field(default=None, description="Milvus database name")
     collection_name: str = Field(description="Milvus collections to write to")
+    record_id_key: str = Field(
+        default=RECORD_ID_LABEL,
+        description="searchable key to find entries for the same record on previous runs",
+    )
 
 
 @dataclass
@@ -156,6 +174,16 @@ class MilvusUploader(Uploader):
     upload_config: MilvusUploaderConfig
     connector_type: str = CONNECTOR_TYPE
 
+    @contextmanager
+    def get_client(self) -> Generator["MilvusClient", None, None]:
+        client = self.connection_config.get_client()
+        if db_name := self.upload_config.db_name:
+            client.using_database(db_name=db_name)
+        try:
+            yield client
+        finally:
+            client.close()
+
     def upload(self, content: UploadContent) -> None:
         file_extension = content.path.suffix
         if file_extension == ".json":
@@ -165,23 +193,39 @@ class MilvusUploader(Uploader):
         else:
             raise ValueError(f"Unsupported file extension: {file_extension}")
 
+    def delete_by_record_id(self, file_data: FileData) -> None:
+        logger.info(
+            f"deleting any content with metadata {RECORD_ID_LABEL}={file_data.identifier} "
+            f"from milvus collection {self.upload_config.collection_name}"
+        )
+        with self.get_client() as client:
+            delete_filter = f'{self.upload_config.record_id_key} == "{file_data.identifier}"'
+            resp = client.delete(
+                collection_name=self.upload_config.collection_name, filter=delete_filter
+            )
+            logger.info(
+                "deleted {} records from milvus collection {}".format(
+                    resp["delete_count"], self.upload_config.collection_name
+                )
+            )
+
     @requires_dependencies(["pymilvus"], extras="milvus")
     def insert_results(self, data: Union[dict, list[dict]]):
         from pymilvus import MilvusException
 
-        logger.debug(
+        logger.info(
             f"uploading {len(data)} entries to {self.connection_config.db_name} "
             f"db in collection {self.upload_config.collection_name}"
         )
-        client = self.connection_config.get_client()
+        with self.get_client() as client:
 
-        try:
-            res = client.insert(collection_name=self.upload_config.collection_name, data=data)
-        except MilvusException as milvus_exception:
-            raise WriteError("failed to upload records to milvus") from milvus_exception
-        if "err_count" in res and isinstance(res["err_count"], int) and res["err_count"] > 0:
-            err_count = res["err_count"]
-            raise WriteError(f"failed to upload {err_count} docs")
+            try:
+                res = client.insert(collection_name=self.upload_config.collection_name, data=data)
+            except MilvusException as milvus_exception:
+                raise WriteError("failed to upload records to milvus") from milvus_exception
+            if "err_count" in res and isinstance(res["err_count"], int) and res["err_count"] > 0:
+                err_count = res["err_count"]
+                raise WriteError(f"failed to upload {err_count} docs")
 
     def upload_csv(self, content: UploadContent) -> None:
         df = pd.read_csv(content.path)
@@ -194,6 +238,7 @@ class MilvusUploader(Uploader):
         self.insert_results(data=data)
 
     def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        self.delete_by_record_id(file_data=file_data)
         self.upload(content=UploadContent(path=path, file_data=file_data))