codekingpro/portable-devtools

1import json2import urllib.request3from typing import Any, Iterator4 5from langchain_core.documents import Document6 7from langchain_community.document_loaders.base import BaseLoader8 9 10class LarkSuiteDocLoader(BaseLoader):11    """Load from `LarkSuite` (`FeiShu`)."""12 13    def __init__(self, domain: str, access_token: str, document_id: str):14        """Initialize with domain, access_token (tenant / user), and document_id.15 16        Args:17            domain: The domain to load the LarkSuite.18            access_token: The access_token to use.19            document_id: The document_id to load.20        """21        self.domain = domain22        self.access_token = access_token23        self.document_id = document_id24 25    def _get_larksuite_api_json_data(self, api_url: str) -> Any:26        """Get LarkSuite (FeiShu) API response json data."""27        headers = {"Authorization": f"Bearer {self.access_token}"}28        request = urllib.request.Request(api_url, headers=headers)29        with urllib.request.urlopen(request) as response:30            json_data = json.loads(response.read().decode())31            return json_data32 33    def lazy_load(self) -> Iterator[Document]:34        """Lazy load LarkSuite (FeiShu) document."""35        api_url_prefix = f"{self.domain}/open-apis/docx/v1/documents"36        metadata_json = self._get_larksuite_api_json_data(37            f"{api_url_prefix}/{self.document_id}"38        )39        raw_content_json = self._get_larksuite_api_json_data(40            f"{api_url_prefix}/{self.document_id}/raw_content"41        )42        text = raw_content_json["data"]["content"]43        metadata = {44            "document_id": self.document_id,45            "revision_id": metadata_json["data"]["document"]["revision_id"],46            "title": metadata_json["data"]["document"]["title"],47        }48        yield Document(page_content=text, metadata=metadata)49 50 51class LarkSuiteWikiLoader(LarkSuiteDocLoader):52    """Load from `LarkSuite` (`FeiShu`) wiki."""53 54    def __init__(self, domain: str, access_token: str, wiki_id: str):55        """Initialize with domain, access_token (tenant / user), and wiki_id.56 57        Args:58            domain: The domain to load the LarkSuite.59            access_token: The access_token to use.60            wiki_id: The wiki_id to load.61        """62        self.domain = domain63        self.access_token = access_token64        self.wiki_id = wiki_id65        self.document_id = ""66 67    def lazy_load(self) -> Iterator[Document]:68        """Lazy load LarkSuite (FeiShu) wiki document."""69 70        # convert Feishu wiki id to document id71        if not self.document_id:72            wiki_url_prefix = f"{self.domain}/open-apis/wiki/v2/spaces/get_node"73            wiki_node_info_json = self._get_larksuite_api_json_data(74                f"{wiki_url_prefix}?token={self.wiki_id}"75            )76            self.document_id = wiki_node_info_json["data"]["node"]["obj_token"]77 78        yield from super().lazy_load()79