unstructured-ingest 0.3.13__py3-none-any.whl

unstructured_ingest/connector/slack.py

@@ -0,0 +1,224 @@
+import typing as t
+import xml.etree.ElementTree as ET
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+
+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.data_prep import validate_date_args
+from unstructured_ingest.utils.dep_check import (
+    requires_dependencies,
+)
+
+DATE_FORMATS = ("%Y-%m-%d", "%Y-%m-%dT%H:%M:%S", "%Y-%m-%dT%H:%M:%S%z")
+
+
+@dataclass
+class SlackAccessConfig(AccessConfig):
+    token: str = enhanced_field(sensitive=True)
+
+
+@dataclass
+class SimpleSlackConfig(BaseConnectorConfig):
+    """Connector config to process all messages by channel id's."""
+
+    access_config: SlackAccessConfig
+    channels: t.List[str]
+    start_date: t.Optional[str] = None
+    end_date: t.Optional[str] = None
+
+    def validate_inputs(self):
+        oldest_valid = True
+        latest_valid = True
+
+        if self.start_date:
+            oldest_valid = validate_date_args(self.start_date)
+
+        if self.end_date:
+            latest_valid = validate_date_args(self.end_date)
+
+        return oldest_valid, latest_valid
+
+    def __post_init__(self):
+        oldest_valid, latest_valid = self.validate_inputs()
+        if not oldest_valid and not latest_valid:
+            raise ValueError(
+                "Start and/or End dates are not valid. ",
+            )
+
+
+@dataclass
+class SlackIngestDoc(IngestDocCleanupMixin, BaseSingleIngestDoc):
+    """Class encapsulating fetching a doc and writing processed results (but not
+    doing the processing!).
+
+    Also includes a cleanup method. When things go wrong and the cleanup
+    method is not called, the file is left behind on the filesystem to assist debugging.
+    """
+
+    connector_config: SimpleSlackConfig
+    channel: str
+    registry_name: str = "slack"
+
+    # NOTE(crag): probably doesn't matter,  but intentionally not defining tmp_download_file
+    # __post_init__ for multiprocessing simplicity (no Path objects in initially
+    # instantiated object)
+    def _tmp_download_file(self):
+        channel_file = self.channel + ".xml"
+        return Path(self.read_config.download_dir) / channel_file
+
+    @property
+    def _output_filename(self):
+        output_file = self.channel + ".json"
+        return Path(self.processor_config.output_dir) / output_file
+
+    @property
+    def version(self) -> t.Optional[str]:
+        return None
+
+    @property
+    def source_url(self) -> t.Optional[str]:
+        return None
+
+    def _create_full_tmp_dir_path(self):
+        self._tmp_download_file().parent.mkdir(parents=True, exist_ok=True)
+
+    @SourceConnectionNetworkError.wrap
+    @requires_dependencies(dependencies=["slack_sdk"], extras="slack")
+    def _fetch_messages(self):
+        from slack_sdk import WebClient
+
+        self.client = WebClient(token=self.connector_config.access_config.token)
+        oldest = "0"
+        latest = "0"
+        if self.connector_config.start_date:
+            oldest = self.convert_datetime(self.connector_config.start_date)
+
+        if self.connector_config.end_date:
+            latest = self.convert_datetime(self.connector_config.end_date)
+
+        result = self.client.conversations_history(
+            channel=self.channel,
+            oldest=oldest,
+            latest=latest,
+        )
+        return result
+
+    def update_source_metadata(self, **kwargs):
+        result = kwargs.get("result", self._fetch_messages())
+        if result is None:
+            self.source_metadata = SourceMetadata(
+                exists=True,
+            )
+            return
+        timestamps = [m["ts"] for m in result["messages"]]
+        timestamps.sort()
+        date_created = None
+        date_modified = None
+        if len(timestamps) > 0:
+            date_created = datetime.fromtimestamp(float(timestamps[0])).isoformat()
+            date_modified = datetime.fromtimestamp(
+                float(timestamps[len(timestamps) - 1]),
+            ).isoformat()
+
+        self.source_metadata = SourceMetadata(
+            date_created=date_created,
+            date_modified=date_modified,
+            exists=True,
+        )
+
+    @SourceConnectionError.wrap
+    @BaseSingleIngestDoc.skip_if_file_exists
+    @requires_dependencies(dependencies=["slack_sdk"], extras="slack")
+    def get_file(self):
+        from slack_sdk.errors import SlackApiError
+
+        """Fetches the data from a slack channel and stores it locally."""
+
+        self._create_full_tmp_dir_path()
+
+        result = self._fetch_messages()
+        self.update_source_metadata(result=result)
+        root = ET.Element("messages")
+        for message in result["messages"]:
+            message_elem = ET.SubElement(root, "message")
+            text_elem = ET.SubElement(message_elem, "text")
+            text_elem.text = message.get("text")
+
+            cursor = None
+            while True:
+                try:
+                    response = self.client.conversations_replies(
+                        channel=self.channel,
+                        ts=message["ts"],
+                        cursor=cursor,
+                    )
+
+                    for reply in response["messages"]:
+                        reply_msg = reply.get("text")
+                        text_elem.text = "".join([str(text_elem.text), " <reply> ", reply_msg])
+
+                    if not response["has_more"]:
+                        break
+
+                    cursor = response["response_metadata"]["next_cursor"]
+
+                except SlackApiError as e:
+                    logger.error(f"Error retrieving replies: {e.response['error']}")
+        tree = ET.ElementTree(root)
+        tree.write(self._tmp_download_file(), encoding="utf-8", xml_declaration=True)
+
+    def convert_datetime(self, date_time):
+        for format in DATE_FORMATS:
+            try:
+                return datetime.strptime(date_time, format).timestamp()
+            except ValueError:
+                pass
+
+    @property
+    def filename(self):
+        """The filename of the file created from a slack channel"""
+        return self._tmp_download_file()
+
+
+class SlackSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
+    """Objects of this class support fetching document(s) from"""
+
+    connector_config: SimpleSlackConfig
+
+    @requires_dependencies(dependencies=["slack_sdk"], extras="slack")
+    def check_connection(self):
+        from slack_sdk import WebClient
+        from slack_sdk.errors import SlackClientError
+
+        try:
+            client = WebClient(token=self.connector_config.access_config.token)
+            client.users_identity()
+        except SlackClientError as slack_error:
+            logger.error(f"failed to validate connection: {slack_error}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {slack_error}")
+
+    def initialize(self):
+        """Verify that can get metadata for an object, validates connections info."""
+
+    def get_ingest_docs(self):
+        return [
+            SlackIngestDoc(
+                connector_config=self.connector_config,
+                processor_config=self.processor_config,
+                read_config=self.read_config,
+                channel=channel,
+            )
+            for channel in self.connector_config.channels
+        ]

