unstructured-ingest 0.3.13__py3-none-any.whl

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,169 @@
+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.constants import RECORD_ID_LABEL
+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
+
+        with await lancedb.connect_async(
+            self.uri,
+            storage_options=self.get_storage_options(),
+        ) as connection:
+            yield connection
+
+
+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_dict(element_dict=element_dict, file_data=file_data)
+                for element_dict in elements_contents
+            ]
+        )
+
+        output_path = (output_dir / output_filename).with_suffix(".feather")
+        df.to_feather(output_path)
+
+        return output_path
+
+    def conform_dict(self, element_dict: dict, file_data: FileData) -> dict:
+        data = element_dict.copy()
+        return {
+            "vector": data.pop("embeddings", None),
+            RECORD_ID_LABEL: file_data.identifier,
+            **flatten_dict(data, 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)
+            if RECORD_ID_LABEL not in schema.names:
+                logger.warning(
+                    f"Designated table doesn't contain {RECORD_ID_LABEL} column of type"
+                    " string which is required to support overwriting updates on subsequent"
+                    " uploads of the same record. New rows will be appended instead."
+                )
+            else:
+                await table.delete(f'{RECORD_ID_LABEL} = "{file_data.identifier}"')
+            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/local.py

@@ -0,0 +1,217 @@
+import glob
+import json
+import shutil
+from dataclasses import dataclass, field
+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,
+    Downloader,
+    DownloaderConfig,
+    DownloadResponse,
+    FileData,
+    FileDataSourceMetadata,
+    Indexer,
+    IndexerConfig,
+    SourceIdentifiers,
+    Uploader,
+    UploaderConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+
+CONNECTOR_TYPE = "local"
+
+
+class LocalAccessConfig(AccessConfig):
+    pass
+
+
+class LocalConnectionConfig(ConnectionConfig):
+    access_config: Secret[LocalAccessConfig] = Field(
+        default=LocalAccessConfig(), validate_default=True
+    )
+
+
+class LocalIndexerConfig(IndexerConfig):
+    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:
+        return Path(self.input_path).resolve()
+
+
+@dataclass
+class LocalIndexer(Indexer):
+    index_config: LocalIndexerConfig
+    connection_config: LocalConnectionConfig = field(
+        default_factory=lambda: LocalConnectionConfig()
+    )
+    connector_type: str = CONNECTOR_TYPE
+
+    def list_files(self) -> list[Path]:
+        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:
+            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()
+        try:
+            date_modified = str(stats.st_mtime)
+        except Exception as e:
+            logger.warning(f"Couldn't detect date modified: {e}")
+            date_modified = None
+
+        try:
+            date_created = str(stats.st_birthtime)
+        except Exception as e:
+            logger.warning(f"Couldn't detect date created: {e}")
+            date_created = None
+
+        try:
+            mode = stats.st_mode
+            permissions_data = [{"mode": mode}]
+        except Exception as e:
+            logger.warning(f"Couldn't detect file mode: {e}")
+            permissions_data = None
+
+        try:
+            filesize_bytes = stats.st_size
+        except Exception as e:
+            logger.warning(f"Couldn't detect file size: {e}")
+            filesize_bytes = None
+
+        return FileDataSourceMetadata(
+            date_modified=date_modified,
+            date_created=date_created,
+            date_processed=str(time()),
+            permissions_data=permissions_data,
+            record_locator={"path": str(path.resolve())},
+            filesize_bytes=filesize_bytes,
+        )
+
+    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+        for file_path in self.list_files():
+            file_data = FileData(
+                identifier=str(file_path.resolve()),
+                connector_type=CONNECTOR_TYPE,
+                source_identifiers=SourceIdentifiers(
+                    fullpath=str(file_path.resolve()),
+                    filename=file_path.name,
+                    rel_path=(
+                        str(file_path.resolve()).replace(str(self.index_config.path.resolve()), "")[
+                            1:
+                        ]
+                        if not self.index_config.path.is_file()
+                        else self.index_config.path.name
+                    ),
+                ),
+                metadata=self.get_file_metadata(path=file_path),
+            )
+            yield file_data
+
+
+class LocalDownloaderConfig(DownloaderConfig):
+    pass
+
+
+@dataclass
+class LocalDownloader(Downloader):
+    connector_type: str = CONNECTOR_TYPE
+    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)
+
+    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
+        return DownloadResponse(
+            file_data=file_data, path=Path(file_data.source_identifiers.fullpath)
+        )
+
+
+class LocalUploaderConfig(UploaderConfig):
+    output_dir: str = Field(
+        default="structured-output", description="Local path to write partitioned output to"
+    )
+
+    @property
+    def output_path(self) -> Path:
+        return Path(self.output_dir).resolve()
+
+    def __post_init__(self):
+        if self.output_path.exists() and self.output_path.is_file():
+            raise ValueError("output path already exists as a file")
+
+
+@dataclass
+class LocalUploader(Uploader):
+    connector_type: str = CONNECTOR_TYPE
+    upload_config: LocalUploaderConfig = field(default_factory=LocalUploaderConfig)
+    connection_config: LocalConnectionConfig = field(
+        default_factory=lambda: LocalConnectionConfig()
+    )
+
+    def is_async(self) -> bool:
+        return False
+
+    def get_destination_path(self, file_data: FileData) -> Path:
+        if source_identifiers := file_data.source_identifiers:
+            rel_path = (
+                source_identifiers.relative_path[1:]
+                if source_identifiers.relative_path.startswith("/")
+                else source_identifiers.relative_path
+            )
+            new_path = self.upload_config.output_path / Path(rel_path)
+            final_path = str(new_path).replace(
+                source_identifiers.filename, f"{source_identifiers.filename}.json"
+            )
+        else:
+            final_path = self.upload_config.output_path / Path(f"{file_data.identifier}.json")
+        final_path = Path(final_path)
+        final_path.parent.mkdir(parents=True, exist_ok=True)
+        return final_path
+
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        final_path = self.get_destination_path(file_data=file_data)
+        with final_path.open("w") as f:
+            json.dump(data, f)
+
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        final_path = self.get_destination_path(file_data=file_data)
+        logger.debug(f"copying file from {path} to {final_path}")
+        shutil.copy(src=str(path), dst=str(final_path))
+
+
+local_source_entry = SourceRegistryEntry(
+    indexer=LocalIndexer,
+    indexer_config=LocalIndexerConfig,
+    downloader=LocalDownloader,
+    downloader_config=LocalDownloaderConfig,
+    connection_config=LocalConnectionConfig,
+)
+
+local_destination_entry = DestinationRegistryEntry(
+    uploader=LocalUploader, uploader_config=LocalUploaderConfig
+)

