unstructured-ingest 0.3.13__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/embed/azure_openai.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from pydantic import Field
+
+from unstructured_ingest.embed.openai import OpenAIEmbeddingConfig, OpenAIEmbeddingEncoder
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if TYPE_CHECKING:
+    from openai import AzureOpenAI
+
+
+class AzureOpenAIEmbeddingConfig(OpenAIEmbeddingConfig):
+    api_version: str = Field(description="Azure API version", default="2024-06-01")
+    azure_endpoint: str
+    embedder_model_name: str = Field(default="text-embedding-ada-002", alias="model_name")
+
+    @requires_dependencies(["openai"], extras="openai")
+    def get_client(self) -> "AzureOpenAI":
+        from openai import AzureOpenAI
+
+        return AzureOpenAI(
+            api_key=self.api_key.get_secret_value(),
+            api_version=self.api_version,
+            azure_endpoint=self.azure_endpoint,
+        )
+
+
+@dataclass
+class AzureOpenAIEmbeddingEncoder(OpenAIEmbeddingEncoder):
+    config: AzureOpenAIEmbeddingConfig

unstructured_ingest/embed/bedrock.py

@@ -0,0 +1,193 @@
+import asyncio
+import json
+import os
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, AsyncIterable
+
+from pydantic import Field, SecretStr
+
+from unstructured_ingest.embed.interfaces import (
+    AsyncBaseEmbeddingEncoder,
+    BaseEmbeddingEncoder,
+    EmbeddingConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.errors import ProviderError, RateLimitError, UserAuthError, UserError
+
+if TYPE_CHECKING:
+    from botocore.client import BaseClient
+
+    class BedrockClient(BaseClient):
+        def invoke_model(self, body: str, modelId: str, accept: str, contentType: str) -> dict:
+            pass
+
+    class AsyncBedrockClient(BaseClient):
+        async def invoke_model(
+            self, body: str, modelId: str, accept: str, contentType: str
+        ) -> dict:
+            pass
+
+
+def conform_query(query: str, provider: str) -> dict:
+    # replace newlines, which can negatively affect performance.
+    text = query.replace(os.linesep, " ")
+
+    # format input body for provider
+    input_body = {}
+    if provider == "cohere":
+        if "input_type" not in input_body:
+            input_body["input_type"] = "search_document"
+        input_body["texts"] = [text]
+    else:
+        # includes common provider == "amazon"
+        input_body["inputText"] = text
+    return input_body
+
+
+class BedrockEmbeddingConfig(EmbeddingConfig):
+    aws_access_key_id: SecretStr
+    aws_secret_access_key: SecretStr
+    region_name: str = "us-west-2"
+    embed_model_name: str = Field(default="amazon.titan-embed-text-v1", alias="model_name")
+
+    def wrap_error(self, e: Exception) -> Exception:
+        from botocore.exceptions import ClientError
+
+        if isinstance(e, ClientError):
+            # https://docs.aws.amazon.com/awssupport/latest/APIReference/CommonErrors.html
+            http_response = e.response
+            meta = http_response["ResponseMetadata"]
+            http_response_code = meta["HTTPStatusCode"]
+            error_code = http_response["Error"]["Code"]
+            if http_response_code == 400:
+                if error_code == "ValidationError":
+                    return UserError(http_response["Error"])
+                elif error_code == "ThrottlingException":
+                    return RateLimitError(http_response["Error"])
+                elif error_code == "NotAuthorized" or error_code == "AccessDeniedException":
+                    return UserAuthError(http_response["Error"])
+            if http_response_code == 403:
+                return UserAuthError(http_response["Error"])
+            if 400 <= http_response_code < 500:
+                return UserError(http_response["Error"])
+            if http_response_code >= 500:
+                return ProviderError(http_response["Error"])
+
+        logger.error(f"unhandled exception from bedrock: {e}", exc_info=True)
+        return e
+
+    @requires_dependencies(
+        ["boto3", "numpy", "botocore"],
+        extras="bedrock",
+    )
+    def get_client(self) -> "BedrockClient":
+        import boto3
+
+        bedrock_client = boto3.client(
+            service_name="bedrock-runtime",
+            aws_access_key_id=self.aws_access_key_id.get_secret_value(),
+            aws_secret_access_key=self.aws_secret_access_key.get_secret_value(),
+            region_name=self.region_name,
+        )
+
+        return bedrock_client
+
+    @requires_dependencies(
+        ["aioboto3"],
+        extras="bedrock",
+    )
+    @asynccontextmanager
+    async def get_async_client(self) -> AsyncIterable["AsyncBedrockClient"]:
+        import aioboto3
+
+        session = aioboto3.Session()
+        async with session.client(
+            "bedrock-runtime",
+            aws_access_key_id=self.aws_access_key_id.get_secret_value(),
+            aws_secret_access_key=self.aws_secret_access_key.get_secret_value(),
+            region_name=self.region_name,
+        ) as aws_bedrock:
+            yield aws_bedrock
+
+
+@dataclass
+class BedrockEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: BedrockEmbeddingConfig
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    def embed_query(self, query: str) -> list[float]:
+        """Call out to Bedrock embedding endpoint."""
+        provider = self.config.embed_model_name.split(".")[0]
+        body = conform_query(query=query, provider=provider)
+
+        bedrock_client = self.config.get_client()
+        # invoke bedrock API
+        try:
+            response = bedrock_client.invoke_model(
+                body=json.dumps(body),
+                modelId=self.config.embed_model_name,
+                accept="application/json",
+                contentType="application/json",
+            )
+        except Exception as e:
+            raise self.wrap_error(e=e)
+
+        # format output based on provider
+        response_body = json.loads(response.get("body").read())
+        if provider == "cohere":
+            return response_body.get("embeddings")[0]
+        else:
+            # includes common provider == "amazon"
+            return response_body.get("embedding")
+
+    def embed_documents(self, elements: list[dict]) -> list[dict]:
+        embeddings = [self.embed_query(query=e.get("text", "")) for e in elements]
+        elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
+        return elements_with_embeddings
+
+
+@dataclass
+class AsyncBedrockEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
+    config: BedrockEmbeddingConfig
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    async def embed_query(self, query: str) -> list[float]:
+        """Call out to Bedrock embedding endpoint."""
+        provider = self.config.embed_model_name.split(".")[0]
+        body = conform_query(query=query, provider=provider)
+        try:
+            async with self.config.get_async_client() as bedrock_client:
+                # invoke bedrock API
+                try:
+                    response = await bedrock_client.invoke_model(
+                        body=json.dumps(body),
+                        modelId=self.config.embed_model_name,
+                        accept="application/json",
+                        contentType="application/json",
+                    )
+                except Exception as e:
+                    raise self.wrap_error(e=e)
+                async with response.get("body") as client_response:
+                    response_body = await client_response.json()
+
+            # format output based on provider
+            if provider == "cohere":
+                return response_body.get("embeddings")[0]
+            else:
+                # includes common provider == "amazon"
+                return response_body.get("embedding")
+        except Exception as e:
+            raise ValueError(f"Error raised by inference endpoint: {e}")
+
+    async def embed_documents(self, elements: list[dict]) -> list[dict]:
+        embeddings = await asyncio.gather(
+            *[self.embed_query(query=e.get("text", "")) for e in elements]
+        )
+        elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
+        return elements_with_embeddings

unstructured_ingest/embed/huggingface.py

@@ -0,0 +1,52 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+from pydantic import Field
+
+from unstructured_ingest.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if TYPE_CHECKING:
+    from sentence_transformers import SentenceTransformer
+
+
+class HuggingFaceEmbeddingConfig(EmbeddingConfig):
+    embedder_model_name: Optional[str] = Field(
+        default="sentence-transformers/all-MiniLM-L6-v2", alias="model_name"
+    )
+    embedder_model_kwargs: Optional[dict] = Field(
+        default_factory=lambda: {"device": "cpu"}, alias="model_kwargs"
+    )
+    encode_kwargs: Optional[dict] = Field(default_factory=lambda: {"normalize_embeddings": False})
+    cache_folder: Optional[str] = Field(default=None)
+
+    @requires_dependencies(
+        ["sentence_transformers"],
+        extras="embed-huggingface",
+    )
+    def get_client(self) -> "SentenceTransformer":
+        from sentence_transformers import SentenceTransformer
+
+        return SentenceTransformer(
+            model_name_or_path=self.embedder_model_name,
+            cache_folder=self.cache_folder,
+            **self.embedder_model_kwargs,
+        )
+
+
+@dataclass
+class HuggingFaceEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: HuggingFaceEmbeddingConfig
+
+    def embed_query(self, query: str) -> list[float]:
+        return self._embed_documents(texts=[query])[0]
+
+    def _embed_documents(self, texts: list[str]) -> list[list[float]]:
+        client = self.config.get_client()
+        embeddings = client.encode(texts, **self.config.encode_kwargs)
+        return embeddings.tolist()
+
+    def embed_documents(self, elements: list[dict]) -> list[dict]:
+        embeddings = self._embed_documents([e.get("text", "") for e in elements])
+        elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
+        return elements_with_embeddings