GitbookLoader now inherits from BaseLoader and uses WebLoader thread-safely, CVE support added, reverted pyproject.toml and uv.lock to master

Author: andrasfe

libs/community/langchain_community/document_loaders/gitbook.py

@@ -1,14 +1,15 @@
 import warnings
-from typing import Any, AsyncIterator, Iterator, List, Optional, Set
-from urllib.parse import urljoin, urlparse
+from typing import Any, AsyncIterator, Iterator, List, Optional, Set, Union
+from urllib.parse import urlparse
 
 from bs4 import BeautifulSoup
 from langchain_core.documents import Document
 
+from langchain_community.document_loaders.base import BaseLoader
 from langchain_community.document_loaders.web_base import WebBaseLoader
 
 
-class GitbookLoader(WebBaseLoader):
+class GitbookLoader(BaseLoader):
     """Load `GitBook` data.
 
     1. load from either a single page, or
@@ -25,6 +26,7 @@ def __init__(
         show_progress: bool = True,
         *,
         sitemap_url: Optional[str] = None,
+        allowed_domains: Optional[Set[str]] = None,
     ):
         """Initialize with web page and whether to load all paths.
 
@@ -44,25 +46,86 @@ def __init__(
             show_progress: whether to show a progress bar while loading. Default: True
             sitemap_url: Custom sitemap URL to use when load_all_paths is True.
                 Defaults to "{base_url}/sitemap.xml".
+            allowed_domains: Optional set of allowed domains to fetch from.
+                If provided, only URLs from these domains will be processed.
+                Helps prevent SSRF vulnerabilities in server environments.
+                Defaults to None (all domains allowed).
         """
         self.base_url = base_url or web_page
         if self.base_url.endswith("/"):
             self.base_url = self.base_url[:-1]
 
-        if load_all_paths:
-            # set web_path to the sitemap if we want to crawl all paths
-            if sitemap_url:
-                web_page = sitemap_url
-            else:
-                web_page = f"{self.base_url}/sitemap.xml"
-
-        super().__init__(
-            web_paths=(web_page,),
-            continue_on_failure=continue_on_failure,
-            show_progress=show_progress,
-        )
+        self.web_page = web_page
         self.load_all_paths = load_all_paths
         self.content_selector = content_selector
+        self.continue_on_failure = continue_on_failure
+        self.show_progress = show_progress
+        self.allowed_domains = allowed_domains
+
+        # If allowed_domains is not specified, extract domain from web_page as default
+        if self.allowed_domains is None:
+            initial_domain = urlparse(web_page).netloc
+            if initial_domain:
+                self.allowed_domains = {initial_domain}
+
+        # Determine the starting URL (either a sitemap or a direct page)
+        if load_all_paths:
+            self.start_url = sitemap_url or f"{self.base_url}/sitemap.xml"
+        else:
+            self.start_url = web_page
+
+        # Validate the start_url is allowed
+        if not self._is_url_allowed(self.start_url):
+            raise ValueError(
+                f"Domain in {self.start_url} is not in the allowed domains list: "
+                f"{self.allowed_domains}"
+            )
+
+    def _is_url_allowed(self, url: str) -> bool:
+        """Check if a URL's domain is allowed for processing.
+
+        Args:
+            url: The URL to check
+
+        Returns:
+            bool: True if the domain is allowed or if no allowed_domains set is defined
+        """
+        if self.allowed_domains is None:
+            return True
+
+        netloc = urlparse(url).netloc
+        return netloc in self.allowed_domains
+
+    def _safe_add_url(
+        self, url_list: List[str], url: str, url_type: str = "URL"
+    ) -> bool:
+        """Safely add a URL to a list if it's from an allowed domain.
+
+        Args:
+            url_list: The list to add the URL to
+            url: The URL to add
+            url_type: Type of URL for warning message (e.g., "sitemap", "content")
+
+        Returns:
+            bool: True if URL was added, False if skipped
+        """
+        if self._is_url_allowed(url):
+            url_list.append(url)
+            return True
+        else:
+            warnings.warn(f"Skipping disallowed {url_type} URL: {url}")
+            return False
+
+    def _create_web_loader(self, url_or_urls: Union[str, List[str]]) -> WebBaseLoader:
+        """Create a new WebBaseLoader instance for the given URL(s).
+
+        This ensures each operation gets its own isolated WebBaseLoader.
+        """
+        return WebBaseLoader(
+            web_path=url_or_urls,
+            continue_on_failure=self.continue_on_failure,
+            show_progress=self.show_progress,
+        )
 
     def _is_sitemap_index(self, soup: BeautifulSoup) -> bool:
         """Check if the soup contains a sitemap index."""
@@ -71,19 +134,30 @@ def _is_sitemap_index(self, soup: BeautifulSoup) -> bool:
     def _extract_sitemap_urls(self, soup: BeautifulSoup) -> List[str]:
         """Extract sitemap URLs from a sitemap index."""
         sitemap_tags = soup.find_all("sitemap")
-        urls = []
+        urls: List[str] = []
         for sitemap in sitemap_tags:
             loc = sitemap.find("loc")
             if loc and loc.text:
-                urls.append(loc.text)
+                self._safe_add_url(urls, loc.text, "sitemap")
         return urls
 
     def _process_sitemap(
-        self, soup: BeautifulSoup, processed_urls: Optional[Set[str]] = None
+        self,
+        soup: BeautifulSoup,
+        processed_urls: Set[str],
+        web_loader: Optional[WebBaseLoader] = None,
     ) -> List[str]:
-        """Process a sitemap, handling both direct content URLs and sitemap indexes."""
-        if processed_urls is None:
-            processed_urls = set()
+        """Process a sitemap, handling both direct content URLs and sitemap indexes.
+
+        Args:
+            soup: The BeautifulSoup object of the sitemap
+            processed_urls: Set of already processed URLs to avoid cycles
+            web_loader: WebBaseLoader instance to reuse for all requests,
+                created if None
+        """
+        # Create a loader if not provided
+        if web_loader is None:
+            web_loader = self._create_web_loader(self.start_url)
 
         # If it's a sitemap index, recursively process each sitemap URL
         if self._is_sitemap_index(soup):
@@ -99,13 +173,20 @@ def _process_sitemap(
 
                 processed_urls.add(sitemap_url)
                 try:
-                    # We need to temporarily set the web_paths to the sitemap URL
-                    original_web_paths = self.web_paths
-                    self.web_paths = [sitemap_url]
-                    sitemap_soup = self.scrape(parser="xml")
+                    # Temporarily override the web_path of the loader
+                    original_web_paths = web_loader.web_paths
+                    web_loader.web_paths = [sitemap_url]
+
+                    # Reuse the same loader for the next sitemap
+                    sitemap_soup = web_loader.scrape(parser="xml")
+
                     # Restore original web_paths
-                    self.web_paths = original_web_paths
-                    content_urls = self._process_sitemap(sitemap_soup, processed_urls)
+                    web_loader.web_paths = original_web_paths
+
+                    # Recursive call with the same loader
+                    content_urls = self._process_sitemap(
+                        sitemap_soup, processed_urls, web_loader
+                    )
                     all_content_urls.extend(content_urls)
                 except Exception as e:
                     if self.continue_on_failure:
@@ -122,28 +203,49 @@ async def _aprocess_sitemap(
         self,
         soup: BeautifulSoup,
         base_url: str,
-        processed_urls: Optional[Set[str]] = None,
+        processed_urls: Set[str],
+        web_loader: Optional[WebBaseLoader] = None,
     ) -> List[str]:
-        """Async version of _process_sitemap."""
-        if processed_urls is None:
-            processed_urls = set()
+        """Async version of _process_sitemap.
+
+        Args:
+            soup: The BeautifulSoup object of the sitemap
+            base_url: The base URL for relative paths
+            processed_urls: Set of already processed URLs to avoid cycles
+            web_loader: WebBaseLoader instance to reuse for all requests,
+                created if None
+        """
+        # Create a loader if not provided
+        if web_loader is None:
+            web_loader = self._create_web_loader(self.start_url)
 
         # If it's a sitemap index, recursively process each sitemap URL
         if self._is_sitemap_index(soup):
             sitemap_urls = self._extract_sitemap_urls(soup)
             all_content_urls = []
 
-            # Use base class's ascrape_all for efficient parallel fetching
-            soups = await self.ascrape_all(
-                [url for url in sitemap_urls if url not in processed_urls], parser="xml"
-            )
-            for sitemap_url, sitemap_soup in zip(
-                [url for url in sitemap_urls if url not in processed_urls], soups
-            ):
+            # Filter out already processed URLs
+            new_urls = [url for url in sitemap_urls if url not in processed_urls]
+
+            if not new_urls:
+                return []
+
+            # Update the web_paths of the loader to fetch all sitemaps at once
+            original_web_paths = web_loader.web_paths
+            web_loader.web_paths = new_urls
+
+            # Use the same WebBaseLoader's ascrape_all for efficient parallel fetching
+            soups = await web_loader.ascrape_all(new_urls, parser="xml")
+
+            # Restore original web_paths
+            web_loader.web_paths = original_web_paths
+
+            for sitemap_url, sitemap_soup in zip(new_urls, soups):
                 processed_urls.add(sitemap_url)
                 try:
+                    # Recursive call with the same loader
                     content_urls = await self._aprocess_sitemap(
-                        sitemap_soup, base_url, processed_urls
+                        sitemap_soup, base_url, processed_urls, web_loader
                     )
                     all_content_urls.extend(content_urls)
                 except Exception as e:
@@ -159,53 +261,84 @@ async def _aprocess_sitemap(
 
     def lazy_load(self) -> Iterator[Document]:
         """Fetch text from one single GitBook page or recursively from sitemap."""
-        if self.load_all_paths:
-            # Get initial sitemap
-            soup_info = self.scrape()
+        if not self.load_all_paths:
+            # Simple case: load a single page
+            temp_loader = self._create_web_loader(self.web_page)
+            soup = temp_loader.scrape()
+            doc = self._get_document(soup, self.web_page)
+            if doc:
+                yield doc
+        else:
+            # Get initial sitemap using the recursive method
+            temp_loader = self._create_web_loader(self.start_url)
+            soup_info = temp_loader.scrape(parser="xml")
 
             # Process sitemap(s) recursively to get all content URLs
-            relative_paths = self._process_sitemap(soup_info)
+            processed_urls: Set[str] = set()
+            relative_paths = self._process_sitemap(soup_info, processed_urls)
+
             if not relative_paths and self.show_progress:
-                warnings.warn(
-                    f"No content URLs found in sitemap at {self.web_paths[0]}"
-                )
+                warnings.warn(f"No content URLs found in sitemap at {self.start_url}")
+
+            # Build full URLs from relative paths
+            urls: List[str] = []
+            for url in relative_paths:
+                # URLs are now already absolute from _get_paths
+                self._safe_add_url(urls, url, "content")
 
-            urls = [urljoin(self.base_url, path) for path in relative_paths]
+            if not urls:
+                return
 
-            # Use base class's scrape_all to efficiently fetch all pages
-            soup_infos = self.scrape_all(urls)
+            # Create a loader for content pages
+            content_loader = self._create_web_loader(urls)
+
+            # Use WebBaseLoader to fetch all pages
+            soup_infos = content_loader.scrape_all(urls)
 
             for soup_info, url in zip(soup_infos, urls):
                 doc = self._get_document(soup_info, url)
                 if doc:
                     yield doc
-        else:
-            # Use base class functionality directly for single page
-            for doc in super().lazy_load():
-                yield doc
 
     async def alazy_load(self) -> AsyncIterator[Document]:
         """Asynchronously fetch text from GitBook page(s)."""
         if not self.load_all_paths:
-            # For single page case, use the parent class implementation
-            async for doc in super().alazy_load():
+            # Simple case: load a single page asynchronously
+            temp_loader = self._create_web_loader(self.web_page)
+            soups = await temp_loader.ascrape_all([self.web_page])
+            soup_info = soups[0]
+            doc = self._get_document(soup_info, self.web_page)
+            if doc:
                 yield doc
         else:
-            # Fetch initial sitemap using base class's functionality
-            soups = await self.ascrape_all(self.web_paths, parser="xml")
+            # Get initial sitemap - web_loader will be created in _aprocess_sitemap
+            temp_loader = self._create_web_loader(self.start_url)
+            soups = await temp_loader.ascrape_all([self.start_url], parser="xml")
             soup_info = soups[0]
 
             # Process sitemap(s) recursively to get all content URLs
-            relative_paths = await self._aprocess_sitemap(soup_info, self.base_url)
+            processed_urls: Set[str] = set()
+            relative_paths = await self._aprocess_sitemap(
+                soup_info, self.base_url, processed_urls
+            )
+
             if not relative_paths and self.show_progress:
-                warnings.warn(
-                    f"No content URLs found in sitemap at {self.web_paths[0]}"
-                )
+                warnings.warn(f"No content URLs found in sitemap at {self.start_url}")
+
+            # Build full URLs from relative paths
+            urls: List[str] = []
+            for url in relative_paths:
+                # URLs are now already absolute from _get_paths
+                self._safe_add_url(urls, url, "content")
+
+            if not urls:
+                return
 
-            urls = [urljoin(self.base_url, path) for path in relative_paths]
+            # Create a loader for content pages
+            content_loader = self._create_web_loader(urls)
 
-            # Use base class's ascrape_all for efficient parallel fetching
-            soup_infos = await self.ascrape_all(urls)
+            # Use WebBaseLoader's ascrape_all for efficient parallel fetching
+            soup_infos = await content_loader.ascrape_all(urls)
 
             for soup_info, url in zip(soup_infos, urls):
                 maybe_doc = self._get_document(soup_info, url)
@@ -222,9 +355,15 @@ def _get_document(
         content = page_content_raw.get_text(separator="\n").strip()
         title_if_exists = page_content_raw.find("h1")
         title = title_if_exists.text if title_if_exists else ""
-        metadata = {"source": custom_url or self.web_path, "title": title}
+        metadata = {"source": custom_url or self.web_page, "title": title}
         return Document(page_content=content, metadata=metadata)
 
     def _get_paths(self, soup: Any) -> List[str]:
-        """Fetch all relative paths in the sitemap."""
-        return [urlparse(loc.text).path for loc in soup.find_all("loc")]
+        """Fetch all URLs in the sitemap."""
+        urls = []
+        for loc in soup.find_all("loc"):
+            if loc.text:
+                # Instead of extracting just the path, keep the full URL
+                # to preserve domain information
+                urls.append(loc.text)
+        return urls