feat(deep_crawling): allow configuring batch_size and dispatcher

crawl4ai/deep_crawling/bff_strategy.py

@@ -10,12 +10,12 @@
 from .scorers import URLScorer
 from . import DeepCrawlStrategy
 
-from ..types import AsyncWebCrawler, CrawlerRunConfig, CrawlResult, RunManyReturn
+from ..types import AsyncWebCrawler, BaseDispatcher, CrawlerRunConfig, CrawlResult, RunManyReturn
 from ..utils import normalize_url_for_deep_crawl
 
 from math import inf as infinity
 
-# Configurable batch size for processing items from the priority queue
+# Default batch size for processing items from the priority queue
 BATCH_SIZE = 10
 
 
@@ -47,13 +47,20 @@ def __init__(
         on_state_change: Optional[Callable[[Dict[str, Any]], Awaitable[None]]] = None,
         # Optional cancellation callback - checked before each URL is processed
         should_cancel: Optional[Callable[[], Union[bool, Awaitable[bool]]]] = None,
+        # Number of items pulled from the priority queue per round and handed
+        # to arun_many() at once. Defaults to 10.
+        batch_size: int = BATCH_SIZE,
+        # Optional dispatcher forwarded to arun_many() for each batch
+        dispatcher: Optional[BaseDispatcher] = None,
     ):
         self.max_depth = max_depth
         self.filter_chain = filter_chain
         self.url_scorer = url_scorer
         self.include_external = include_external
         self.score_threshold = score_threshold
         self.max_pages = max_pages
+        self.batch_size = batch_size
+        self.dispatcher = dispatcher
         # self.logger = logger or logging.getLogger(__name__)
         # Ensure logger is always a Logger instance, not a dict from serialization
         if isinstance(logger, logging.Logger):
