codekingpro/portable-devtools

1"""Base class for all loaders that uses O365 Package"""2 3from __future__ import annotations4 5import logging6import mimetypes7import os8import re9import tempfile10import urllib11from abc import abstractmethod12from datetime import datetime13from pathlib import Path, PurePath14from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Sequence, Union15 16from pydantic import (17    BaseModel,18    Field,19    FilePath,20    PrivateAttr,21    SecretStr,22)23from pydantic_settings import BaseSettings, SettingsConfigDict24 25from langchain_community.document_loaders.base import BaseBlobParser, BaseLoader26from langchain_community.document_loaders.blob_loaders.file_system import (27    FileSystemBlobLoader,28)29from langchain_community.document_loaders.blob_loaders.schema import Blob30from langchain_community.document_loaders.parsers.generic import MimeTypeBasedParser31from langchain_community.document_loaders.parsers.registry import get_parser32 33if TYPE_CHECKING:34    from O365 import Account35    from O365.drive import Drive, Folder36 37logger = logging.getLogger(__name__)38 39CHUNK_SIZE = 1024 * 1024 * 540 41 42class _O365Settings(BaseSettings):43    client_id: str = Field(..., alias="O365_CLIENT_ID")44    client_secret: SecretStr = Field(..., alias="O365_CLIENT_SECRET")45 46    model_config = SettingsConfigDict(47        case_sensitive=False, env_file=".env", env_prefix="", extra="ignore"48    )49 50 51class _O365TokenStorage(BaseSettings):52    token_path: FilePath = Path.home() / ".credentials" / "o365_token.txt"53 54 55def fetch_mime_types(file_types: Sequence[str]) -> Dict[str, str]:56    """Fetch the mime types for the specified file types."""57    mime_types_mapping = {}58    for ext in file_types:59        mime_type, _ = mimetypes.guess_type(f"file.{ext}")60        if mime_type:61            mime_types_mapping[ext] = mime_type62        else:63            raise ValueError(f"Unknown mimetype of extension {ext}")64    return mime_types_mapping65 66 67def fetch_extensions(mime_types: Sequence[str]) -> Dict[str, str]:68    """Fetch the mime types for the specified file types."""69    mime_types_mapping = {}70    for mime_type in mime_types:71        ext = mimetypes.guess_extension(mime_type)72        if ext:73            mime_types_mapping[ext[1:]] = mime_type  # ignore leading `.`74        else:75            raise ValueError(f"Unknown mimetype {mime_type}")76    return mime_types_mapping77 78 79class O365BaseLoader(BaseLoader, BaseModel):80    """Base class for all loaders that uses O365 Package"""81 82    settings: _O365Settings = Field(default_factory=_O365Settings)  # type: ignore[arg-type]83    """Settings for the Office365 API client."""84    auth_with_token: bool = False85    """Whether to authenticate with a token or not. Defaults to False."""86    chunk_size: Union[int, str] = CHUNK_SIZE87    """Number of bytes to retrieve from each api call to the server. int or 'auto'."""88    recursive: bool = False89    """Should the loader recursively load subfolders?"""90    modified_since: Optional[datetime] = None91    """Only fetch documents modified since given datetime. The datetime object92    must be timezone aware."""93    handlers: Optional[Dict[str, Any]] = {}94    """95    Provide custom handlers for MimeTypeBasedParser.96 97    Pass a dictionary mapping either file extensions (like "doc", "pdf", etc.) 98    or MIME types (like "application/pdf", "text/plain", etc.) to parsers. 99    Note that you must use either file extensions or MIME types exclusively and 100    cannot mix them.101 102    Do not include the leading dot for file extensions.103    104    Example using file extensions:105    ```python106        handlers = {107            "doc": MsWordParser(),108            "pdf": PDFMinerParser(),109            "txt": TextParser()110        }111    ```112    113    Example using MIME types:114    ```python115        handlers = {116            "application/msword": MsWordParser(),117            "application/pdf": PDFMinerParser(),118            "text/plain": TextParser()119        }120    ```121    """122 123    _blob_parser: BaseBlobParser = PrivateAttr()124    _file_types: Sequence[str] = PrivateAttr()125    _mime_types: Dict[str, str] = PrivateAttr()126 127    def __init__(self, **kwargs: Any) -> None:128        super().__init__(**kwargs)129        if self.handlers:130            handler_keys = list(self.handlers.keys())131            try:132                # assume handlers.keys() are file extensions133                self._mime_types = fetch_mime_types(handler_keys)134                self._file_types = list(set(handler_keys))135                mime_handlers = {136                    self._mime_types[extension]: handler137                    for extension, handler in self.handlers.items()138                }139            except ValueError:140                try:141                    # assume handlers.keys() are mime types142                    self._mime_types = fetch_extensions(handler_keys)143                    self._file_types = list(set(self._mime_types.keys()))144                    mime_handlers = self.handlers145                except ValueError:146                    raise ValueError(147                        "`handlers` keys must be either file extensions or mimetypes.\n"148                        f"{handler_keys} could not be interpreted as either.\n"149                        "File extensions and mimetypes cannot mix. "150                        "Use either one or the other"151                    )152 153            self._blob_parser = MimeTypeBasedParser(154                handlers=mime_handlers, fallback_parser=None155            )156        else:157            self._blob_parser = get_parser("default")158            if not isinstance(self._blob_parser, MimeTypeBasedParser):159                raise TypeError(160                    'get_parser("default) was supposed to return MimeTypeBasedParser.'161                    f"It returned {type(self._blob_parser)}"162                )163            self._mime_types = fetch_extensions(list(self._blob_parser.handlers.keys()))164 165    @property166    def _fetch_mime_types(self) -> Dict[str, str]:167        """Return a dict of supported file types to corresponding mime types."""168        return self._mime_types169 170    @property171    @abstractmethod172    def _scopes(self) -> List[str]:173        """Return required scopes."""174 175    def _load_from_folder(self, folder: Folder) -> Iterable[Blob]:176        """Lazily load all files from a specified folder of the configured MIME type.177 178        Args:179            folder: The Folder instance from which the files are to be loaded. This180                Folder instance should represent a directory in a file system where the181                files are stored.182 183        Yields:184            An iterator that yields Blob instances, which are binary representations of185                the files loaded from the folder.186        """187        file_mime_types = self._fetch_mime_types188        items = folder.get_items()189        metadata_dict: Dict[str, Dict[str, Any]] = {}190        with tempfile.TemporaryDirectory() as temp_dir:191            os.makedirs(os.path.dirname(temp_dir), exist_ok=True)192            for file in items:193                if file.is_file:194                    if file.mime_type in list(file_mime_types.values()):195                        if (not self.modified_since) or (196                            file.modified > self.modified_since197                        ):198                            source = file.web_url199                            if re.search(200                                r"Doc.aspx\?sourcedoc=.*file=([^&]+)", file.web_url201                            ):202                                source = (203                                    file._parent.web_url204                                    + "/"205                                    + urllib.parse.quote(file.name)206                                )207                            file.download(to_path=temp_dir, chunk_size=self.chunk_size)208                            metadata_dict[file.name] = {209                                "source": source,210                                "mime_type": file.mime_type,211                                "created": str(file.created),212                                "modified": str(file.modified),213                                "created_by": str(file.created_by),214                                "modified_by": str(file.modified_by),215                                "description": file.description,216                                "id": str(file.object_id),217                            }218 219            loader = FileSystemBlobLoader(path=temp_dir)220            for blob in loader.yield_blobs():221                if not isinstance(blob.path, PurePath):222                    raise NotImplementedError("Expected blob path to be a PurePath")223                if blob.path:224                    file_metadata_ = metadata_dict.get(str(blob.path.name), {})225                    blob.metadata.update(file_metadata_)226                yield blob227        if self.recursive:228            for subfolder in folder.get_child_folders():229                yield from self._load_from_folder(subfolder)230 231    def _load_from_object_ids(232        self, drive: Drive, object_ids: List[str]233    ) -> Iterable[Blob]:234        """Lazily load files specified by their object_ids from a drive.235 236        Load files into the system as binary large objects (Blobs) and return Iterable.237 238        Args:239            drive: The Drive instance from which the files are to be loaded. This Drive240                instance should represent a cloud storage service or similar storage241                system where the files are stored.242            object_ids: A list of object_id strings. Each object_id represents a unique243                identifier for a file in the drive.244 245        Yields:246            An iterator that yields Blob instances, which are binary representations of247            the files loaded from the drive using the specified object_ids.248        """249        file_mime_types = self._fetch_mime_types250        metadata_dict: Dict[str, Dict[str, Any]] = {}251        with tempfile.TemporaryDirectory() as temp_dir:252            for object_id in object_ids:253                file = drive.get_item(object_id)254                if not file:255                    logging.warning(256                        "There isn't a file with"257                        f"object_id {object_id} in drive {drive}."258                    )259                    continue260                if file.is_file:261                    if file.mime_type in list(file_mime_types.values()):262                        source = file.web_url263                        if re.search(264                            r"Doc.aspx\?sourcedoc=.*file=([^&]+)", file.web_url265                        ):266                            source = (267                                file._parent.web_url268                                + "/"269                                + urllib.parse.quote(file.name)270                            )271                        file.download(to_path=temp_dir, chunk_size=self.chunk_size)272                        metadata_dict[file.name] = {273                            "source": source,274                            "mime_type": file.mime_type,275                            "created": file.created,276                            "modified": file.modified,277                            "created_by": str(file.created_by),278                            "modified_by": str(file.modified_by),279                            "description": file.description,280                            "id": str(file.object_id),281                        }282 283            loader = FileSystemBlobLoader(path=temp_dir)284            for blob in loader.yield_blobs():285                if not isinstance(blob.path, PurePath):286                    raise NotImplementedError("Expected blob path to be a PurePath")287                if blob.path:288                    file_metadata_ = metadata_dict.get(str(blob.path.name), {})289                    blob.metadata.update(file_metadata_)290                yield blob291 292    def _auth(self) -> Account:293        """Authenticates the OneDrive API client294 295        Returns:296            The authenticated Account object.297        """298        try:299            from O365 import Account, FileSystemTokenBackend300        except ImportError:301            raise ImportError(302                "O365 package not found, please install it with `pip install o365`"303            )304        if self.auth_with_token:305            token_storage = _O365TokenStorage()306            token_path = token_storage.token_path307            token_backend = FileSystemTokenBackend(308                token_path=token_path.parent, token_filename=token_path.name309            )310            account = Account(311                credentials=(312                    self.settings.client_id,313                    self.settings.client_secret.get_secret_value(),314                ),315                scopes=self._scopes,316                token_backend=token_backend,317                **{"raise_http_errors": False},318            )319        else:320            token_backend = FileSystemTokenBackend(321                token_path=Path.home() / ".credentials"322            )323            account = Account(324                credentials=(325                    self.settings.client_id,326                    self.settings.client_secret.get_secret_value(),327                ),328                scopes=self._scopes,329                token_backend=token_backend,330                **{"raise_http_errors": False},331            )332            # make the auth333            account.authenticate()334        return account335