codekingpro/portable-devtools

1import logging2from typing import Any, List3 4import requests5from langchain_core.documents import Document6 7from langchain_community.document_loaders.base import BaseLoader8 9logger = logging.getLogger(__name__)10 11 12class DiffbotLoader(BaseLoader):13    """Load `Diffbot` json file."""14 15    def __init__(16        self, api_token: str, urls: List[str], continue_on_failure: bool = True17    ):18        """Initialize with API token, ids, and key.19 20        Args:21            api_token: Diffbot API token.22            urls: List of URLs to load.23            continue_on_failure: Whether to continue loading other URLs if one fails.24               Defaults to True.25        """26        self.api_token = api_token27        self.urls = urls28        self.continue_on_failure = continue_on_failure29 30    def _diffbot_api_url(self, diffbot_api: str) -> str:31        return f"https://api.diffbot.com/v3/{diffbot_api}"32 33    def _get_diffbot_data(self, url: str) -> Any:34        """Get Diffbot file from Diffbot REST API."""35        # TODO: Add support for other Diffbot APIs36        diffbot_url = self._diffbot_api_url("article")37        params = {38            "token": self.api_token,39            "url": url,40        }41        response = requests.get(diffbot_url, params=params, timeout=10)42 43        # TODO: handle non-ok errors44        return response.json() if response.ok else {}45 46    def load(self) -> List[Document]:47        """Extract text from Diffbot on all the URLs and return Documents"""48        docs: List[Document] = list()49 50        for url in self.urls:51            try:52                data = self._get_diffbot_data(url)53                text = data["objects"][0]["text"] if "objects" in data else ""54                metadata = {"source": url}55                docs.append(Document(page_content=text, metadata=metadata))56            except Exception as e:57                if self.continue_on_failure:58                    logger.error(f"Error fetching or processing {url}, exception: {e}")59                else:60                    raise e61        return docs62