unstructured-ingest 0.0.0__py3-none-any.whl

unstructured_ingest/connector/chroma.py

@@ -0,0 +1,158 @@
+import copy
+import typing as t
+import uuid
+from dataclasses import dataclass
+
+from unstructured_ingest.enhanced_dataclass.core import _asdict
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    BaseConnectorConfig,
+    BaseDestinationConnector,
+    WriteConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.data_prep import batch_generator, flatten_dict
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if t.TYPE_CHECKING:
+    from chromadb import Collection as ChromaCollection
+
+
+@dataclass
+class ChromaAccessConfig(AccessConfig):
+    settings: t.Optional[t.Dict[str, str]] = None
+    headers: t.Optional[t.Dict[str, str]] = None
+
+
+@dataclass
+class SimpleChromaConfig(BaseConnectorConfig):
+    access_config: ChromaAccessConfig
+    collection_name: str
+    path: t.Optional[str] = None
+    tenant: t.Optional[str] = "default_tenant"
+    database: t.Optional[str] = "default_database"
+    host: t.Optional[str] = None
+    port: t.Optional[int] = None
+    ssl: bool = False
+
+
+@dataclass
+class ChromaWriteConfig(WriteConfig):
+    batch_size: int = 100
+
+
+@dataclass
+class ChromaDestinationConnector(BaseDestinationConnector):
+    write_config: ChromaWriteConfig
+    connector_config: SimpleChromaConfig
+    _collection: t.Optional["ChromaCollection"] = None
+
+    @property
+    def chroma_collection(self):
+        if self._collection is None:
+            self._collection = self.create_collection()
+        return self._collection
+
+    def initialize(self):
+        pass
+
+    @DestinationConnectionError.wrap
+    def check_connection(self):
+        _ = self.chroma_collection
+
+    def to_dict(self, **kwargs):
+        """
+        The _collection variable in this dataclass breaks deepcopy due to:
+        TypeError: cannot pickle 'module' object
+        When serializing, remove it, meaning collection data will need to be reinitialized
+        when deserialized
+        """
+        self_cp = copy.copy(self)
+        if hasattr(self_cp, "_collection"):
+            setattr(self_cp, "_collection", None)
+        return _asdict(self_cp, **kwargs)
+
+    @requires_dependencies(["chromadb"], extras="chroma")
+    def create_collection(self) -> "ChromaCollection":
+        import chromadb
+
+        if self.connector_config.path:
+            chroma_client = chromadb.PersistentClient(
+                path=self.connector_config.path,
+                settings=self.connector_config.settings,
+                tenant=self.connector_config.tenant,
+                database=self.connector_config.database,
+            )
+
+        elif self.connector_config.host and self.connector_config.port:
+            chroma_client = chromadb.HttpClient(
+                host=self.connector_config.host,
+                port=self.connector_config.port,
+                ssl=self.connector_config.ssl,
+                headers=self.connector_config.access_config.headers,
+                settings=self.connector_config.access_config.settings,
+                tenant=self.connector_config.tenant,
+                database=self.connector_config.database,
+            )
+        else:
+            raise ValueError("Chroma connector requires either path or host and port to be set.")
+
+        collection = chroma_client.get_or_create_collection(
+            name=self.connector_config.collection_name
+        )
+        return collection
+
+    @DestinationConnectionError.wrap
+    @requires_dependencies(["chromadb"], extras="chroma")
+    def upsert_batch(self, batch):
+        collection = self.chroma_collection
+
+        try:
+            # Chroma wants lists even if there is only one element
+            # Upserting to prevent duplicates
+            collection.upsert(
+                ids=batch["ids"],
+                documents=batch["documents"],
+                embeddings=batch["embeddings"],
+                metadatas=batch["metadatas"],
+            )
+        except Exception as e:
+            raise ValueError(f"chroma error: {e}") from e
+
+    @staticmethod
+    def prepare_chroma_list(chunk: t.Tuple[t.Dict[str, t.Any]]) -> t.Dict[str, t.List[t.Any]]:
+        """Helper function to break a tuple of dicts into list of parallel lists for ChromaDb.
+        ({'id':1}, {'id':2}, {'id':3}) -> {'ids':[1,2,3]}"""
+        chroma_dict = {}
+        chroma_dict["ids"] = [x.get("id") for x in chunk]
+        chroma_dict["documents"] = [x.get("document") for x in chunk]
+        chroma_dict["embeddings"] = [x.get("embedding") for x in chunk]
+        chroma_dict["metadatas"] = [x.get("metadata") for x in chunk]
+        # Make sure all lists are of the same length
+        assert (
+            len(chroma_dict["ids"])
+            == len(chroma_dict["documents"])
+            == len(chroma_dict["embeddings"])
+            == len(chroma_dict["metadatas"])
+        )
+        return chroma_dict
+
+    def write_dict(self, *args, elements_dict: t.List[t.Dict[str, t.Any]], **kwargs) -> None:
+        logger.info(f"Inserting / updating {len(elements_dict)} documents to destination ")
+
+        chroma_batch_size = self.write_config.batch_size
+
+        for chunk in batch_generator(elements_dict, chroma_batch_size):
+            self.upsert_batch(self.prepare_chroma_list(chunk))
+
+    def normalize_dict(self, element_dict: dict) -> dict:
+        element_id = element_dict.get("element_id", str(uuid.uuid4()))
+        return {
+            "id": element_id,
+            "embedding": element_dict.pop("embeddings", None),
+            "document": element_dict.pop("text", None),
+            "metadata": flatten_dict(
+                element_dict, separator="-", flatten_lists=True, remove_none=True
+            ),
+        }