unstructured_ingest/connector/sql.py

@@ -0,0 +1,199 @@
+import copy
+import json
+import typing as t
+import uuid
+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
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    BaseConnectorConfig,
+    BaseDestinationConnector,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+ELEMENTS_TABLE_NAME = "elements"
+
+
+@dataclass
+class SqlAccessConfig(AccessConfig):
+    username: t.Optional[str]
+    password: t.Optional[str] = enhanced_field(sensitive=True)
+
+
+@dataclass
+class SimpleSqlConfig(BaseConnectorConfig):
+    db_type: t.Optional[str]
+    host: t.Optional[str]
+    database: t.Optional[str]
+    port: t.Optional[int]
+    access_config: SqlAccessConfig
+
+    def __post_init__(self):
+        if (self.db_type == "sqlite") and (self.database is None):
+            raise ValueError(
+                "A sqlite connection requires a path to a *.db file "
+                "through the `database` argument"
+            )
+
+    @property
+    def connection(self):
+        if self.db_type == "postgresql":
+            return self._make_psycopg_connection
+        elif self.db_type == "sqlite":
+            return self._make_sqlite_connection
+        raise ValueError(f"Unsupported database {self.db_type} connection.")
+
+    def _make_sqlite_connection(self):
+        from sqlite3 import connect
+
+        return connect(database=self.database)
+
+    @requires_dependencies(["psycopg2"], extras="postgres")
+    def _make_psycopg_connection(self):
+        from psycopg2 import connect
+
+        return connect(
+            user=self.access_config.username,
+            password=self.access_config.password,
+            dbname=self.database,
+            host=self.host,
+            port=self.port,
+        )
+
+
+@dataclass
+class SqlDestinationConnector(BaseDestinationConnector):
+    connector_config: SimpleSqlConfig
+    _client: t.Optional[t.Any] = 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
+    def client(self):
+        if self._client is None:
+            self._client = self.connector_config.connection()
+        return self._client
+
+    @DestinationConnectionError.wrap
+    def initialize(self):
+        _ = self.client
+
+    def check_connection(self):
+        try:
+            cursor = self.client.cursor()
+            cursor.execute("SELECT 1;")
+            cursor.close()
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def conform_dict(self, data: dict) -> None:
+        """
+        Updates the element dictionary to conform to the sql schema
+        """
+        from datetime import datetime
+
+        data["id"] = str(uuid.uuid4())
+
+        # 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 (embeddings := data.get("embeddings")) and (
+            self.connector_config.db_type != "postgresql"
+        ):
+            data["embeddings"] = str(json.dumps(embeddings))
+
+        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)
+
+        if sent_from := data.get("metadata", {}).get("sent_from", {}):
+            data["metadata"]["sent_from"] = str(json.dumps(sent_from))
+
+        if sent_to := data.get("metadata", {}).get("sent_to", {}):
+            data["metadata"]["sent_to"] = str(json.dumps(sent_to))
+
+        # Datetime formatting
+        if date_created := data.get("metadata", {}).get("data_source", {}).get("date_created"):
+            data["metadata"]["data_source"]["date_created"] = datetime.fromisoformat(date_created)
+
+        if date_modified := data.get("metadata", {}).get("data_source", {}).get("date_modified"):
+            data["metadata"]["data_source"]["date_modified"] = datetime.fromisoformat(date_modified)
+
+        if date_processed := data.get("metadata", {}).get("data_source", {}).get("date_processed"):
+            data["metadata"]["data_source"]["date_processed"] = datetime.fromisoformat(
+                date_processed
+            )
+
+        if last_modified := data.get("metadata", {}).get("last_modified", {}):
+            data["metadata"]["last_modified"] = datetime.fromisoformat(last_modified)
+
+        # 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))
+
+        if data.get("metadata", {}).get("data_source", None):
+            data.update(data.get("metadata", {}).pop("data_source", None))
+        if data.get("metadata", {}).get("coordinates", None):
+            data.update(data.get("metadata", {}).pop("coordinates", None))
+        if data.get("metadata", {}):
+            data.update(data.pop("metadata", None))
+
+    @DestinationConnectionError.wrap
+    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 database {self.connector_config.database} "
+            f"at {self.connector_config.host}"
+        )
+
+        with self.client as conn:
+            cursor = conn.cursor()
+
+            # Since we have no guarantee that each element will have the same keys
+            # we insert each element individually
+            for elem in elements_dict:
+                query = f"INSERT INTO {ELEMENTS_TABLE_NAME} ({','.join(elem.keys())}) \
+                VALUES({','.join(['?' if self.connector_config.db_type=='sqlite' else '%s' for x in elem])})"  # noqa E501
+                values = []
+                for v in elem.values():
+                    if self.connector_config.db_type == "sqlite" and isinstance(v, list):
+                        values.append(json.dumps(v))
+                    else:
+                        values.append(v)
+                cursor.execute(query, values)
+
+            conn.commit()
+            cursor.close()
+
+        # Leaving contexts doesn't close the connection, so doing it here
+        conn.close()

