unstructured-ingest 0.0.0__py3-none-any.whl

unstructured_ingest/connector/weaviate.py

@@ -0,0 +1,190 @@
+import copy
+import json
+import typing as t
+from dataclasses import dataclass, field
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.enhanced_dataclass.core import _asdict
+from unstructured_ingest.error import DestinationConnectionError, SourceConnectionError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    BaseConnectorConfig,
+    BaseDestinationConnector,
+    WriteConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if t.TYPE_CHECKING:
+    from weaviate import Client
+
+
+@dataclass
+class WeaviateAccessConfig(AccessConfig):
+    access_token: t.Optional[str] = enhanced_field(default=None, sensitive=True)
+    refresh_token: t.Optional[str] = enhanced_field(default=None, sensitive=True)
+    api_key: t.Optional[str] = enhanced_field(default=None, sensitive=True)
+    client_secret: t.Optional[str] = enhanced_field(default=None, sensitive=True)
+    scope: t.Optional[t.List[str]] = None
+    username: t.Optional[str] = None
+    password: t.Optional[str] = enhanced_field(default=None, sensitive=True)
+    anonymous: bool = False
+
+
+@dataclass
+class SimpleWeaviateConfig(BaseConnectorConfig):
+    access_config: WeaviateAccessConfig
+    host_url: str
+    class_name: str
+
+
+@dataclass
+class WeaviateWriteConfig(WriteConfig):
+    batch_size: int = 100
+
+
+@dataclass
+class WeaviateDestinationConnector(BaseDestinationConnector):
+    write_config: WeaviateWriteConfig
+    connector_config: SimpleWeaviateConfig
+    _client: t.Optional["Client"] = field(init=False, default=None)
+
+    def to_dict(self, **kwargs):
+        """
+        The _client variable in this dataclass breaks deepcopy due to:
+        TypeError: cannot pickle '_thread.lock' object
+        When serializing, remove it, meaning client data will need to be reinitialized
+        when deserialized
+        """
+        self_cp = copy.copy(self)
+        if hasattr(self_cp, "_client"):
+            setattr(self_cp, "_client", None)
+        return _asdict(self_cp, **kwargs)
+
+    @property
+    @requires_dependencies(["weaviate"], extras="weaviate")
+    def client(self) -> "Client":
+        if self._client is None:
+            from weaviate import Client
+
+            auth = self._resolve_auth_method()
+            self._client = Client(url=self.connector_config.host_url, auth_client_secret=auth)
+        return self._client
+
+    @requires_dependencies(["weaviate"], extras="weaviate")
+    @DestinationConnectionError.wrap
+    def initialize(self):
+        _ = self.client
+
+    @requires_dependencies(["weaviate"], extras="weaviate")
+    def check_connection(self):
+        try:
+            _ = self.client
+        except Exception as e:
+            logger.error(f"Failed to validate connection {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+    def _resolve_auth_method(self):
+        access_configs = self.connector_config.access_config
+        if access_configs.anonymous:
+            return None
+
+        if access_configs.access_token:
+            from weaviate.auth import AuthBearerToken
+
+            return AuthBearerToken(
+                access_token=access_configs.access_token,
+                refresh_token=access_configs.refresh_token,
+            )
+        elif access_configs.api_key:
+            from weaviate.auth import AuthApiKey
+
+            return AuthApiKey(api_key=access_configs.api_key)
+        elif access_configs.client_secret:
+            from weaviate.auth import AuthClientCredentials
+
+            return AuthClientCredentials(
+                client_secret=access_configs.client_secret, scope=access_configs.scope
+            )
+        elif access_configs.username and access_configs.password:
+            from weaviate.auth import AuthClientPassword
+
+            return AuthClientPassword(
+                username=access_configs.username,
+                password=access_configs.password,
+                scope=access_configs.scope,
+            )
+        return None
+
+    def conform_dict(self, data: dict) -> None:
+        """
+        Updates the element dictionary to conform to the Weaviate schema
+        """
+        from dateutil import parser
+
+        # Dict as string formatting
+        if record_locator := data.get("metadata", {}).get("data_source", {}).get("record_locator"):
+            # Explicit casting otherwise fails schema type checking
+            data["metadata"]["data_source"]["record_locator"] = str(json.dumps(record_locator))
+
+        # Array of items as string formatting
+        if points := data.get("metadata", {}).get("coordinates", {}).get("points"):
+            data["metadata"]["coordinates"]["points"] = str(json.dumps(points))
+
+        if links := data.get("metadata", {}).get("links", {}):
+            data["metadata"]["links"] = str(json.dumps(links))
+
+        if permissions_data := (
+            data.get("metadata", {}).get("data_source", {}).get("permissions_data")
+        ):
+            data["metadata"]["data_source"]["permissions_data"] = json.dumps(permissions_data)
+
+        # Datetime formatting
+        if date_created := data.get("metadata", {}).get("data_source", {}).get("date_created"):
+            data["metadata"]["data_source"]["date_created"] = parser.parse(date_created).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            )
+
+        if date_modified := data.get("metadata", {}).get("data_source", {}).get("date_modified"):
+            data["metadata"]["data_source"]["date_modified"] = parser.parse(date_modified).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            )
+
+        if date_processed := data.get("metadata", {}).get("data_source", {}).get("date_processed"):
+            data["metadata"]["data_source"]["date_processed"] = parser.parse(
+                date_processed
+            ).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            )
+
+        if last_modified := data.get("metadata", {}).get("last_modified", {}):
+            data["metadata"]["last_modified"] = parser.parse(last_modified).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            )
+
+        # String casting
+        if version := data.get("metadata", {}).get("data_source", {}).get("version"):
+            data["metadata"]["data_source"]["version"] = str(version)
+
+        if page_number := data.get("metadata", {}).get("page_number"):
+            data["metadata"]["page_number"] = str(page_number)
+
+        if regex_metadata := data.get("metadata", {}).get("regex_metadata"):
+            data["metadata"]["regex_metadata"] = str(json.dumps(regex_metadata))
+
+    def write_dict(self, *args, elements_dict: t.List[t.Dict[str, t.Any]], **kwargs) -> None:
+        logger.info(
+            f"writing {len(elements_dict)} objects to destination "
+            f"class {self.connector_config.class_name} "
+            f"at {self.connector_config.host_url}",
+        )
+
+        self.client.batch.configure(batch_size=self.write_config.batch_size)
+        with self.client.batch as b:
+            for e in elements_dict:
+                vector = e.pop("embeddings", None)
+                b.add_data_object(
+                    e,
+                    self.connector_config.class_name,
+                    vector=vector,
+                )