unstructured_ingest/connector/clarifai.py

@@ -0,0 +1,122 @@
+import typing as t
+import uuid
+from dataclasses import dataclass, field
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    BaseConnectorConfig,
+    BaseDestinationConnector,
+    WriteConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.data_prep import flatten_dict
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if t.TYPE_CHECKING:
+    from clarifai.client.input import Inputs
+
+
+@dataclass
+class ClarifaiAccessConfig(AccessConfig):
+    api_key: str = enhanced_field(sensitive=True)
+
+
+@dataclass
+class SimpleClarifaiConfig(BaseConnectorConfig):
+    access_config: ClarifaiAccessConfig
+    app_id: str
+    user_id: str
+    dataset_id: t.Optional[str] = None
+
+
+@dataclass
+class ClarifaiWriteConfig(WriteConfig):
+    batch_size: int = 50
+
+
+@dataclass
+class ClarifaiDestinationConnector(BaseDestinationConnector):
+    write_config: ClarifaiWriteConfig
+    connector_config: SimpleClarifaiConfig
+    _client: t.Optional["Inputs"] = field(init=False, default=None)
+
+    @property
+    @requires_dependencies(["clarifai"], extras="clarifai")
+    def client(self) -> "Inputs":
+        if self._client is None:
+            from clarifai.client.input import Inputs
+
+            access_conf = self.connector_config.access_config
+            try:
+                if access_conf.api_key is not None:
+                    clarifai_pat = access_conf.api_key
+            except Exception as e:
+                raise (f"please provide clarifai PAT key : {e}")
+
+            self._client = Inputs(
+                app_id=self.connector_config.app_id,
+                user_id=self.connector_config.user_id,
+                pat=clarifai_pat,
+            )
+        return self._client
+
+    @requires_dependencies(["clarifai"], extras="clarifai")
+    @DestinationConnectionError.wrap
+    def initialize(self):
+        _ = self.client
+
+    def check_connection(self):
+        try:
+            _ = [inp for inp in self.client.list_inputs(page_no=1, per_page=1)]  # noqa: C416
+        except Exception as e:
+            logger.error(f"Failed to validate connection {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def normalize_dict(self, element_dict: dict) -> dict:
+        """Modifying schema of the dict in order to compile with clarifai input formats"""
+        return {
+            "input_id": str(uuid.uuid4().hex),
+            "text": element_dict.pop("text", None),
+            "metadata": {
+                **flatten_dict(
+                    element_dict,
+                    separator="_",
+                    flatten_lists=True,
+                    remove_none=True,
+                ),
+            },
+        }
+
+    def write_dict(self, *args, elements_dict: t.List[t.Dict[str, t.Any]], **kwargs) -> None:
+        from google.protobuf.struct_pb2 import Struct
+
+        logger.info(
+            f"writing {len(elements_dict)} objects to destination "
+            f"app {self.connector_config.app_id} "
+        )
+        try:
+            batch_size = self.write_config.batch_size
+            for idx in range(0, len(elements_dict), batch_size):
+                batch_dict = elements_dict[idx : batch_size + idx]
+                input_batch = []
+                for elem in batch_dict:
+                    meta_struct = Struct()
+                    meta_struct.update(elem["metadata"])
+                    input_batch.append(
+                        self._client.get_text_input(
+                            input_id=elem["input_id"],
+                            raw_text=elem["text"],
+                            dataset_id=self.connector_config.dataset_id,
+                            metadata=meta_struct,
+                        )
+                    )
+                result_id = self._client.upload_inputs(inputs=input_batch)
+                logger.debug(
+                    f"Input posted successfully into {self.connector_config.app_id}. \
+                    Result id: {result_id}"
+                )
+
+        except Exception as e:
+            raise e

unstructured_ingest/connector/confluence.py

@@ -0,0 +1,285 @@
+import math
+import typing as t
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+import requests
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.error import SourceConnectionError, SourceConnectionNetworkError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    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 atlassian import Confluence
+
+
+@dataclass
+class ConfluenceAccessConfig(AccessConfig):
+    api_token: str = enhanced_field(sensitive=True)
+
+
+@dataclass
+class SimpleConfluenceConfig(BaseConnectorConfig):
+    """Connector config where:
+    user_email is the email to authenticate into Confluence Cloud,
+    api_token is the api token to authenticate into Confluence Cloud,
+    and url is the URL pointing to the Confluence Cloud instance.
+
+    Check https://developer.atlassian.com/cloud/confluence/basic-auth-for-rest-apis/
+    for more info on the api_token.
+    """
+
+    user_email: str
+    access_config: ConfluenceAccessConfig
+    url: str
+    max_num_of_spaces: int = 500
+    max_num_of_docs_from_each_space: int = 100
+    spaces: t.List[str] = field(default_factory=list)
+
+
+@dataclass
+class ConfluenceDocumentMeta:
+    """Metadata specifying:
+    id for the confluence space that the document locates in,
+    and the id of document that is being reached to.
+    """
+
+    space_id: str
+    document_id: str
+
+
+def scroll_wrapper(func):
+    def wrapper(*args, **kwargs):
+        """Wraps a function to obtain scroll functionality."""
+        number_of_items_to_fetch = kwargs["number_of_items_to_fetch"]
+        del kwargs["number_of_items_to_fetch"]
+
+        kwargs["limit"] = min(100, number_of_items_to_fetch)
+        kwargs["start"] = kwargs.get("start", 0)
+
+        all_results = []
+        num_iterations = math.ceil(number_of_items_to_fetch / kwargs["limit"])
+
+        for _ in range(num_iterations):
+            response = func(*args, **kwargs)
+            if isinstance(response, list):
+                all_results += func(*args, **kwargs)
+            elif isinstance(response, dict):
+                all_results += func(*args, **kwargs)["results"]
+
+            kwargs["start"] += kwargs["limit"]
+
+        return all_results[:number_of_items_to_fetch]
+
+    return wrapper
+
+
+@dataclass
+class ConfluenceIngestDoc(IngestDocCleanupMixin, BaseSingleIngestDoc):
+    """Class encapsulating fetching a doc and writing processed results (but not
+    doing the processing).
+
+    Current implementation creates a Confluence connection object
+    to fetch each doc, rather than creating a it for each thread.
+    """
+
+    connector_config: SimpleConfluenceConfig
+    document_meta: ConfluenceDocumentMeta
+    registry_name: str = "confluence"
+
+    # TODO: remove one of filename or _tmp_download_file, using a wrapper
+    @property
+    def filename(self):
+        if not self.read_config.download_dir:
+            return None
+        return (
+            Path(self.read_config.download_dir)
+            / self.document_meta.space_id
+            / f"{self.document_meta.document_id}.html"
+        ).resolve()
+
+    @property
+    def _output_filename(self):
+        """Create output file path based on output directory, space id and document id."""
+        output_file = f"{self.document_meta.document_id}.json"
+        return Path(self.processor_config.output_dir) / self.document_meta.space_id / output_file
+
+    @property
+    def record_locator(self) -> t.Optional[t.Dict[str, t.Any]]:
+        return {
+            "url": self.connector_config.url,
+            "page_id": self.document_meta.document_id,
+        }
+
+    @SourceConnectionNetworkError.wrap
+    @requires_dependencies(["atlassian"], extras="Confluence")
+    def _get_page(self):
+        from atlassian import Confluence
+        from atlassian.errors import ApiError
+
+        try:
+            confluence = Confluence(
+                self.connector_config.url,
+                username=self.connector_config.user_email,
+                password=self.connector_config.access_config.api_token,
+            )
+            result = confluence.get_page_by_id(
+                page_id=self.document_meta.document_id,
+                expand="history.lastUpdated,version,body.view",
+            )
+        except ApiError as e:
+            logger.error(e)
+            return None
+        return result
+
+    def update_source_metadata(self, **kwargs):
+        """Fetches file metadata from the current page."""
+        page = kwargs.get("page", self._get_page())
+        if page is None:
+            self.source_metadata = SourceMetadata(
+                exists=False,
+            )
+            return
+        document_history = page["history"]
+        date_created = datetime.strptime(
+            document_history["createdDate"],
+            "%Y-%m-%dT%H:%M:%S.%fZ",
+        ).isoformat()
+        if last_updated := document_history.get("lastUpdated", {}).get("when", ""):
+            date_modified = datetime.strptime(
+                last_updated,
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            ).isoformat()
+        else:
+            date_modified = date_created
+        version = page["version"]["number"]
+        self.source_metadata = SourceMetadata(
+            date_created=date_created,
+            date_modified=date_modified,
+            version=version,
+            source_url=page["_links"].get("self", None),
+            exists=True,
+        )
+
+    @SourceConnectionError.wrap
+    @requires_dependencies(["atlassian"], extras="confluence")
+    @BaseSingleIngestDoc.skip_if_file_exists
+    def get_file(self):
+        # TODO: instead of having a separate connection object for each doc,
+        # have a separate connection object for each process
+
+        result = self._get_page()
+        self.update_source_metadata(page=result)
+        if result is None:
+            raise ValueError(f"Failed to retrieve page with ID {self.document_meta.document_id}")
+        self.document = result["body"]["view"]["value"]
+        self.filename.parent.mkdir(parents=True, exist_ok=True)
+        with open(self.filename, "w", encoding="utf8") as f:
+            f.write(self.document)
+
+
+@dataclass
+class ConfluenceSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
+    """Fetches body fields from all documents within all spaces in a Confluence Cloud instance."""
+
+    connector_config: SimpleConfluenceConfig
+    _confluence: t.Optional["Confluence"] = field(init=False, default=None)
+
+    @property
+    def confluence(self) -> "Confluence":
+        from atlassian import Confluence
+
+        if self._confluence is None:
+            self._confluence = Confluence(
+                url=self.connector_config.url,
+                username=self.connector_config.user_email,
+                password=self.connector_config.access_config.api_token,
+            )
+        return self._confluence
+
+    @requires_dependencies(["atlassian"], extras="Confluence")
+    def check_connection(self):
+        url = "rest/api/space"
+        try:
+            self.confluence.request(method="HEAD", path=url)
+        except requests.HTTPError as http_error:
+            logger.error(f"failed to validate connection: {http_error}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {http_error}")
+
+    @requires_dependencies(["atlassian"], extras="Confluence")
+    def initialize(self):
+        self.list_of_spaces = None
+        if self.connector_config.spaces:
+            self.list_of_spaces = self.connector_config.spaces
+            if self.connector_config.max_num_of_spaces:
+                logger.warning(
+                    """--confluence-list-of-spaces and --confluence-num-of-spaces cannot
+                    be used at the same time. Connector will only fetch the
+                    --confluence-list-of-spaces that you've provided.""",
+                )
+
+    @requires_dependencies(["atlassian"], extras="Confluence")
+    def _get_space_ids(self):
+        """Fetches spaces in a confluence domain."""
+
+        get_spaces_with_scroll = scroll_wrapper(self.confluence.get_all_spaces)
+
+        all_results = get_spaces_with_scroll(
+            number_of_items_to_fetch=self.connector_config.max_num_of_spaces,
+        )
+
+        space_ids = [space["key"] for space in all_results]
+        return space_ids
+
+    @requires_dependencies(["atlassian"], extras="Confluence")
+    def _get_docs_ids_within_one_space(
+        self,
+        space_id: str,
+        content_type: str = "page",
+    ):
+        get_pages_with_scroll = scroll_wrapper(self.confluence.get_all_pages_from_space)
+        results = get_pages_with_scroll(
+            space=space_id,
+            number_of_items_to_fetch=self.connector_config.max_num_of_docs_from_each_space,
+            content_type=content_type,
+        )
+
+        doc_ids = [(space_id, doc["id"]) for doc in results]
+        return doc_ids
+
+    @requires_dependencies(["atlassian"], extras="Confluence")
+    def _get_doc_ids_within_spaces(self):
+        space_ids = self._get_space_ids() if not self.list_of_spaces else self.list_of_spaces
+
+        doc_ids_all = [self._get_docs_ids_within_one_space(space_id=id) for id in space_ids]
+
+        doc_ids_flattened = [
+            (space_id, doc_id)
+            for doc_ids_space in doc_ids_all
+            for space_id, doc_id in doc_ids_space
+        ]
+        return doc_ids_flattened
+
+    def get_ingest_docs(self):
+        """Fetches all documents in a confluence space."""
+        doc_ids = self._get_doc_ids_within_spaces()
+        return [
+            ConfluenceIngestDoc(
+                connector_config=self.connector_config,
+                processor_config=self.processor_config,
+                read_config=self.read_config,
+                document_meta=ConfluenceDocumentMeta(space_id, doc_id),
+            )
+            for space_id, doc_id in doc_ids
+        ]

unstructured_ingest/connector/databricks_volumes.py

@@ -0,0 +1,137 @@
+import copy
+import json
+import os
+import typing as t
+from dataclasses import dataclass, field
+from io import BytesIO
+from pathlib import PurePath
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.enhanced_dataclass.core import _asdict
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    BaseConnectorConfig,
+    BaseDestinationConnector,
+    BaseSingleIngestDoc,
+    WriteConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if t.TYPE_CHECKING:
+    from databricks.sdk import WorkspaceClient
+
+
+@dataclass
+class DatabricksVolumesAccessConfig(AccessConfig):
+    account_id: t.Optional[str] = None
+    username: t.Optional[str] = None
+    password: t.Optional[str] = enhanced_field(default=None, sensitive=True)
+    client_id: t.Optional[str] = None
+    client_secret: t.Optional[str] = enhanced_field(default=None, sensitive=True)
+    token: t.Optional[str] = enhanced_field(default=None, sensitive=True)
+    profile: t.Optional[str] = None
+    azure_workspace_resource_id: t.Optional[str] = None
+    azure_client_secret: t.Optional[str] = enhanced_field(default=None, sensitive=True)
+    azure_client_id: t.Optional[str] = None
+    azure_tenant_id: t.Optional[str] = None
+    azure_environment: t.Optional[str] = None
+    auth_type: t.Optional[str] = None
+    cluster_id: t.Optional[str] = None
+    google_credentials: t.Optional[str] = None
+    google_service_account: t.Optional[str] = None
+
+
+@dataclass
+class SimpleDatabricksVolumesConfig(BaseConnectorConfig):
+    access_config: DatabricksVolumesAccessConfig
+    host: t.Optional[str] = None
+
+
+@dataclass
+class DatabricksVolumesWriteConfig(WriteConfig):
+    volume: str
+    catalog: str
+    volume_path: t.Optional[str] = None
+    overwrite: bool = False
+    encoding: str = "utf-8"
+    schema: str = "default"
+
+    @property
+    def path(self) -> str:
+        path = f"/Volumes/{self.catalog}/{self.schema}/{self.volume}"
+        if self.volume_path:
+            path = f"{path}/{self.volume_path}"
+        return path
+
+
+@dataclass
+class DatabricksVolumesDestinationConnector(BaseDestinationConnector):
+    write_config: DatabricksVolumesWriteConfig
+    connector_config: SimpleDatabricksVolumesConfig
+    _client: t.Optional["WorkspaceClient"] = field(init=False, default=None)
+
+    def to_dict(self, **kwargs):
+        self_cp = copy.copy(self)
+        if hasattr(self_cp, "_client"):
+            setattr(self_cp, "_client", None)
+        return _asdict(self_cp, **kwargs)
+
+    @requires_dependencies(dependencies=["databricks.sdk"], extras="databricks-volumes")
+    def generate_client(self) -> "WorkspaceClient":
+        from databricks.sdk import WorkspaceClient
+
+        return WorkspaceClient(
+            host=self.connector_config.host, **self.connector_config.access_config.to_dict()
+        )
+
+    @property
+    def client(self) -> "WorkspaceClient":
+        if self._client is None:
+            self._client = self.generate_client()
+        return self._client
+
+    def check_connection(self):
+        try:
+            assert self.client.current_user.me().active
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def initialize(self):
+        _ = self.client
+
+    def write_dict(
+        self,
+        *args,
+        elements_dict: t.List[t.Dict[str, t.Any]],
+        filename: t.Optional[str] = None,
+        indent: int = 4,
+        encoding: str = "utf-8",
+        **kwargs,
+    ) -> None:
+        output_folder = self.write_config.path
+        output_folder = os.path.join(output_folder)  # Make sure folder ends with file seperator
+        filename = (
+            filename.strip(os.sep) if filename else filename
+        )  # Make sure filename doesn't begin with file seperator
+        output_path = str(PurePath(output_folder, filename)) if filename else output_folder
+        logger.debug(f"uploading content to {output_path}")
+        self.client.files.upload(
+            file_path=output_path,
+            contents=BytesIO(json.dumps(elements_dict).encode(encoding=self.write_config.encoding)),
+            overwrite=self.write_config.overwrite,
+        )
+
+    def get_elements_dict(self, docs: t.List[BaseSingleIngestDoc]) -> t.List[t.Dict[str, t.Any]]:
+        pass
+
+    def write(self, docs: t.List[BaseSingleIngestDoc]) -> None:
+        for doc in docs:
+            file_path = doc.base_output_filename
+            filename = file_path if file_path else None
+            with open(doc._output_filename) as json_file:
+                logger.debug(f"uploading content from {doc._output_filename}")
+                json_list = json.load(json_file)
+                self.write_dict(elements_dict=json_list, filename=filename)