unstructured-ingest 0.0.0__py3-none-any.whl

unstructured_ingest/connector/astra.py

@@ -0,0 +1,237 @@
+import copy
+import typing as t
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from unstructured_ingest import __name__ as integration_name
+from unstructured_ingest.__version__ import __version__ as integration_version
+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,
+    BaseSingleIngestDoc,
+    BaseSourceConnector,
+    IngestDocCleanupMixin,
+    SourceConnectorCleanupMixin,
+    SourceMetadata,
+    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 astrapy.db import AstraDB, AstraDBCollection
+
+NON_INDEXED_FIELDS = ["metadata._node_content", "content"]
+
+
+@dataclass
+class AstraAccessConfig(AccessConfig):
+    token: str = enhanced_field(sensitive=True)
+    api_endpoint: str = enhanced_field(sensitive=True)
+
+
+@dataclass
+class SimpleAstraConfig(BaseConnectorConfig):
+    access_config: AstraAccessConfig
+    collection_name: str
+    namespace: t.Optional[str] = None
+
+
+@dataclass
+class AstraIngestDoc(IngestDocCleanupMixin, BaseSingleIngestDoc):
+    connector_config: SimpleAstraConfig
+    metadata: t.Dict[str, str] = field(default_factory=dict)
+    registry_name: str = "astra"
+
+    @property
+    def filename(self):
+        return (
+            Path(self.read_config.download_dir)
+            / self.connector_config.collection_name
+            / f"{self.metadata['_id']}.txt"
+        ).resolve()
+
+    @property
+    def _output_filename(self):
+        return (
+            Path(self.processor_config.output_dir)
+            / self.connector_config.collection_name
+            / f"{self.metadata['_id']}.json"
+        ).resolve()
+
+    def update_source_metadata(self, **kwargs):
+        if not self.metadata:
+            self.source_metadata = SourceMetadata(
+                exists=False,
+            )
+            return
+        self.source_metadata = SourceMetadata(
+            exists=True,
+        )
+
+    @SourceConnectionError.wrap
+    @requires_dependencies(["astrapy"], extras="astra")
+    @BaseSingleIngestDoc.skip_if_file_exists
+    def get_file(self):
+        self.filename.parent.mkdir(parents=True, exist_ok=True)
+
+        flattened_dict = flatten_dict(dictionary=self.metadata)
+        str_values = [str(value) for value in flattened_dict.values()]
+        concatenated_values = "\n".join(str_values)
+
+        with open(self.filename, "w") as f:
+            f.write(concatenated_values)
+
+
+@dataclass
+class AstraSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
+    connector_config: SimpleAstraConfig
+    _astra_db: t.Optional["AstraDB"] = field(init=False, default=None)
+    _astra_db_collection: t.Optional["AstraDBCollection"] = field(init=False, default=None)
+
+    @property
+    @requires_dependencies(["astrapy"], extras="astra")
+    def astra_db_collection(self) -> "AstraDBCollection":
+        if self._astra_db_collection is None:
+            from astrapy.db import AstraDB
+
+            # Build the Astra DB object.
+            # caller_name/version for AstraDB tracking
+            self._astra_db = AstraDB(
+                api_endpoint=self.connector_config.access_config.api_endpoint,
+                token=self.connector_config.access_config.token,
+                namespace=self.connector_config.namespace,
+                caller_name=integration_name,
+                caller_version=integration_version,
+            )
+
+            # Create and connect to the collection
+            self._astra_db_collection = self._astra_db.collection(
+                collection_name=self.connector_config.collection_name,
+            )
+        return self._astra_db_collection  # type: ignore
+
+    @requires_dependencies(["astrapy"], extras="astra")
+    @SourceConnectionError.wrap  # type: ignore
+    def initialize(self):
+        _ = self.astra_db_collection
+
+    @requires_dependencies(["astrapy"], extras="astra")
+    def check_connection(self):
+        try:
+            _ = self.astra_db_collection
+        except Exception as e:
+            logger.error(f"Failed to validate connection {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+    @requires_dependencies(["astrapy"], extras="astra")
+    def get_ingest_docs(self):  # type: ignore
+        # Perform the find operation
+        astra_docs = list(self.astra_db_collection.paginated_find())
+
+        doc_list = []
+        for record in astra_docs:
+            doc = AstraIngestDoc(
+                connector_config=self.connector_config,
+                processor_config=self.processor_config,
+                read_config=self.read_config,
+                metadata=record,
+            )
+
+            doc.update_source_metadata()
+
+            doc_list.append(doc)
+
+        return doc_list
+
+
+@dataclass
+class AstraWriteConfig(WriteConfig):
+    embedding_dimension: int
+    requested_indexing_policy: t.Optional[t.Dict[str, t.Any]] = None
+    batch_size: int = 20
+
+
+@dataclass
+class AstraDestinationConnector(BaseDestinationConnector):
+    write_config: AstraWriteConfig
+    connector_config: SimpleAstraConfig
+    _astra_db: t.Optional["AstraDB"] = field(init=False, default=None)
+    _astra_db_collection: t.Optional["AstraDBCollection"] = field(init=False, default=None)
+
+    def to_dict(self, **kwargs):
+        """
+        The _astra_db_collection 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, "_astra_db_collection"):
+            setattr(self_cp, "_astra_db_collection", None)
+
+        return _asdict(self_cp, **kwargs)
+
+    @property
+    @requires_dependencies(["astrapy"], extras="astra")
+    def astra_db_collection(self) -> "AstraDBCollection":
+        if self._astra_db_collection is None:
+            from astrapy.db import AstraDB
+
+            collection_name = self.connector_config.collection_name
+            embedding_dimension = self.write_config.embedding_dimension
+
+            # If the user has requested an indexing policy, pass it to the AstraDB
+            requested_indexing_policy = self.write_config.requested_indexing_policy
+            options = {"indexing": requested_indexing_policy} if requested_indexing_policy else None
+
+            # caller_name/version for AstraDB tracking
+            self._astra_db = AstraDB(
+                api_endpoint=self.connector_config.access_config.api_endpoint,
+                token=self.connector_config.access_config.token,
+                namespace=self.connector_config.namespace,
+                caller_name=integration_name,
+                caller_version=integration_version,
+            )
+
+            # Create and connect to the newly created collection
+            self._astra_db_collection = self._astra_db.create_collection(
+                collection_name=collection_name,
+                dimension=embedding_dimension,
+                options=options,
+            )
+        return self._astra_db_collection
+
+    @requires_dependencies(["astrapy"], extras="astra")
+    @DestinationConnectionError.wrap
+    def initialize(self):
+        _ = self.astra_db_collection
+
+    @requires_dependencies(["astrapy"], extras="astra")
+    def check_connection(self):
+        try:
+            _ = self.astra_db_collection
+        except Exception as e:
+            logger.error(f"Failed to validate connection {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    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 Astra.")
+
+        astra_batch_size = self.write_config.batch_size
+
+        for batch in batch_generator(elements_dict, astra_batch_size):
+            self._astra_db_collection.insert_many(batch)
+
+    def normalize_dict(self, element_dict: dict) -> dict:
+        return {
+            "$vector": element_dict.pop("embeddings", None),
+            "content": element_dict.pop("text", None),
+            "metadata": element_dict,
+        }

unstructured_ingest/connector/azure_cognitive_search.py

@@ -0,0 +1,144 @@
+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, WriteError
+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 azure.search.documents import SearchClient
+
+
+@dataclass
+class AzureCognitiveSearchAccessConfig(AccessConfig):
+    key: str = enhanced_field(sensitive=True)
+
+
+@dataclass
+class SimpleAzureCognitiveSearchStorageConfig(BaseConnectorConfig):
+    endpoint: str
+    access_config: AzureCognitiveSearchAccessConfig
+
+
+@dataclass
+class AzureCognitiveSearchWriteConfig(WriteConfig):
+    index: str
+
+
+@dataclass
+class AzureCognitiveSearchDestinationConnector(BaseDestinationConnector):
+    write_config: AzureCognitiveSearchWriteConfig
+    connector_config: SimpleAzureCognitiveSearchStorageConfig
+    _client: t.Optional["SearchClient"] = field(init=False, default=None)
+
+    @requires_dependencies(["azure.search"], extras="azure-cognitive-search")
+    def generate_client(self) -> "SearchClient":
+        from azure.core.credentials import AzureKeyCredential
+        from azure.search.documents import SearchClient
+
+        # Create a client
+        credential = AzureKeyCredential(self.connector_config.access_config.key)
+        return SearchClient(
+            endpoint=self.connector_config.endpoint,
+            index_name=self.write_config.index,
+            credential=credential,
+        )
+
+    @property
+    def client(self) -> "SearchClient":
+        if self._client is None:
+            self._client = self.generate_client()
+        return self._client
+
+    def check_connection(self):
+        try:
+            self.client.get_document_count()
+        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 conform_dict(self, data: dict) -> None:
+        """
+        updates the dictionary that is from each Element being converted into a dict/json
+        into a dictionary that conforms to the schema expected by the
+        Azure Cognitive Search index
+        """
+        from dateutil import parser  # type: ignore
+
+        data["id"] = str(uuid.uuid4())
+
+        if points := data.get("metadata", {}).get("coordinates", {}).get("points"):
+            data["metadata"]["coordinates"]["points"] = json.dumps(points)
+        if version := data.get("metadata", {}).get("data_source", {}).get("version"):
+            data["metadata"]["data_source"]["version"] = str(version)
+        if record_locator := data.get("metadata", {}).get("data_source", {}).get("record_locator"):
+            data["metadata"]["data_source"]["record_locator"] = json.dumps(record_locator)
+        if permissions_data := (
+            data.get("metadata", {}).get("data_source", {}).get("permissions_data")
+        ):
+            data["metadata"]["data_source"]["permissions_data"] = json.dumps(permissions_data)
+        if links := data.get("metadata", {}).get("links"):
+            data["metadata"]["links"] = [json.dumps(link) for link in links]
+        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",
+            )
+        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 regex_metadata := data.get("metadata", {}).get("regex_metadata"):
+            data["metadata"]["regex_metadata"] = json.dumps(regex_metadata)
+        if page_number := data.get("metadata", {}).get("page_number"):
+            data["metadata"]["page_number"] = str(page_number)
+
+    @requires_dependencies(["azure"], extras="azure-cognitive-search")
+    def write_dict(self, *args, elements_dict: t.List[t.Dict[str, t.Any]], **kwargs) -> None:
+        import azure.core.exceptions
+
+        logger.info(
+            f"writing {len(elements_dict)} documents to destination "
+            f"index at {self.write_config.index}",
+        )
+        try:
+            results = self.client.upload_documents(documents=elements_dict)
+
+        except azure.core.exceptions.HttpResponseError as http_error:
+            raise WriteError(f"http error: {http_error}") from http_error
+        errors = []
+        success = []
+        for result in results:
+            if result.succeeded:
+                success.append(result)
+            else:
+                errors.append(result)
+        logger.debug(f"results: {len(success)} successes, {len(errors)} failures")
+        if errors:
+            raise WriteError(
+                ", ".join(
+                    [
+                        f"{error.key}: [{error.status_code}] {error.error_message}"
+                        for error in errors
+                    ],
+                ),
+            )

unstructured_ingest/connector/biomed.py

@@ -0,0 +1,313 @@
+import os
+import typing as t
+import urllib.request
+from dataclasses import dataclass
+from ftplib import FTP, error_perm
+from pathlib import Path
+
+import requests
+from requests.adapters import HTTPAdapter
+
+from unstructured_ingest.error import SourceConnectionError, SourceConnectionNetworkError
+from unstructured_ingest.interfaces import (
+    BaseConnectorConfig,
+    BaseSingleIngestDoc,
+    BaseSourceConnector,
+    IngestDocCleanupMixin,
+    SourceConnectorCleanupMixin,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.data_prep import (
+    validate_date_args,
+)
+
+DOMAIN = "ftp.ncbi.nlm.nih.gov"
+FTP_DOMAIN = f"ftp://{DOMAIN}"
+PMC_DIR = "pub/pmc"
+PDF_DIR = "oa_pdf"
+
+
+@dataclass
+class BiomedFileMeta:
+    ftp_path: str
+    download_filepath: str
+    output_filepath: str
+
+
+@dataclass
+class SimpleBiomedConfig(BaseConnectorConfig):
+    """Connector config where path is the FTP directory path and
+    id_, from_, until, format are API parameters."""
+
+    path: t.Optional[str] = None
+    # OA Web Service API Options
+    api_id: t.Optional[str] = None
+    api_from: t.Optional[str] = None
+    api_until: t.Optional[str] = None
+    max_request_time: int = 45
+
+    def validate_api_inputs(self):
+        valid = False
+
+        if self.api_from:
+            valid = validate_date_args(self.api_from)
+
+        if self.api_until:
+            valid = validate_date_args(self.api_until)
+
+        return valid
+
+    def __post_init__(self):
+        self.is_file = False
+        self.is_dir = False
+        self.is_api = False
+
+        if not self.path:
+            is_valid = self.validate_api_inputs()
+            if not is_valid:
+                raise ValueError(
+                    "Path argument or at least one of the "
+                    "OA Web Service arguments MUST be provided.",
+                )
+
+            self.is_api = True
+        else:
+            self.path = self.path.strip("/")
+            is_valid = self.path.lower().startswith(PDF_DIR)
+
+            if not is_valid:
+                raise ValueError(f"Path MUST start with {PDF_DIR}")
+
+            ftp = FTP(DOMAIN)
+            ftp.login()
+
+            path = Path(PMC_DIR) / self.path
+            response = ""
+            try:
+                if path.suffix == ".pdf":
+                    response = ftp.cwd(str(path.parent))
+                    self.is_file = True
+                else:
+                    response = ftp.cwd(str(path))
+            except error_perm as exc:
+                if "no such file or directory" in exc.args[0].lower():
+                    raise ValueError(f"The path: {path} is not valid.")
+                elif "not a directory" in exc.args[0].lower():
+                    self.is_file = True
+                elif "command successful" in response:
+                    self.is_dir = True
+                else:
+                    raise ValueError(
+                        "Something went wrong when validating the path: {path}.",
+                    )
+
+
+@dataclass
+class BiomedIngestDoc(IngestDocCleanupMixin, BaseSingleIngestDoc):
+    connector_config: SimpleBiomedConfig
+    file_meta: BiomedFileMeta
+    registry_name: str = "biomed"
+
+    @property
+    def filename(self):
+        return Path(self.file_meta.download_filepath).resolve()  # type: ignore
+
+    @property
+    def _output_filename(self):
+        return Path(f"{self.file_meta.output_filepath}.json").resolve()
+
+    def cleanup_file(self):
+        if (
+            not self.read_config.preserve_downloads
+            and self.filename.is_file()
+            and not self.read_config.download_only
+        ):
+            logger.debug(f"Cleaning up {self}")
+            Path.unlink(self.filename)
+
+    @SourceConnectionError.wrap
+    @BaseSingleIngestDoc.skip_if_file_exists
+    def get_file(self):
+        download_path = self.file_meta.download_filepath  # type: ignore
+        dir_ = Path(os.path.dirname(download_path))  # type: ignore
+        if not dir_.is_dir():
+            logger.debug(f"Creating directory: {dir_}")
+
+            if dir_:
+                dir_.mkdir(parents=True, exist_ok=True)
+        self._retrieve()
+        logger.debug(f"File downloaded: {self.file_meta.download_filepath}")
+
+    @SourceConnectionNetworkError.wrap
+    def _retrieve(self):
+        urllib.request.urlretrieve(
+            self.file_meta.ftp_path,  # type: ignore
+            self.file_meta.download_filepath,
+        )
+
+
+class BiomedSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
+    """Objects of this class support fetching documents from Biomedical literature FTP directory"""
+
+    connector_config: SimpleBiomedConfig
+
+    def get_base_endpoints_url(self) -> str:
+        endpoint_url = "https://www.ncbi.nlm.nih.gov/pmc/utils/oa/oa.fcgi?format=pdf"
+
+        if self.connector_config.api_id:
+            endpoint_url += f"&id={self.connector_config.api_id}"
+
+        if self.connector_config.api_from:
+            endpoint_url += f"&from={self.connector_config.api_from}"
+
+        if self.connector_config.api_until:
+            endpoint_url += f"&until={self.connector_config.api_until}"
+
+        return endpoint_url
+
+    def _list_objects_api(self) -> t.List[BiomedFileMeta]:
+        from bs4 import BeautifulSoup
+
+        def urls_to_metadata(urls):
+            files = []
+            for url in urls:
+                parts = url.split(PDF_DIR)
+                if len(parts) > 1:
+                    local_path = parts[1].strip("/")
+                    files.append(
+                        BiomedFileMeta(
+                            ftp_path=url,
+                            download_filepath=(Path(self.read_config.download_dir) / local_path)
+                            .resolve()
+                            .as_posix(),
+                            output_filepath=(Path(self.processor_config.output_dir) / local_path)
+                            .resolve()
+                            .as_posix(),
+                        ),
+                    )
+
+            return files
+
+        files: t.List[BiomedFileMeta] = []
+
+        endpoint_url = self.get_base_endpoints_url()
+
+        while endpoint_url:
+            session = requests.Session()
+            adapter = HTTPAdapter()
+            session.mount("http://", adapter)
+            session.mount("https://", adapter)
+            response = self._get_request(session=session, endpoint_url=endpoint_url)
+            soup = BeautifulSoup(response.content, features="lxml")
+            urls = [link["href"] for link in soup.find_all("link")]
+
+            if not urls:
+                return files
+
+            endpoint_url = urls[-1] if "resumptiontoken" in urls[-1].lower() else None
+            if endpoint_url:
+                urls = urls[:-1]
+
+            files.extend(urls_to_metadata(urls))
+
+        return files
+
+    @SourceConnectionNetworkError.wrap
+    def _get_request(self, session: requests.Session, endpoint_url: str) -> requests.Response:
+        return session.get(endpoint_url, timeout=self.connector_config.max_request_time)
+
+    def _list_objects(self) -> t.List[BiomedFileMeta]:
+        files = []
+
+        # Conform to mypy, null check performed elsewhere.
+        # Wouldn't be in this method unless self.config.path exists
+        path: str = self.connector_config.path if self.connector_config.path else ""
+
+        def traverse(path, download_dir, output_dir):
+            full_path = Path(PMC_DIR) / path
+            logger.debug(f"Traversing directory: {full_path}")
+
+            ftp = FTP(DOMAIN)
+            ftp.login()
+
+            try:
+                response = ftp.cwd(str(full_path))
+            except error_perm:
+                raise ValueError(f"{full_path} is not a valid directory.")
+
+            if "command successful" in response.lower():
+                sub_paths = [path / p for p in ftp.nlst()]
+
+                if not sub_paths:
+                    return
+
+                ext = Path(sub_paths[0]).suffix
+                if ext:
+                    for sub_path in sub_paths:
+                        ftp_path = f"{FTP_DOMAIN}/{PMC_DIR}/{sub_path}"
+                        local_path = "/".join(str(sub_path).split("/")[1:])
+                        files.append(
+                            BiomedFileMeta(
+                                ftp_path=ftp_path,
+                                download_filepath=(Path(self.read_config.download_dir) / local_path)
+                                .resolve()
+                                .as_posix(),
+                                output_filepath=(
+                                    Path(self.processor_config.output_dir) / local_path
+                                )
+                                .resolve()
+                                .as_posix(),
+                            ),
+                        )
+
+                else:
+                    for sub_path in sub_paths:
+                        traverse(sub_path, download_dir, output_dir)
+
+            else:
+                raise ValueError(f"{full_path} is not a valid directory.")
+
+        ftp_path = f"{FTP_DOMAIN}/{PMC_DIR}/{self.connector_config.path}"
+        if self.connector_config.is_file:
+            local_path = "/".join(path.split("/")[1:])
+            return [
+                BiomedFileMeta(
+                    ftp_path=ftp_path,
+                    download_filepath=(Path(self.read_config.download_dir) / local_path)
+                    .resolve()
+                    .as_posix(),
+                    output_filepath=(Path(self.processor_config.output_dir) / local_path)
+                    .resolve()
+                    .as_posix(),
+                ),
+            ]
+        else:
+            traverse(
+                Path(path),
+                Path(self.read_config.download_dir),
+                Path(self.processor_config.output_dir),
+            )
+
+        return files
+
+    def initialize(self):
+        pass
+
+    def check_connection(self):
+        resp = requests.head(self.get_base_endpoints_url())
+        try:
+            resp.raise_for_status()
+        except requests.HTTPError as http_error:
+            raise SourceConnectionError(f"failed to validate connection: {http_error}")
+
+    def get_ingest_docs(self):
+        files = self._list_objects_api() if self.connector_config.is_api else self._list_objects()
+        return [
+            BiomedIngestDoc(
+                processor_config=self.processor_config,
+                connector_config=self.connector_config,
+                read_config=self.read_config,
+                file_meta=file,
+            )
+            for file in files
+        ]