unstructured_ingest/connector/wikipedia.py

@@ -0,0 +1,208 @@
+import typing as t
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from unstructured_ingest.error import SourceConnectionError, SourceConnectionNetworkError
+from unstructured_ingest.interfaces import (
+    BaseConnectorConfig,
+    BaseSingleIngestDoc,
+    BaseSourceConnector,
+    IngestDocCleanupMixin,
+    SourceConnectorCleanupMixin,
+    SourceMetadata,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if t.TYPE_CHECKING:
+    from wikipedia import WikipediaPage
+
+
+@dataclass
+class SimpleWikipediaConfig(BaseConnectorConfig):
+    page_title: str
+    auto_suggest: bool = False
+
+
+@dataclass
+class WikipediaIngestDoc(IngestDocCleanupMixin, BaseSingleIngestDoc):
+    connector_config: SimpleWikipediaConfig = field(repr=False)
+
+    @property
+    @requires_dependencies(["wikipedia"], extras="wikipedia")
+    def page(self) -> "WikipediaPage":
+        import wikipedia
+
+        return wikipedia.page(
+            self.connector_config.page_title,
+            auto_suggest=self.connector_config.auto_suggest,
+        )
+
+    def get_filename_prefix(self) -> str:
+        title: str = str(self.connector_config.page_title)
+        title = " ".join(title.split()).replace(" ", "-")
+        return title
+
+    @property
+    def filename(self) -> Path:
+        raise NotImplementedError()
+
+    @property
+    def text(self) -> str:
+        raise NotImplementedError()
+
+    @property
+    def _output_filename(self):
+        raise NotImplementedError()
+
+    @property
+    def date_created(self) -> t.Optional[str]:
+        return None
+
+    @property
+    def date_modified(self) -> t.Optional[str]:
+        return None
+
+    @property
+    def record_locator(self) -> t.Optional[t.Dict[str, t.Any]]:
+        return {
+            "page_title": self.connector_config.page_title,
+            "page_url": self.source_metadata.source_url,  # type: ignore
+        }
+
+    def _create_full_tmp_dir_path(self):
+        self.filename.parent.mkdir(parents=True, exist_ok=True)
+
+    @requires_dependencies(["wikipedia"], extras="wikipedia")
+    def update_source_metadata(self):
+        from wikipedia.exceptions import PageError
+
+        try:
+            page = self.page
+        except PageError:
+            self.source_metadata = SourceMetadata(
+                exists=False,
+            )
+            return
+
+        self.source_metadata = SourceMetadata(
+            version=page.revision_id,
+            source_url=page.url,
+            exists=True,
+        )
+
+    @SourceConnectionError.wrap
+    @BaseSingleIngestDoc.skip_if_file_exists
+    def get_file(self):
+        """Fetches the "remote" doc and stores it locally on the filesystem."""
+        self._create_full_tmp_dir_path()
+        self.update_source_metadata()
+        with open(self.filename, "w", encoding="utf8") as f:
+            f.write(self.text)
+
+
+@dataclass
+class WikipediaIngestHTMLDoc(WikipediaIngestDoc):
+    registry_name: str = "wikipedia_html"
+
+    @property
+    def filename(self) -> Path:
+        return (
+            Path(self.read_config.download_dir) / f"{self.get_filename_prefix()}.html"
+        ).resolve()
+
+    @property
+    def text(self):
+        return self._get_html()
+
+    @SourceConnectionNetworkError.wrap
+    def _get_html(self):
+        return self.page.html()
+
+    @property
+    def _output_filename(self):
+        return Path(self.processor_config.output_dir) / f"{self.get_filename_prefix()}-html.json"
+
+
+@dataclass
+class WikipediaIngestTextDoc(WikipediaIngestDoc):
+    registry_name: str = "wikipedia_text"
+
+    @property
+    def filename(self) -> Path:
+        return (Path(self.read_config.download_dir) / f"{self.get_filename_prefix()}.txt").resolve()
+
+    @property
+    def text(self):
+        return self._get_content()
+
+    @SourceConnectionNetworkError.wrap
+    def _get_content(self):
+        return self.page.content
+
+    @property
+    def _output_filename(self):
+        return Path(self.processor_config.output_dir) / f"{self.get_filename_prefix()}-txt.json"
+
+
+@dataclass
+class WikipediaIngestSummaryDoc(WikipediaIngestDoc):
+    registry_name: str = "wikipedia_summary"
+
+    @property
+    def filename(self) -> Path:
+        return (
+            Path(self.read_config.download_dir) / f"{self.get_filename_prefix()}-summary.txt"
+        ).resolve()
+
+    @property
+    def text(self):
+        return self._get_summary()
+
+    @SourceConnectionNetworkError.wrap
+    def _get_summary(self):
+        return self.page.summary
+
+    @property
+    def _output_filename(self):
+        return Path(self.processor_config.output_dir) / f"{self.get_filename_prefix()}-summary.json"
+
+
+@dataclass
+class WikipediaSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
+    connector_config: SimpleWikipediaConfig
+
+    def initialize(self):
+        pass
+
+    @requires_dependencies(["wikipedia"], extras="wikipedia")
+    def check_connection(self):
+        import wikipedia
+
+        try:
+            wikipedia.page(
+                self.connector_config.page_title,
+                auto_suggest=self.connector_config.auto_suggest,
+            )
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+    def get_ingest_docs(self):
+        return [
+            WikipediaIngestTextDoc(
+                processor_config=self.processor_config,
+                connector_config=self.connector_config,
+                read_config=self.read_config,
+            ),
+            WikipediaIngestHTMLDoc(
+                processor_config=self.processor_config,
+                connector_config=self.connector_config,
+                read_config=self.read_config,
+            ),
+            WikipediaIngestSummaryDoc(
+                processor_config=self.processor_config,
+                connector_config=self.connector_config,
+                read_config=self.read_config,
+            ),
+        ]

unstructured_ingest/enhanced_dataclass/__init__.py

@@ -0,0 +1,4 @@
+from .dataclasses import enhanced_field
+from .json_mixin import EnhancedDataClassJsonMixin
+
+__all__ = ["enhanced_field", "EnhancedDataClassJsonMixin"]

unstructured_ingest/enhanced_dataclass/core.py

@@ -0,0 +1,99 @@
+import _thread
+import copy
+import functools
+from dataclasses import fields
+
+from dataclasses_json.core import (
+    Collection,
+    Enum,
+    Mapping,
+    _encode_overrides,
+    _handle_undefined_parameters_safe,
+    _user_overrides_or_exts,
+    is_dataclass,
+)
+
+
+def _recursive_repr(user_function):
+    # Copied from dataclasses as this method isn't exposed for importing
+    repr_running = set()
+
+    @functools.wraps(user_function)
+    def wrapper(self):
+        key = id(self), _thread.get_ident()
+        if key in repr_running:
+            return "..."
+        repr_running.add(key)
+        try:
+            result = user_function(self)
+        finally:
+            repr_running.discard(key)
+        return result
+
+    return wrapper
+
+
+def _asdict(
+    obj,
+    encode_json=False,
+    redact_sensitive=False,
+    redacted_text="***REDACTED***",
+    apply_name_overload: bool = True,
+):
+    """
+    A re-implementation of `asdict` (based on the original in the `dataclasses`
+    source) to support arbitrary Collection and Mapping types.
+    """
+    if is_dataclass(obj):
+        result = []
+        overrides = _user_overrides_or_exts(obj)
+        for field in fields(obj):
+            if overrides[field.name].encoder:
+                value = getattr(obj, field.name)
+            else:
+                value = _asdict(
+                    getattr(obj, field.name),
+                    encode_json=encode_json,
+                    redact_sensitive=redact_sensitive,
+                    redacted_text=redacted_text,
+                    apply_name_overload=apply_name_overload,
+                )
+            if getattr(field, "sensitive", False) and redact_sensitive and value:
+                value = redacted_text
+            if getattr(field, "overload_name", None) and apply_name_overload:
+                overload_name = getattr(field, "overload_name")
+                result.append((overload_name, value))
+            else:
+                result.append((field.name, value))
+
+        result = _handle_undefined_parameters_safe(cls=obj, kvs=dict(result), usage="to")
+        return _encode_overrides(
+            dict(result), _user_overrides_or_exts(obj), encode_json=encode_json
+        )
+    elif isinstance(obj, Mapping):
+        return {
+            _asdict(
+                k,
+                encode_json=encode_json,
+                redact_sensitive=redact_sensitive,
+                redacted_text=redacted_text,
+            ): _asdict(
+                v,
+                encode_json=encode_json,
+                redact_sensitive=redact_sensitive,
+                redacted_text=redacted_text,
+            )
+            for k, v in obj.items()
+        }
+    elif isinstance(obj, Collection) and not isinstance(obj, (str, bytes, Enum)):
+        return [
+            _asdict(
+                v,
+                encode_json=encode_json,
+                redact_sensitive=redact_sensitive,
+                redacted_text=redacted_text,
+            )
+            for v in obj
+        ]
+    else:
+        return copy.deepcopy(obj)

unstructured_ingest/enhanced_dataclass/dataclasses.py

@@ -0,0 +1,54 @@
+import typing as t
+from dataclasses import MISSING, Field
+
+from unstructured_ingest.enhanced_dataclass.core import _recursive_repr
+
+
+class EnhancedField(Field):
+    def __init__(self, *args, sensitive=False, overload_name: t.Optional[str] = None):
+        super().__init__(*args)
+        self.sensitive = sensitive
+        self.overload_name = overload_name
+
+    @_recursive_repr
+    def __repr__(self):
+        # Support for kw_only added in 3.10, to support as low as 3.8, need to dynamically map
+        fields_array = [
+            f"name={self.name!r}",
+            f"type={self.type!r}",
+            f"default={self.default!r}",
+            f"default_factory={self.default_factory!r}",
+            f"init={self.init!r}",
+            f"repr={self.repr!r}",
+            f"hash={self.hash!r}",
+            f"compare={self.compare!r}",
+            f"metadata={self.metadata!r}",
+            f"sensitive={self.sensitive!r}",
+            f"overload_name={self.overload_name!r}",
+            f"_field_type={self._field_type}",
+        ]
+        if kw_only := getattr(self, "kw_only", None):
+            fields_array.append(f"kw_only={kw_only!r}")
+        return "Field({})".format(",".join(fields_array))
+
+
+def enhanced_field(
+    *,
+    default=MISSING,
+    default_factory=MISSING,
+    init: bool = True,
+    repr: bool = True,
+    hash=None,
+    compare: bool = True,
+    metadata=None,
+    kw_only=MISSING,
+    sensitive: bool = False,
+    overload_name: t.Optional[str] = None,
+):
+    if default is not MISSING and default_factory is not MISSING:
+        raise ValueError("cannot specify both default and default_factory")
+    args = [default, default_factory, init, repr, hash, compare, metadata]
+    # Support for kw_only added in 3.10, to support as low as 3.8, need to dynamically map
+    if "kw_only" in EnhancedField.__slots__:
+        args.append(kw_only)
+    return EnhancedField(*args, sensitive=sensitive, overload_name=overload_name)

unstructured_ingest/enhanced_dataclass/json_mixin.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+import json
+from dataclasses import InitVar, fields
+from typing import Any, Callable, Optional, Type, TypeVar, Union
+
+import dataclasses_json.core as dataclasses_json_core
+from dataclasses_json import DataClassJsonMixin
+
+from unstructured_ingest.enhanced_dataclass.core import _asdict
+
+A = TypeVar("A", bound="EnhancedDataClassJsonMixin")
+
+# Monkey-patch _decode_dataclass class to support name override
+og_decode_dataclass = dataclasses_json_core._decode_dataclass
+
+
+def custom_decode_dataclass(cls, kvs, infer_missing):
+    dataclass_fields = fields(cls)
+    for f in [
+        field
+        for field in dataclass_fields
+        if hasattr(field, "overload_name") and getattr(field, "overload_name", None)
+    ]:
+        field_name = f.name
+        overload_name = getattr(f, "overload_name")
+        if isinstance(kvs, dict) and overload_name in kvs:
+            kvs[field_name] = kvs.pop(overload_name)
+    return og_decode_dataclass(cls, kvs, infer_missing)
+
+
+dataclasses_json_core._decode_dataclass = custom_decode_dataclass
+
+
+class EnhancedDataClassJsonMixin(DataClassJsonMixin):
+    """A mixin class extending DataClassJsonMixin.
+
+    This class extends the functionality of DataClassJsonMixin to provide enhanced functionality
+    for JSON serialization and deserialization. It introduces options for redacting sensitive
+    information, custom encoding, and more advanced schema handling.
+
+    Attributes:
+        N/A (No additional attributes)
+
+    Methods:
+        to_json: Serialize the object to JSON format with customizable options.
+        from_dict: Deserialize a dictionary into an object of this class.
+        to_dict: Convert the object to a dictionary with customizable options.
+        schema: Generate a schema for validating and parsing JSON data based on this class.
+    """
+
+    @classmethod
+    def check_init_var(cls):
+        ann = cls.__dict__.get("__annotations__", {})
+        init_vars = {k: v for k, v in ann.items() if isinstance(v, InitVar)}
+        if init_vars:
+            raise TypeError(
+                "Class {} has the following fields defined with an InitVar which "
+                "cannot be used with EnhancedDataClassJsonMixin: {}".format(
+                    cls.__name__, ", ".join(init_vars.keys())
+                )
+            )
+
+    def to_json(
+        self,
+        *,
+        skipkeys: bool = False,
+        ensure_ascii: bool = True,
+        check_circular: bool = True,
+        allow_nan: bool = True,
+        indent: Optional[Union[int, str]] = None,
+        separators: Optional[tuple[str, str]] = None,
+        default: Optional[Callable[..., Any]] = None,
+        sort_keys: bool = False,
+        redact_sensitive: bool = False,
+        redacted_text: str = "***REDACTED***",
+        apply_name_overload: bool = True,
+        **kw: Any,
+    ) -> str:
+        self.check_init_var()
+        return json.dumps(
+            self.to_dict(
+                encode_json=False,
+                redact_sensitive=redact_sensitive,
+                redacted_text=redacted_text,
+                apply_name_overload=apply_name_overload,
+            ),
+            cls=dataclasses_json_core._ExtendedEncoder,
+            skipkeys=skipkeys,
+            ensure_ascii=ensure_ascii,
+            check_circular=check_circular,
+            allow_nan=allow_nan,
+            indent=indent,
+            separators=separators,
+            default=default,
+            sort_keys=sort_keys,
+            **kw,
+        )
+
+    @classmethod
+    def from_dict(
+        cls: Type[A],
+        kvs: dataclasses_json_core.Json,
+        *,
+        infer_missing=False,
+        apply_name_overload=False,
+    ) -> A:
+        cls.check_init_var()
+        return dataclasses_json_core._decode_dataclass(cls, kvs, infer_missing)
+
+    def to_dict(
+        self,
+        encode_json: bool = False,
+        redact_sensitive: bool = False,
+        redacted_text: str = "***REDACTED***",
+        apply_name_overload: bool = True,
+    ) -> dict[str, dataclasses_json_core.Json]:
+        self.check_init_var()
+        return _asdict(
+            self,
+            encode_json=encode_json,
+            redact_sensitive=redact_sensitive,
+            redacted_text=redacted_text,
+            apply_name_overload=apply_name_overload,
+        )