unstructured-ingest 0.3.13__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/weaviate/weaviate.py

@@ -0,0 +1,267 @@
+import json
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from datetime import date, datetime
+from typing import TYPE_CHECKING, Any, Generator, Optional
+
+from dateutil import parser
+from pydantic import Field, Secret
+
+from unstructured_ingest.error import DestinationConnectionError, WriteError
+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
+
+if TYPE_CHECKING:
+    from weaviate.classes.init import Timeout
+    from weaviate.client import WeaviateClient
+    from weaviate.collections.batch.client import BatchClient
+
+CONNECTOR_TYPE = "weaviate"
+
+
+class WeaviateAccessConfig(AccessConfig, ABC):
+    pass
+
+
+class WeaviateConnectionConfig(ConnectionConfig, ABC):
+    init_timeout: int = Field(default=2, ge=0, description="Timeout for initialization checks")
+    insert_timeout: int = Field(default=90, ge=0, description="Timeout for insert operations")
+    query_timeout: int = Field(default=30, ge=0, description="Timeout for query operations")
+    access_config: Secret[WeaviateAccessConfig] = Field(
+        default=WeaviateAccessConfig(), validate_default=True
+    )
+
+    @requires_dependencies(["weaviate"], extras="weaviate")
+    def get_timeout(self) -> "Timeout":
+        from weaviate.classes.init import Timeout
+
+        return Timeout(init=self.init_timeout, query=self.query_timeout, insert=self.insert_timeout)
+
+    @abstractmethod
+    @contextmanager
+    def get_client(self) -> Generator["WeaviateClient", None, None]:
+        pass
+
+
+class WeaviateUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class WeaviateUploadStager(UploadStager):
+    upload_stager_config: WeaviateUploadStagerConfig = field(
+        default_factory=lambda: WeaviateUploadStagerConfig()
+    )
+
+    @staticmethod
+    def parse_date_string(date_string: str) -> date:
+        try:
+            timestamp = float(date_string)
+            return datetime.fromtimestamp(timestamp)
+        except Exception as e:
+            logger.debug(f"date {date_string} string not a timestamp: {e}")
+        return parser.parse(date_string)
+
+    def conform_dict(self, element_dict: dict, file_data: FileData) -> dict:
+        """
+        Updates the element dictionary to conform to the Weaviate schema
+        """
+        data = element_dict.copy()
+        working_data = data.copy()
+        # Dict as string formatting
+        if (
+            record_locator := working_data.get("metadata", {})
+            .get("data_source", {})
+            .get("record_locator")
+        ):
+            # Explicit casting otherwise fails schema type checking
+            working_data["metadata"]["data_source"]["record_locator"] = str(
+                json.dumps(record_locator)
+            )
+
+        # Array of items as string formatting
+        if points := working_data.get("metadata", {}).get("coordinates", {}).get("points"):
+            working_data["metadata"]["coordinates"]["points"] = str(json.dumps(points))
+
+        if links := working_data.get("metadata", {}).get("links", {}):
+            working_data["metadata"]["links"] = str(json.dumps(links))
+
+        if permissions_data := (
+            working_data.get("metadata", {}).get("data_source", {}).get("permissions_data")
+        ):
+            working_data["metadata"]["data_source"]["permissions_data"] = json.dumps(
+                permissions_data
+            )
+
+        # Datetime formatting
+        if (
+            date_created := working_data.get("metadata", {})
+            .get("data_source", {})
+            .get("date_created")
+        ):
+            working_data["metadata"]["data_source"]["date_created"] = self.parse_date_string(
+                date_created
+            ).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            )
+
+        if (
+            date_modified := working_data.get("metadata", {})
+            .get("data_source", {})
+            .get("date_modified")
+        ):
+            working_data["metadata"]["data_source"]["date_modified"] = self.parse_date_string(
+                date_modified
+            ).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            )
+
+        if (
+            date_processed := working_data.get("metadata", {})
+            .get("data_source", {})
+            .get("date_processed")
+        ):
+            working_data["metadata"]["data_source"]["date_processed"] = self.parse_date_string(
+                date_processed
+            ).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            )
+
+        if last_modified := working_data.get("metadata", {}).get("last_modified"):
+            working_data["metadata"]["last_modified"] = self.parse_date_string(
+                last_modified
+            ).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            )
+
+        # String casting
+        if version := working_data.get("metadata", {}).get("data_source", {}).get("version"):
+            working_data["metadata"]["data_source"]["version"] = str(version)
+
+        if page_number := working_data.get("metadata", {}).get("page_number"):
+            working_data["metadata"]["page_number"] = str(page_number)
+
+        if regex_metadata := working_data.get("metadata", {}).get("regex_metadata"):
+            working_data["metadata"]["regex_metadata"] = str(json.dumps(regex_metadata))
+
+        working_data[RECORD_ID_LABEL] = file_data.identifier
+        return working_data
+
+
+class WeaviateUploaderConfig(UploaderConfig):
+    collection: str = Field(description="The name of the collection this object belongs to")
+    batch_size: Optional[int] = Field(default=None, description="Number of records per batch")
+    requests_per_minute: Optional[int] = Field(default=None, description="Rate limit for upload")
+    dynamic_batch: bool = Field(default=True, description="Whether to use dynamic batch")
+    record_id_key: str = Field(
+        default=RECORD_ID_LABEL,
+        description="searchable key to find entries for the same record on previous runs",
+    )
+
+    def model_post_init(self, __context: Any) -> None:
+        batch_types = {
+            "fixed_size": self.batch_size is not None,
+            "rate_limited": self.requests_per_minute is not None,
+            "dynamic": self.dynamic_batch,
+        }
+
+        enabled_batch_modes = [batch_key for batch_key, flag in batch_types.items() if flag]
+        if not enabled_batch_modes:
+            raise ValueError("No batch mode enabled")
+        if len(enabled_batch_modes) > 1:
+            raise ValueError(
+                "Multiple batch modes enabled, only one mode can be used: {}".format(
+                    ", ".join(enabled_batch_modes)
+                )
+            )
+        logger.info(f"Uploader config instantiated with {enabled_batch_modes[0]} batch mode")
+
+    @contextmanager
+    def get_batch_client(self, client: "WeaviateClient") -> Generator["BatchClient", None, None]:
+        if self.dynamic_batch:
+            with client.batch.dynamic() as batch_client:
+                yield batch_client
+        elif self.batch_size:
+            with client.batch.fixed_size(batch_size=self.batch_size) as batch_client:
+                yield batch_client
+        elif self.requests_per_minute:
+            with client.batch.rate_limit(
+                requests_per_minute=self.requests_per_minute
+            ) as batch_client:
+                yield batch_client
+        else:
+            raise ValueError("No batch mode enabled")
+
+
+@dataclass
+class WeaviateUploader(Uploader, ABC):
+    upload_config: WeaviateUploaderConfig
+    connection_config: WeaviateConnectionConfig
+
+    def precheck(self) -> None:
+        try:
+            self.connection_config.get_client()
+        except Exception as e:
+            logger.error(f"Failed to validate connection {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def check_for_errors(self, client: "WeaviateClient") -> None:
+        failed_uploads = client.batch.failed_objects
+        if failed_uploads:
+            for failure in failed_uploads:
+                logger.error(
+                    f"Failed to upload object with id {failure.original_uuid}: {failure.message}"
+                )
+            raise WriteError("Failed to upload to weaviate")
+
+    @requires_dependencies(["weaviate"], extras="weaviate")
+    def delete_by_record_id(self, client: "WeaviateClient", file_data: FileData) -> None:
+        from weaviate.classes.query import Filter
+
+        record_id = file_data.identifier
+        collection = client.collections.get(self.upload_config.collection)
+        delete_filter = Filter.by_property(name=self.upload_config.record_id_key).equal(
+            val=record_id
+        )
+        # There is a configurable maximum limit (QUERY_MAXIMUM_RESULTS) on the number of
+        # objects that can be deleted in a single query (default 10,000). To delete
+        # more objects than the limit, re-run the query until nothing is deleted.
+        while True:
+            resp = collection.data.delete_many(where=delete_filter)
+            if resp.failed:
+                raise WriteError(
+                    f"failed to delete records in collection "
+                    f"{self.upload_config.collection} with record "
+                    f"id property {record_id}"
+                )
+            if not resp.failed and not resp.successful:
+                break
+
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        logger.info(
+            f"writing {len(data)} objects to destination "
+            f"class {self.connection_config.access_config} "
+        )
+
+        with self.connection_config.get_client() as weaviate_client:
+            self.delete_by_record_id(client=weaviate_client, file_data=file_data)
+            with self.upload_config.get_batch_client(client=weaviate_client) as batch_client:
+                for e in data:
+                    vector = e.pop("embeddings", None)
+                    batch_client.add_object(
+                        collection=self.upload_config.collection,
+                        properties=e,
+                        vector=vector,
+                    )
+            self.check_for_errors(client=weaviate_client)

unstructured_ingest/v2/processes/embedder.py

@@ -0,0 +1,195 @@
+import json
+from abc import ABC
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Literal, Optional
+
+from pydantic import BaseModel, Field, SecretStr
+
+from unstructured_ingest.v2.interfaces.process import BaseProcess
+
+if TYPE_CHECKING:
+    from unstructured_ingest.embed.interfaces import BaseEmbeddingEncoder
+
+
+class EmbedderConfig(BaseModel):
+    embedding_provider: Optional[
+        Literal[
+            "openai",
+            "azure-openai",
+            "huggingface",
+            "aws-bedrock",
+            "vertexai",
+            "voyageai",
+            "octoai",
+            "mixedbread-ai",
+            "togetherai",
+        ]
+    ] = Field(default=None, description="Type of the embedding class to be used.")
+    embedding_api_key: Optional[SecretStr] = Field(
+        default=None,
+        description="API key for the embedding model, for the case an API key is needed.",
+    )
+    embedding_model_name: Optional[str] = Field(
+        default=None,
+        description="Embedding model name, if needed. "
+        "Chooses a particular LLM between different options, to embed with it.",
+    )
+    embedding_aws_access_key_id: Optional[str] = Field(
+        default=None, description="AWS access key used for AWS-based embedders, such as bedrock"
+    )
+    embedding_aws_secret_access_key: Optional[SecretStr] = Field(
+        default=None, description="AWS secret key used for AWS-based embedders, such as bedrock"
+    )
+    embedding_aws_region: Optional[str] = Field(
+        default="us-west-2", description="AWS region used for AWS-based embedders, such as bedrock"
+    )
+    embedding_azure_endpoint: Optional[str] = Field(
+        default=None,
+        description="Your Azure endpoint, including the resource, "
+        "e.g. `https://example-resource.azure.openai.com/`",
+    )
+    embedding_azure_api_version: Optional[str] = Field(
+        description="Azure API version", default=None
+    )
+
+    def get_huggingface_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
+        from unstructured_ingest.embed.huggingface import (
+            HuggingFaceEmbeddingConfig,
+            HuggingFaceEmbeddingEncoder,
+        )
+
+        return HuggingFaceEmbeddingEncoder(
+            config=HuggingFaceEmbeddingConfig.model_validate(embedding_kwargs)
+        )
+
+    def get_openai_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
+        from unstructured_ingest.embed.openai import OpenAIEmbeddingConfig, OpenAIEmbeddingEncoder
+
+        return OpenAIEmbeddingEncoder(config=OpenAIEmbeddingConfig.model_validate(embedding_kwargs))
+
+    def get_azure_openai_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
+        from unstructured_ingest.embed.azure_openai import (
+            AzureOpenAIEmbeddingConfig,
+            AzureOpenAIEmbeddingEncoder,
+        )
+
+        config_kwargs = {
+            "api_key": self.embedding_api_key,
+            "azure_endpoint": self.embedding_azure_endpoint,
+        }
+        if api_version := self.embedding_azure_api_version:
+            config_kwargs["api_version"] = api_version
+        if model_name := self.embedding_model_name:
+            config_kwargs["model_name"] = model_name
+
+        return AzureOpenAIEmbeddingEncoder(
+            config=AzureOpenAIEmbeddingConfig.model_validate(config_kwargs)
+        )
+
+    def get_octoai_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
+        from unstructured_ingest.embed.octoai import OctoAiEmbeddingConfig, OctoAIEmbeddingEncoder
+
+        return OctoAIEmbeddingEncoder(config=OctoAiEmbeddingConfig.model_validate(embedding_kwargs))
+
+    def get_bedrock_embedder(self) -> "BaseEmbeddingEncoder":
+        from unstructured_ingest.embed.bedrock import (
+            BedrockEmbeddingConfig,
+            BedrockEmbeddingEncoder,
+        )
+
+        return BedrockEmbeddingEncoder(
+            config=BedrockEmbeddingConfig(
+                aws_access_key_id=self.embedding_aws_access_key_id,
+                aws_secret_access_key=self.embedding_aws_secret_access_key.get_secret_value(),
+                region_name=self.embedding_aws_region,
+            )
+        )
+
+    def get_vertexai_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
+        from unstructured_ingest.embed.vertexai import (
+            VertexAIEmbeddingConfig,
+            VertexAIEmbeddingEncoder,
+        )
+
+        return VertexAIEmbeddingEncoder(
+            config=VertexAIEmbeddingConfig.model_validate(embedding_kwargs)
+        )
+
+    def get_voyageai_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
+        from unstructured_ingest.embed.voyageai import (
+            VoyageAIEmbeddingConfig,
+            VoyageAIEmbeddingEncoder,
+        )
+
+        return VoyageAIEmbeddingEncoder(
+            config=VoyageAIEmbeddingConfig.model_validate(embedding_kwargs)
+        )
+
+    def get_mixedbread_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
+        from unstructured_ingest.embed.mixedbreadai import (
+            MixedbreadAIEmbeddingConfig,
+            MixedbreadAIEmbeddingEncoder,
+        )
+
+        return MixedbreadAIEmbeddingEncoder(
+            config=MixedbreadAIEmbeddingConfig.model_validate(embedding_kwargs)
+        )
+
+    def get_togetherai_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
+        from unstructured_ingest.embed.togetherai import (
+            TogetherAIEmbeddingConfig,
+            TogetherAIEmbeddingEncoder,
+        )
+
+        return TogetherAIEmbeddingEncoder(
+            config=TogetherAIEmbeddingConfig.model_validate(embedding_kwargs)
+        )
+
+    def get_embedder(self) -> "BaseEmbeddingEncoder":
+        kwargs: dict[str, Any] = {}
+        if self.embedding_api_key:
+            kwargs["api_key"] = self.embedding_api_key.get_secret_value()
+        if self.embedding_model_name:
+            kwargs["model_name"] = self.embedding_model_name
+        # TODO make this more dynamic to map to encoder configs
+        if self.embedding_provider == "openai":
+            return self.get_openai_embedder(embedding_kwargs=kwargs)
+
+        if self.embedding_provider == "huggingface":
+            return self.get_huggingface_embedder(embedding_kwargs=kwargs)
+
+        if self.embedding_provider == "octoai":
+            return self.get_octoai_embedder(embedding_kwargs=kwargs)
+
+        if self.embedding_provider == "aws-bedrock":
+            return self.get_bedrock_embedder()
+
+        if self.embedding_provider == "vertexai":
+            return self.get_vertexai_embedder(embedding_kwargs=kwargs)
+
+        if self.embedding_provider == "voyageai":
+            return self.get_voyageai_embedder(embedding_kwargs=kwargs)
+        if self.embedding_provider == "mixedbread-ai":
+            return self.get_mixedbread_embedder(embedding_kwargs=kwargs)
+        if self.embedding_provider == "togetherai":
+            return self.get_togetherai_embedder(embedding_kwargs=kwargs)
+        if self.embedding_provider == "azure-openai":
+            return self.get_azure_openai_embedder(embedding_kwargs=kwargs)
+
+        raise ValueError(f"{self.embedding_provider} not a recognized encoder")
+
+
+@dataclass
+class Embedder(BaseProcess, ABC):
+    config: EmbedderConfig
+
+    def run(self, elements_filepath: Path, **kwargs: Any) -> list[dict]:
+        # TODO update base embedder classes to support async
+        embedder = self.config.get_embedder()
+        with elements_filepath.open("r") as elements_file:
+            elements = json.load(elements_file)
+        if not elements:
+            return [e.to_dict() for e in elements]
+        embedded_elements = embedder.embed_documents(elements=elements)
+        return embedded_elements

unstructured_ingest/v2/processes/filter.py

@@ -0,0 +1,60 @@
+import fnmatch
+from abc import ABC
+from dataclasses import dataclass, field
+from typing import Any, Callable, Optional
+
+from pydantic import BaseModel, Field
+
+from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.interfaces.process import BaseProcess
+from unstructured_ingest.v2.logger import logger
+
+
+class FiltererConfig(BaseModel):
+    file_glob: Optional[list[str]] = Field(
+        default=None,
+        description="file globs to limit which types of " "files are accepted",
+        examples=["*.pdf", "*.html"],
+    )
+    max_file_size: Optional[int] = Field(
+        default=None, description="Max file size to process in bytes"
+    )
+
+
+@dataclass
+class Filterer(BaseProcess, ABC):
+    config: FiltererConfig = field(default_factory=lambda: FiltererConfig())
+    filters: list[Callable[[FileData], bool]] = field(init=False, default_factory=list)
+
+    def __post_init__(self):
+        # Populate the filters based on values in config
+        if self.config.file_glob is not None:
+            self.filters.append(self.glob_filter)
+        if self.config.max_file_size:
+            self.filters.append(self.file_size_filter)
+
+    def is_async(self) -> bool:
+        return False
+
+    def file_size_filter(self, file_data: FileData) -> bool:
+        if filesize_bytes := file_data.metadata.filesize_bytes:
+            return filesize_bytes <= self.config.max_file_size
+        return True
+
+    def glob_filter(self, file_data: FileData) -> bool:
+        patterns = self.config.file_glob
+        path = file_data.source_identifiers.fullpath
+        for pattern in patterns:
+            if fnmatch.filter([path], pattern):
+                return True
+        logger.debug(f"the file {path!r} is discarded as it does not match any given glob.")
+        return False
+
+    def run(self, file_data: FileData, **kwargs: Any) -> Optional[FileData]:
+        for filter in self.filters:
+            if not filter(file_data):
+                logger.debug(
+                    f"filtered out file data due to {filter.__name__}: {file_data.identifier}"
+                )
+                return None
+        return file_data