@@ -245,15 +252,15 @@ async def _arun_best_first(
 
             # Calculate how many more URLs we can process in this batch
             remaining = self.max_pages - self._pages_crawled
-            batch_size = min(BATCH_SIZE, remaining)
-            if batch_size <= 0:
+            effective_batch_size = min(self.batch_size, remaining)
+            if effective_batch_size <= 0:
                 # No more pages to crawl
                 self.logger.info(f"Max pages limit ({self.max_pages}) reached, stopping crawl")
                 break
-                
+
             batch: List[Tuple[float, int, str, Optional[str]]] = []
-            # Retrieve up to BATCH_SIZE items from the priority queue.
-            for _ in range(BATCH_SIZE):
+            # Retrieve up to self.batch_size items from the priority queue.
+            for _ in range(self.batch_size):
                 if queue.empty():
                     break
                 item = await queue.get()
@@ -278,7 +285,12 @@ async def _arun_best_first(
             # make subsequent queue ordering depend on network timing.
             urls = [item[2] for item in batch]
             batch_config = config.clone(deep_crawl_strategy=None, stream=True)
-            stream_gen = await crawler.arun_many(urls=urls, config=batch_config)
+            arun_many_kwargs = (
+                {"dispatcher": self.dispatcher} if self.dispatcher is not None else {}
+            )
+            stream_gen = await crawler.arun_many(
+                urls=urls, config=batch_config, **arun_many_kwargs
+            )
             results_by_url: Dict[str, CrawlResult] = {}
             async for result in stream_gen:
                 results_by_url[result.url] = result

crawl4ai/deep_crawling/bfs_strategy.py

@@ -8,8 +8,8 @@
 from ..models import TraversalStats
 from .filters import FilterChain
 from .scorers import URLScorer
-from . import DeepCrawlStrategy  
-from ..types import AsyncWebCrawler, CrawlerRunConfig, CrawlResult
+from . import DeepCrawlStrategy
+from ..types import AsyncWebCrawler, BaseDispatcher, CrawlerRunConfig, CrawlResult
 from ..utils import normalize_url_for_deep_crawl, efficient_normalize_url_for_deep_crawl
 from math import inf as infinity
 
@@ -36,13 +36,16 @@ def __init__(
         on_state_change: Optional[Callable[[Dict[str, Any]], Awaitable[None]]] = None,
         # Optional cancellation callback - checked before each URL is processed
         should_cancel: Optional[Callable[[], Union[bool, Awaitable[bool]]]] = None,
+        # Optional dispatcher forwarded to arun_many() for each level
+        dispatcher: Optional[BaseDispatcher] = None,
     ):
         self.max_depth = max_depth
         self.filter_chain = filter_chain
         self.url_scorer = url_scorer
         self.include_external = include_external
         self.score_threshold = score_threshold
         self.max_pages = max_pages
+        self.dispatcher = dispatcher
         # self.logger = logger or logging.getLogger(__name__)
         # Ensure logger is always a Logger instance, not a dict from serialization
         if isinstance(logger, logging.Logger):
@@ -251,7 +254,15 @@ async def _arun_batch(
 
             # Clone the config to disable deep crawling recursion and enforce batch mode.
             batch_config = config.clone(deep_crawl_strategy=None, stream=False)
-            batch_results = await crawler.arun_many(urls=urls, config=batch_config)
+            # Only pass `dispatcher` when explicitly set, so the call shape is
+            # unchanged (and test doubles built against the old signature keep
+            # working) for the common case of relying on arun_many()'s own default.
+            arun_many_kwargs = (
+                {"dispatcher": self.dispatcher} if self.dispatcher is not None else {}
+            )
+            batch_results = await crawler.arun_many(
+                urls=urls, config=batch_config, **arun_many_kwargs
+            )
 
             for result in batch_results:
                 url = result.url
@@ -339,8 +350,13 @@ async def _arun_stream(
             visited.update(urls)
 
             stream_config = config.clone(deep_crawl_strategy=None, stream=True)
-            stream_gen = await crawler.arun_many(urls=urls, config=stream_config)
-
+            arun_many_kwargs = (
+                {"dispatcher": self.dispatcher} if self.dispatcher is not None else {}
+            )
+            stream_gen = await crawler.arun_many(
+                urls=urls, config=stream_config, **arun_many_kwargs
+            )
+
             # Keep track of processed results for this batch
             results_count = 0
             async for result in stream_gen:

tests/deep_crawling/test_deep_crawl_dispatcher_batch_size.py

@@ -0,0 +1,259 @@
+"""
+Test Suite: configurable `batch_size` (BestFirstCrawlingStrategy) and
+`dispatcher` (BFSDeepCrawlStrategy, BestFirstCrawlingStrategy) parameters.
+
+Covers:
+1. `batch_size` defaults to the previous hardcoded value (10) and, when
+   overridden, actually changes how many URLs are pulled from the priority
+   queue per round.
+2. `dispatcher` defaults to None and is NOT forwarded to arun_many() in that
+   case, so the call shape is unchanged for existing callers/test doubles.
+3. `dispatcher`, when explicitly set, is forwarded to arun_many() for both
+   BFSDeepCrawlStrategy and BestFirstCrawlingStrategy.
+"""
+
+import pytest
+from unittest.mock import MagicMock
+
+from crawl4ai.deep_crawling import BFSDeepCrawlStrategy, BestFirstCrawlingStrategy
+
+
+def create_mock_config(stream=False):
+    config = MagicMock()
+    config.clone = MagicMock(return_value=config)
+    config.stream = stream
+    return config
+
+
+def create_mock_crawler_old_signature():
+    """Mock crawler whose arun_many() only accepts (urls, config) — the
+    signature every caller used before `dispatcher` existed. If our strategy
+    code unconditionally passed `dispatcher=`, this would raise TypeError."""
+
+    async def mock_arun_many(urls, config):
+        results = []
+        for url in urls:
+            result = MagicMock()
+            result.url = url
+            result.success = True
+            result.metadata = {}
+            result.links = {"internal": [], "external": []}
+            results.append(result)
+        if config.stream:
+
+            async def gen():
+                for r in results:
+                    yield r
+
+            return gen()
+        return results
+
+    crawler = MagicMock()
+    crawler.arun_many = mock_arun_many
+    return crawler
+
+
+def create_mock_crawler_capturing(recorder: dict):
+    """Mock crawler that records the `dispatcher` kwarg and each batch of
+    urls it was called with (accepts the new signature)."""
+
+    async def mock_arun_many(urls, config, dispatcher=None):
+        recorder.setdefault("dispatcher_calls", []).append(dispatcher)
+        recorder.setdefault("batches", []).append(list(urls))
+        results = []
+        for url in urls:
+            result = MagicMock()
+            result.url = url
+            result.success = True
+            result.metadata = {}
+            result.links = {"internal": [], "external": []}
+            results.append(result)
+        if config.stream:
+
+            async def gen():
+                for r in results:
+                    yield r
+
+            return gen()
+        return results
+
+    crawler = MagicMock()
+    crawler.arun_many = mock_arun_many
+    return crawler
+
+
+def make_queue_items(n: int):
+    return [
+        {
+            "score": -i,
+            "depth": 0,
+            "url": f"https://example.com/p{i}",
+            "parent_url": None,
+        }
+        for i in range(n)
+    ]
+
+
+class TestBatchSizeDefaults:
+    def test_defaults_to_previous_hardcoded_value(self):
+        strategy = BestFirstCrawlingStrategy(max_depth=1)
+        assert strategy.batch_size == 10
+
+    def test_overridable_via_constructor(self):
+        strategy = BestFirstCrawlingStrategy(max_depth=1, batch_size=100)
+        assert strategy.batch_size == 100
+
+
+class TestBatchSizeBehavior:
+    @pytest.mark.asyncio
+    async def test_custom_batch_size_changes_round_size(self):
+        """With 12 queued URLs and batch_size=5, rounds should be 5, 5, 2 —
+        not the previous fixed 10, 2."""
+        resume_state = {
+            "visited": [],
+            "depths": {},
+            "pages_crawled": 0,
+            "queue_items": make_queue_items(12),
+        }
+        strategy = BestFirstCrawlingStrategy(
+            max_depth=1, max_pages=12, batch_size=5, resume_state=resume_state
+        )
+        recorder = {}
+        mock_crawler = create_mock_crawler_capturing(recorder)
+        # BestFirstCrawlingStrategy always treats its internal arun_many call as
+        # a stream generator (its own batch_config hardcodes stream=True),
+        # regardless of the outer config — so the mock config must say stream=True
+        # for the mock's async-generator branch to be used, matching reality.
+        mock_config = create_mock_config(stream=True)
+
+        await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
+
+        batch_sizes = [len(b) for b in recorder["batches"]]
+        assert batch_sizes == [5, 5, 2]
+
+    @pytest.mark.asyncio
+    async def test_default_batch_size_matches_old_behavior(self):
+        """With no batch_size override, rounds should still be 10, 2 (old default)."""
+        resume_state = {
+            "visited": [],
+            "depths": {},
+            "pages_crawled": 0,
+            "queue_items": make_queue_items(12),
+        }
+        strategy = BestFirstCrawlingStrategy(
+            max_depth=1, max_pages=12, resume_state=resume_state
+        )
+        recorder = {}
+        mock_crawler = create_mock_crawler_capturing(recorder)
+        # BestFirstCrawlingStrategy always treats its internal arun_many call as
+        # a stream generator (its own batch_config hardcodes stream=True),
+        # regardless of the outer config — so the mock config must say stream=True
+        # for the mock's async-generator branch to be used, matching reality.
+        mock_config = create_mock_config(stream=True)
+
+        await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
+
+        batch_sizes = [len(b) for b in recorder["batches"]]
+        assert batch_sizes == [10, 2]
+
+
+class TestDispatcherDefaultOmitted:
+    """Regression: arun_many() must NOT be called with `dispatcher=` when the
+    strategy's own dispatcher is None, so existing (old-signature) test
+    doubles/integrations keep working."""
+
+    @pytest.mark.asyncio
+    async def test_bfs_batch_mode_works_with_old_signature_mock(self):
+        strategy = BFSDeepCrawlStrategy(max_depth=1, max_pages=5)
+        mock_crawler = create_mock_crawler_old_signature()
+        mock_config = create_mock_config(stream=False)
+
+        results = await strategy._arun_batch(
+            "https://example.com", mock_crawler, mock_config
+        )
+        assert isinstance(results, list)
+        assert len(results) > 0
+
+    @pytest.mark.asyncio
+    async def test_bfs_stream_mode_works_with_old_signature_mock(self):
+        strategy = BFSDeepCrawlStrategy(max_depth=1, max_pages=5)
+        mock_crawler = create_mock_crawler_old_signature()
+        mock_config = create_mock_config(stream=True)
+
+        results = [
+            r
+            async for r in strategy._arun_stream(
+                "https://example.com", mock_crawler, mock_config
+            )
+        ]
+        assert len(results) > 0
+
+    @pytest.mark.asyncio
+    async def test_best_first_works_with_old_signature_mock(self):
+        strategy = BestFirstCrawlingStrategy(max_depth=1, max_pages=5)
+        mock_crawler = create_mock_crawler_old_signature()
+        mock_config = create_mock_config(
+            stream=True
+        )  # BestFirst always streams internally
+
+        results = await strategy._arun_batch(
+            "https://example.com", mock_crawler, mock_config
+        )
+        assert isinstance(results, list)
+        assert len(results) > 0
+
+
+class TestDispatcherForwarded:
+    """When a dispatcher IS set, it must actually reach arun_many()."""
+
+    @pytest.mark.asyncio
+    async def test_bfs_batch_mode_forwards_dispatcher(self):
+        sentinel_dispatcher = object()
+        strategy = BFSDeepCrawlStrategy(
+            max_depth=1, max_pages=5, dispatcher=sentinel_dispatcher
+        )
+        recorder = {}
+        mock_crawler = create_mock_crawler_capturing(recorder)
+        mock_config = create_mock_config(stream=False)
+
+        await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
+
+        assert recorder["dispatcher_calls"]
+        assert all(d is sentinel_dispatcher for d in recorder["dispatcher_calls"])
+
+    @pytest.mark.asyncio
+    async def test_bfs_stream_mode_forwards_dispatcher(self):
+        sentinel_dispatcher = object()
+        strategy = BFSDeepCrawlStrategy(
+            max_depth=1, max_pages=5, dispatcher=sentinel_dispatcher
+        )
+        recorder = {}
+        mock_crawler = create_mock_crawler_capturing(recorder)
+        mock_config = create_mock_config(stream=True)
+
+        results = [
+            r
+            async for r in strategy._arun_stream(
+                "https://example.com", mock_crawler, mock_config
+            )
+        ]
+        assert len(results) > 0
+        assert recorder["dispatcher_calls"]
+        assert all(d is sentinel_dispatcher for d in recorder["dispatcher_calls"])
+
+    @pytest.mark.asyncio
+    async def test_best_first_forwards_dispatcher(self):
+        sentinel_dispatcher = object()
+        strategy = BestFirstCrawlingStrategy(
+            max_depth=1, max_pages=5, dispatcher=sentinel_dispatcher
+        )
+        recorder = {}
+        mock_crawler = create_mock_crawler_capturing(recorder)
+        mock_config = create_mock_config(
+            stream=True
+        )  # BestFirst always streams internally
+
+        await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
+
+        assert recorder["dispatcher_calls"]
+        assert all(d is sentinel_dispatcher for d in recorder["dispatcher_calls"])