codekingpro/portable-devtools

1import os2from typing import Any, Dict, Iterator, List, Optional, Union3 4from langchain_core.documents import Document5from langchain_core.utils import get_from_dict_or_env6 7from langchain_community.document_loaders import PyPDFLoader8from langchain_community.document_loaders.base import BaseLoader9 10 11class RSpaceLoader(BaseLoader):12    """Load content from RSpace notebooks, folders, documents or PDF Gallery files.13 14    Map RSpace document <-> Langchain Document in 1-1. PDFs are imported using PyPDF.15 16    Requirements are rspace_client (`pip install rspace_client`) and PyPDF if importing17     PDF docs (`pip install pypdf`).18 19    """20 21    def __init__(22        self, global_id: str, api_key: Optional[str] = None, url: Optional[str] = None23    ):24        """api_key: RSpace API key - can also be supplied as environment variable25        'RSPACE_API_KEY'26        url: str27        The URL of your RSpace instance - can also be supplied as environment28        variable 'RSPACE_URL'29        global_id: str30         The global ID of the resource to load,31        e.g. 'SD12344' (a single document); 'GL12345'(A PDF file in the gallery);32        'NB4567' (a notebook); 'FL12244' (a folder)33        """34        args: Dict[str, Optional[str]] = {35            "api_key": api_key,36            "url": url,37            "global_id": global_id,38        }39        verified_args: Dict[str, str] = RSpaceLoader.validate_environment(args)40        self.api_key = verified_args["api_key"]41        self.url = verified_args["url"]42        self.global_id: str = verified_args["global_id"]43 44    @classmethod45    def validate_environment(cls, values: Dict) -> Dict:46        """Validate that API key and URL exist in environment."""47        values["api_key"] = get_from_dict_or_env(values, "api_key", "RSPACE_API_KEY")48        values["url"] = get_from_dict_or_env(values, "url", "RSPACE_URL")49        if "global_id" not in values or values["global_id"] is None:50            raise ValueError(51                "No value supplied for global_id. Please supply an RSpace global ID"52            )53        return values54 55    def _create_rspace_client(self) -> Any:56        """Create a RSpace client."""57        try:58            from rspace_client.eln import eln, field_content59 60        except ImportError:61            raise ImportError("You must run `pip install rspace_client`")62 63        try:64            eln = eln.ELNClient(self.url, self.api_key)65            eln.get_status()66 67        except Exception:68            raise Exception(69                f"Unable to initialize client - is url {self.url} or api key  correct?"70            )71 72        return eln, field_content.FieldContent73 74    def _get_doc(self, cli: Any, field_content: Any, d_id: Union[str, int]) -> Document:75        content = ""76        doc = cli.get_document(d_id)77        content += f"<h2>{doc['name']}<h2/>"78        for f in doc["fields"]:79            content += f"{f['name']}\n"80            fc = field_content(f["content"])81            content += fc.get_text()82            content += "\n"83        return Document(84            metadata={"source": f"rspace: {doc['name']}-{doc['globalId']}"},85            page_content=content,86        )87 88    def _load_structured_doc(self) -> Iterator[Document]:89        cli, field_content = self._create_rspace_client()90        yield self._get_doc(cli, field_content, self.global_id)91 92    def _load_folder_tree(self) -> Iterator[Document]:93        cli, field_content = self._create_rspace_client()94        if self.global_id:95            docs_in_folder = cli.list_folder_tree(96                folder_id=self.global_id[2:], typesToInclude=["document"]97            )98        doc_ids: List[int] = [d["id"] for d in docs_in_folder["records"]]99        for doc_id in doc_ids:100            yield self._get_doc(cli, field_content, doc_id)101 102    def _load_pdf(self) -> Iterator[Document]:103        cli, field_content = self._create_rspace_client()104        file_info = cli.get_file_info(self.global_id)105        _, ext = os.path.splitext(file_info["name"])106        if ext.lower() == ".pdf":107            outfile = f"{self.global_id}.pdf"108            cli.download_file(self.global_id, outfile)109            pdf_loader = PyPDFLoader(outfile)110            for pdf in pdf_loader.lazy_load():111                pdf.metadata["rspace_src"] = self.global_id112                yield pdf113 114    def lazy_load(self) -> Iterator[Document]:115        if self.global_id and "GL" in self.global_id:116            for d in self._load_pdf():117                yield d118        elif self.global_id and "SD" in self.global_id:119            for d in self._load_structured_doc():120                yield d121        elif self.global_id and self.global_id[0:2] in ["FL", "NB"]:122            for d in self._load_folder_tree():123                yield d124        else:125            raise ValueError("Unknown global ID type")126