Harrison/diffbot (#2984)

Co-authored-by: Manuel Saelices <msaelices@gmail.com>

langchain/document_loaders/__init__.py

@@ -16,6 +16,7 @@ from langchain.document_loaders.college_confidential import CollegeConfidentialL
 from langchain.document_loaders.conllu import CoNLLULoader
 from langchain.document_loaders.csv_loader import CSVLoader
 from langchain.document_loaders.dataframe import DataFrameLoader
+from langchain.document_loaders.diffbot import DiffbotLoader
 from langchain.document_loaders.directory import DirectoryLoader
 from langchain.document_loaders.duckdb_loader import DuckDBLoader
 from langchain.document_loaders.email import (
@@ -141,6 +142,7 @@ __all__ = [
     "SitemapLoader",
     "DuckDBLoader",
     "BigQueryLoader",
+    "DiffbotLoader",
     "BiliBiliLoader",
     "SlackDirectoryLoader",
     "GitLoader",

langchain/document_loaders/diffbot.py

@@ -0,0 +1,55 @@
+"""Loader that uses Diffbot to load webpages in text format."""
+import logging
+from typing import Any, List
+
+import requests
+
+from langchain.docstore.document import Document
+from langchain.document_loaders.base import BaseLoader
+
+logger = logging.getLogger(__file__)
+
+
+class DiffbotLoader(BaseLoader):
+    """Loader that loads Diffbot file json."""
+
+    def __init__(
+        self, api_token: str, urls: List[str], continue_on_failure: bool = True
+    ):
+        """Initialize with API token, ids, and key."""
+        self.api_token = api_token
+        self.urls = urls
+        self.continue_on_failure = continue_on_failure
+
+    def _diffbot_api_url(self, diffbot_api: str) -> str:
+        return f"https://api.diffbot.com/v3/{diffbot_api}"
+
+    def _get_diffbot_data(self, url: str) -> Any:
+        """Get Diffbot file from Diffbot REST API."""
+        # TODO: Add support for other Diffbot APIs
+        diffbot_url = self._diffbot_api_url("article")
+        params = {
+            "token": self.api_token,
+            "url": url,
+        }
+        response = requests.get(diffbot_url, params=params, timeout=10)
+
+        # TODO: handle non-ok errors
+        return response.json() if response.ok else {}
+
+    def load(self) -> List[Document]:
+        """Extract text from Diffbot on all the URLs and return Document instances"""
+        docs: List[Document] = list()
+
+        for url in self.urls:
+            try:
+                data = self._get_diffbot_data(url)
+                text = data["objects"][0]["text"] if "objects" in data else ""
+                metadata = {"source": url}
+                docs.append(Document(page_content=text, metadata=metadata))
+            except Exception as e:
+                if self.continue_on_failure:
+                    logger.error(f"Error fetching or processing {url}, exception: {e}")
+                else:
+                    raise e
+        return docs