unstructured-ingest 1.2.32__py3-none-any.whl

unstructured_ingest/processes/connectors/lancedb/gcp.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.interfaces.connector import AccessConfig
+from unstructured_ingest.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.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/processes/connectors/lancedb/lancedb.py

@@ -0,0 +1,181 @@
+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
+
+from pydantic import Field
+
+from unstructured_ingest.data_types.file_data import FileData
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.interfaces import (
+    ConnectionConfig,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.constants import RECORD_ID_LABEL
+from unstructured_ingest.utils.data_prep import flatten_dict
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+CONNECTOR_TYPE = "lancedb"
+
+if TYPE_CHECKING:
+    from lancedb import AsyncConnection
+    from lancedb.table import AsyncTable
+    from pandas import DataFrame
+
+
+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
+    )
+
+    @requires_dependencies(["pandas"], extras="lancedb")
+    def run(
+        self,
+        elements_filepath: Path,
+        file_data: FileData,
+        output_dir: Path,
+        output_filename: str,
+        **kwargs: Any,
+    ) -> Path:
+        import pandas as pd
+
+        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()
+
+    @requires_dependencies(["pandas"], extras="lancedb")
+    async def run_async(self, path, file_data, **kwargs):
+        import pandas as pd
+
+        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: "DataFrame", schema) -> "DataFrame":
+        import pandas as pd
+
+        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/processes/connectors/lancedb/local.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.interfaces.connector import AccessConfig
+from unstructured_ingest.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.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/processes/connectors/local.py

@@ -0,0 +1,227 @@
+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.data_types.file_data import (
+    FileData,
+    FileDataSourceMetadata,
+    SourceIdentifiers,
+)
+from unstructured_ingest.error import FileExistsError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    Downloader,
+    DownloaderConfig,
+    DownloadResponse,
+    Indexer,
+    IndexerConfig,
+    Uploader,
+    UploaderConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.processes.utils.blob_storage import (
+    BlobStoreUploadStager,
+    BlobStoreUploadStagerConfig,
+)
+
+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():
+            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
+                ),
+            )
+            file_data = FileData(
+                identifier=str(file_path.resolve()),
+                connector_type=CONNECTOR_TYPE,
+                source_identifiers=source_identifiers,
+                metadata=self.get_file_metadata(path=file_path),
+                display_name=source_identifiers.fullpath,
+            )
+            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 FileExistsError(f"output path {self.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,
+    upload_stager_config=BlobStoreUploadStagerConfig,
+    upload_stager=BlobStoreUploadStager,
+)