unstructured_ingest/connector/vectara.py

@@ -0,0 +1,253 @@
+import datetime
+import json
+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,
+    BaseIngestDoc,
+    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
+
+BASE_URL = "https://api.vectara.io/v1"
+
+
+@dataclass
+class VectaraAccessConfig(AccessConfig):
+    oauth_client_id: str = enhanced_field(sensitive=True)
+    oauth_secret: str = enhanced_field(sensitive=True)
+
+
+@dataclass
+class SimpleVectaraConfig(BaseConnectorConfig):
+    access_config: VectaraAccessConfig
+    customer_id: str
+    corpus_name: t.Optional[str] = None
+    corpus_id: t.Optional[str] = None
+    token_url: str = "https://vectara-prod-{}.auth.us-west-2.amazoncognito.com/oauth2/token"
+
+
+@dataclass
+class VectaraDestinationConnector(BaseDestinationConnector):
+    write_config: WriteConfig
+    connector_config: SimpleVectaraConfig
+    _jwt_token: t.Optional[str] = field(init=False, default=None)
+    _jwt_token_expires_ts: t.Optional[float] = field(init=False, default=None)
+
+    @property
+    def jwt_token(self):
+        if (
+            not self._jwt_token
+            or self._jwt_token_expires_ts - datetime.datetime.now().timestamp() <= 60
+        ):
+            self._jwt_token = self._get_jwt_token()
+        return self._jwt_token
+
+    @DestinationConnectionError.wrap
+    def vectara(self):
+        """
+        Check the connection for Vectara and validate corpus exists.
+        - If more than one corpus with the same name exists - then return a message
+        - If exactly one corpus exists with this name - use it.
+        - If does not exist - create it.
+        """
+        try:
+            # Get token if not already set
+            self.jwt_token
+
+            list_corpora_response = self._request(
+                endpoint="list-corpora",
+                data={"numResults": 1, "filter": self.connector_config.corpus_name},
+            )
+
+            possible_corpora_ids_names_map = {
+                corpus.get("id"): corpus.get("name")
+                for corpus in list_corpora_response.get("corpus")
+                if corpus.get("name") == self.connector_config.corpus_name
+            }
+
+            if len(possible_corpora_ids_names_map) > 1:
+                return f"Multiple Corpora exist with name {self.connector_config.corpus_name}"
+            if len(possible_corpora_ids_names_map) == 1:
+                self.connector_config.corpus_id = list(possible_corpora_ids_names_map.keys())[0]
+            else:
+                data = {
+                    "corpus": {
+                        "name": self.connector_config.corpus_name,
+                    }
+                }
+                create_corpus_response = self._request(endpoint="create-corpus", data=data)
+                self.connector_config.corpus_id = create_corpus_response.get("corpusId")
+
+        except Exception as e:
+            logger.error(f"failed to create Vectara connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to create Vectara connection: {e}")
+
+    def initialize(self):
+        self.vectara()
+
+    @requires_dependencies(["requests"], extras="vectara")
+    def _request(
+        self,
+        endpoint: str,
+        http_method: str = "POST",
+        params: t.Mapping[str, t.Any] = None,
+        data: t.Mapping[str, t.Any] = None,
+    ):
+        import requests
+
+        url = f"{BASE_URL}/{endpoint}"
+
+        headers = {
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+            "Authorization": f"Bearer {self.jwt_token}",
+            "customer-id": self.connector_config.customer_id,
+            "X-source": "unstructured",
+        }
+
+        response = requests.request(
+            method=http_method, url=url, headers=headers, params=params, data=json.dumps(data)
+        )
+        response.raise_for_status()
+        return response.json()
+
+    # Get Oauth2 JWT token
+    @requires_dependencies(["requests"], extras="vectara")
+    def _get_jwt_token(self):
+        import requests
+
+        """Connect to the server and get a JWT token."""
+        token_endpoint = self.connector_config.token_url.format(self.connector_config.customer_id)
+        headers = {
+            "Content-Type": "application/x-www-form-urlencoded",
+        }
+        data = {
+            "grant_type": "client_credentials",
+            "client_id": self.connector_config.access_config.oauth_client_id,
+            "client_secret": self.connector_config.access_config.oauth_secret,
+        }
+
+        response = requests.request(method="POST", url=token_endpoint, headers=headers, data=data)
+        response.raise_for_status()
+        response_json = response.json()
+
+        request_time = datetime.datetime.now().timestamp()
+        self._jwt_token_expires_ts = request_time + response_json.get("expires_in")
+
+        return response_json.get("access_token")
+
+    @DestinationConnectionError.wrap
+    def check_connection(self):
+        try:
+            self.vectara()
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def _delete_doc(self, doc_id: str) -> None:
+        """
+        Delete a document from the Vectara corpus.
+
+        Args:
+            url (str): URL of the page to delete.
+            doc_id (str): ID of the document to delete.
+        """
+        body = {
+            "customer_id": self.connector_config.customer_id,
+            "corpus_id": self.connector_config.corpus_id,
+            "document_id": doc_id,
+        }
+        self._request(endpoint="delete-doc", data=body)
+
+    def _index_document(self, document: t.Dict[str, t.Any]) -> None:
+        """
+        Index a document (by uploading it to the Vectara corpus) from the document dictionary
+        """
+        body = {
+            "customer_id": self.connector_config.customer_id,
+            "corpus_id": self.connector_config.corpus_id,
+            "document": document,
+        }
+
+        try:
+            result = self._request(endpoint="index", data=body, http_method="POST")
+        except Exception as e:
+            logger.info(f"exception {e} while indexing document {document['documentId']}")
+            return
+
+        if (
+            "status" in result
+            and result["status"]
+            and (
+                "ALREADY_EXISTS" in result["status"]["code"]
+                or (
+                    "CONFLICT" in result["status"]["code"]
+                    and "Indexing doesn't support updating documents"
+                    in result["status"]["statusDetail"]
+                )
+            )
+        ):
+            logger.info(f"document {document['documentId']} already exists, re-indexing")
+            self._delete_doc(document["documentId"])
+            result = self._request(endpoint="index", data=body, http_method="POST")
+            return
+
+        if "status" in result and result["status"] and "OK" in result["status"]["code"]:
+            logger.info(f"indexing document {document['documentId']} succeeded")
+        else:
+            logger.info(f"indexing document {document['documentId']} failed, response = {result}")
+
+    def write_dict(self, *args, docs_list: t.List[t.Dict[str, t.Any]], **kwargs) -> None:
+        logger.info(f"inserting / updating {len(docs_list)} documents to Vectara ")
+        for vdoc in docs_list:
+            self._index_document(vdoc)
+
+    def write(self, docs: t.List[BaseIngestDoc]) -> None:
+        docs_list: t.Dict[t.Dict[str, t.Any]] = []
+
+        def get_metadata(element) -> t.Dict[str, t.Any]:
+            """
+            Select which meta-data fields to include and optionally map them to a new new.
+            remove the "metadata-" prefix from the keys
+            """
+            metadata_map = {
+                "page_number": "page_number",
+                "data_source-url": "url",
+                "filename": "filename",
+                "filetype": "filetype",
+                "last_modified": "last_modified",
+            }
+            md = flatten_dict(element, separator="-", flatten_lists=True)
+            md = {k.replace("metadata-", ""): v for k, v in md.items()}
+            md = {metadata_map[k]: v for k, v in md.items() if k in metadata_map}
+            return md
+
+        for doc in docs:
+            local_path = doc._output_filename
+            with open(local_path) as json_file:
+                dict_content = json.load(json_file)
+                vdoc = {
+                    "documentId": str(uuid.uuid4()),
+                    "title": dict_content[0].get("metadata", {}).get("data_source", {}).get("url"),
+                    "section": [
+                        {
+                            "text": element.pop("text", None),
+                            "metadataJson": json.dumps(get_metadata(element)),
+                        }
+                        for element in dict_content
+                    ],
+                }
+                logger.info(
+                    f"Extending {len(vdoc)} json elements from content in {local_path}",
+                )
+                docs_list.append(vdoc)
+        self.write_dict(docs_list=docs_list)