unstructured_ingest/v2/processes/connectors/milvus.py

@@ -0,0 +1,225 @@
+import json
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Generator, Optional, Union
+
+from dateutil import parser
+from pydantic import Field, Secret
+
+from unstructured_ingest.error import DestinationConnectionError, 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,
+    FileData,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+)
+
+if TYPE_CHECKING:
+    from pymilvus import MilvusClient
+
+CONNECTOR_TYPE = "milvus"
+
+
+class MilvusAccessConfig(AccessConfig):
+    password: Optional[str] = Field(default=None, description="Milvus password")
+    token: Optional[str] = Field(default=None, description="Milvus access token")
+
+
+class MilvusConnectionConfig(ConnectionConfig):
+    access_config: Secret[MilvusAccessConfig] = Field(
+        default=MilvusAccessConfig(), validate_default=True
+    )
+    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 = self.access_config.get_secret_value()
+        access_config_dict = access_config.model_dump()
+        connection_config_dict = self.model_dump()
+        connection_config_dict.pop("access_config", None)
+        connection_config_dict.update(access_config_dict)
+        # Drop any that were not set explicitly
+        connection_config_dict = {k: v for k, v in connection_config_dict.items() if v is not None}
+        return connection_config_dict
+
+    @requires_dependencies(["pymilvus"], extras="milvus")
+    @contextmanager
+    def get_client(self) -> Generator["MilvusClient", None, None]:
+        from pymilvus import MilvusClient
+
+        client = None
+        try:
+            client = MilvusClient(**self.get_connection_kwargs())
+            yield client
+        finally:
+            if client:
+                client.close()
+
+
+class MilvusUploadStagerConfig(UploadStagerConfig):
+    fields_to_include: Optional[list[str]] = None
+    """If set - list of fields to include in the output.
+    Unspecified fields are removed from the elements.
+    This action takes place after metadata flattening.
+    Missing fields will cause stager to throw KeyError."""
+
+    flatten_metadata: bool = True
+    """If set - flatten "metadata" key and put contents directly into data"""
+
+
+@dataclass
+class MilvusUploadStager(UploadStager):
+    upload_stager_config: MilvusUploadStagerConfig = field(
+        default_factory=lambda: MilvusUploadStagerConfig()
+    )
+
+    @staticmethod
+    def parse_date_string(date_string: str) -> float:
+        try:
+            timestamp = float(date_string)
+            return timestamp
+        except ValueError:
+            pass
+        return parser.parse(date_string).timestamp()
+
+    def conform_dict(self, element_dict: dict, file_data: FileData) -> dict:
+        working_data = element_dict.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 working_data:
+                working_data[default] = defaults[default]
+
+        if self.upload_stager_config.fields_to_include:
+            data_keys = set(working_data.keys())
+            for data_key in data_keys:
+                if data_key not in self.upload_stager_config.fields_to_include:
+                    working_data.pop(data_key)
+            for field_include_key in self.upload_stager_config.fields_to_include:
+                if field_include_key not in working_data:
+                    raise KeyError(f"Field '{field_include_key}' is missing in data!")
+
+        datetime_columns = [
+            "data_source_date_created",
+            "data_source_date_modified",
+            "data_source_date_processed",
+            "last_modified",
+        ]
+
+        json_dumps_fields = ["languages", "data_source_permissions_data"]
+
+        for datetime_column in datetime_columns:
+            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 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
+
+
+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
+class MilvusUploader(Uploader):
+    connection_config: MilvusConnectionConfig
+    upload_config: MilvusUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @DestinationConnectionError.wrap
+    def precheck(self):
+        from pymilvus import MilvusException
+
+        try:
+            with self.get_client() as client:
+                if not client.has_collection(self.upload_config.collection_name):
+                    raise DestinationConnectionError(
+                        f"Collection '{self.upload_config.collection_name}' does not exist"
+                    )
+        except MilvusException as milvus_exception:
+            raise DestinationConnectionError(
+                f"failed to precheck Milvus: {str(milvus_exception.message)}"
+            ) from milvus_exception
+
+    @contextmanager
+    def get_client(self) -> Generator["MilvusClient", None, None]:
+        with self.connection_config.get_client() as client:
+            if db_name := self.upload_config.db_name:
+                client.using_database(db_name=db_name)
+            yield client
+
+    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.info(
+            f"uploading {len(data)} entries to {self.connection_config.db_name} "
+            f"db in collection {self.upload_config.collection_name}"
+        )
+        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(
+                    f"failed to upload records to Milvus: {str(milvus_exception.message)}"
+                ) 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 run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        self.delete_by_record_id(file_data=file_data)
+        self.insert_results(data=data)
+
+
+milvus_destination_entry = DestinationRegistryEntry(
+    connection_config=MilvusConnectionConfig,
+    uploader=MilvusUploader,
+    uploader_config=MilvusUploaderConfig,
+    upload_stager=MilvusUploadStager,
+    upload_stager_config=MilvusUploadStagerConfig,
+)