unstructured-ingest 0.3.13__py3-none-any.whl

unstructured_ingest/connector/biomed.py

@@ -0,0 +1,320 @@
+import os
+import typing as t
+import urllib.request
+from dataclasses import dataclass
+from ftplib import FTP, error_perm
+from pathlib import Path
+
+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,
+)
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if t.TYPE_CHECKING:
+    from requests import Response, Session
+
+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
+
+    @requires_dependencies(["requests"], extras="biomed")
+    def _list_objects_api(self) -> t.List[BiomedFileMeta]:
+        from bs4 import BeautifulSoup
+        from requests import Session
+        from requests.adapters import HTTPAdapter
+
+        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 = 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: "Session", endpoint_url: str) -> "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
+
+    @requires_dependencies(["requests"], extras="biomed")
+    def check_connection(self):
+        import requests
+
+        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
+        ]

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