Give user more control over docling processing

compass/pipeline/collection/steps.py

@@ -1,5 +1,6 @@
 """Fixed collection steps for the process pipeline"""
 
+import asyncio
 import logging
 from abc import ABC, abstractmethod
 
@@ -267,22 +268,35 @@ async def collect(self, workflow):  # ruff:ignore[no-self-use]
             )
             return []
 
+        crawl_timeout_s = (
+            workflow.runtime.search_params.website_crawl_timeout_seconds
+        )
         logger.debug(
             "Collecting documents using ELM web crawl for: %s",
             workflow.jurisdiction.full_name,
         )
         try:
-            out = await download_jurisdiction_ordinances_from_website(
-                workflow.jurisdiction_website,
-                heuristic=await workflow.extractor.get_heuristic(),
-                keyword_points=(
-                    await workflow.extractor.get_website_keywords()
-                ),
-                file_loader_kwargs=workflow.runtime.file_loader_kwargs,
-                crawl_semaphore=workflow.runtime.crawl_semaphore,
-                pb_jurisdiction_name=workflow.jurisdiction.full_name,
-                return_c4ai_results=True,
+            async with workflow.runtime.crawl_semaphore:
+                async with asyncio.timeout(crawl_timeout_s):
+                    out = await download_jurisdiction_ordinances_from_website(
+                        workflow.jurisdiction_website,
+                        heuristic=await workflow.extractor.get_heuristic(),
+                        keyword_points=(
+                            await workflow.extractor.get_website_keywords()
+                        ),
+                        file_loader_kwargs=workflow.runtime.file_loader_kwargs,
+                        pb_jurisdiction_name=workflow.jurisdiction.full_name,
+                        return_c4ai_results=True,
+                    )
+        except TimeoutError:
+            logger.exception(
+                "ELM Website crawl deadline (%s) exceeded for %s; "
+                "continuing with no crawl docs",
+                f"{int(crawl_timeout_s):,d}s",
+                workflow.jurisdiction.full_name,
             )
+            workflow.last_scrape_results = []
+            return []
         except Exception:
             logger.exception(
                 "Error collecting documents using ELM web crawl for %s",
@@ -358,17 +372,32 @@ async def collect(self, workflow):  # ruff:ignore[no-self-use]
         for scrape_result in workflow.last_scrape_results:
             checked_urls.update({sub_res.url for sub_res in scrape_result})
 
+        crawl_timeout_s = (
+            workflow.runtime.search_params.website_crawl_timeout_seconds
+        )
         func = download_jurisdiction_ordinances_from_website_compass_crawl
         try:
-            docs = await func(
-                workflow.jurisdiction_website,
-                heuristic=await workflow.extractor.get_heuristic(),
-                keyword_points=await workflow.extractor.get_website_keywords(),
-                file_loader_kwargs=workflow.runtime.file_loader_kwargs,
-                already_visited=checked_urls,
-                crawl_semaphore=workflow.runtime.crawl_semaphore,
-                pb_jurisdiction_name=workflow.jurisdiction.full_name,
+            async with workflow.runtime.crawl_semaphore:
+                async with asyncio.timeout(crawl_timeout_s):
+                    docs = await func(
+                        workflow.jurisdiction_website,
+                        heuristic=await workflow.extractor.get_heuristic(),
+                        keyword_points=(
+                            await workflow.extractor.get_website_keywords()
+                        ),
+                        file_loader_kwargs=workflow.runtime.file_loader_kwargs,
+                        already_visited=checked_urls,
+                        pb_jurisdiction_name=workflow.jurisdiction.full_name,
+                    )
+        except TimeoutError:
+            logger.exception(
+                "COMPASS Website crawl deadline (%s) exceeded for %s; "
+                "continuing with no crawl docs",
+                f"{int(crawl_timeout_s):,d}s",
+                workflow.jurisdiction.full_name,
             )
+            workflow.last_scrape_results = []
+            return []
         except Exception:
             logger.exception(
                 "Error collecting documents using COMPASS web crawl for %s",

compass/pipeline/data_classes.py

@@ -181,6 +181,7 @@ def __init__(
         num_urls_to_check_per_jurisdiction=5,
         max_num_concurrent_browsers=10,
         max_num_concurrent_website_searches=None,
+        website_crawl_timeout_seconds=3600,
         url_ignore_substrings=None,
         url_keep_substrings=None,
         search_engines=None,
@@ -205,6 +206,11 @@ def __init__(
             simultaneously. Increasing this value can speed up searches,
             but may lead to timeouts or performance issues on machines
             with limited resources. By default, ``10``.
+        website_crawl_timeout_seconds : int, default=3600
+            Maximum number of seconds to allow for a website crawl to
+            complete before timing out. If the crawl exceeds this time,
+            it will be terminated and no documents will be returned for
+            the crawl step for the jurisdiction. By default, ``3600``
         url_ignore_substrings : list of str, optional
             A list of substrings that, if found in any URL, will cause
             the URL to be excluded from consideration. This can be used
@@ -273,6 +279,7 @@ def __init__(
         self.max_num_concurrent_website_searches = (
             max_num_concurrent_website_searches
         )
+        self.website_crawl_timeout_seconds = website_crawl_timeout_seconds
         self.url_ignore_substrings = _DOMAINS["blacklist"]
         self.url_ignore_substrings += url_ignore_substrings or []
         self.url_keep_substrings = _DOMAINS["whitelist"]
@@ -334,6 +341,7 @@ def __init__(  # ruff:ignore[too-many-arguments]
         max_num_concurrent_browsers=10,
         max_num_concurrent_website_searches=10,
         max_num_concurrent_jurisdictions=25,
+        website_crawl_timeout_seconds=3600,
         url_ignore_substrings=None,
         url_keep_substrings=None,
         known_local_docs=None,
@@ -474,6 +482,11 @@ def __init__(  # ruff:ignore[too-many-arguments]
             Maximum number of jurisdictions to process concurrently.
             Limiting this can help manage memory usage when dealing with
             a large number of documents. By default, ``25``.
+        website_crawl_timeout_seconds : int, default=3600
+            Maximum number of seconds to allow for a website crawl to
+            complete before timing out. If the crawl exceeds this time,
+            it will be terminated and no documents will be returned for
+            the crawl step for the jurisdiction. By default, ``3600``
         url_ignore_substrings : list of str, optional
             A list of substrings that, if found in any URL, will cause
             the URL to be excluded from consideration. This can be used
@@ -634,6 +647,7 @@ def __init__(  # ruff:ignore[too-many-arguments]
             max_num_concurrent_website_searches=(
                 max_num_concurrent_website_searches
             ),
+            website_crawl_timeout_seconds=website_crawl_timeout_seconds,
             url_ignore_substrings=url_ignore_substrings,
             url_keep_substrings=url_keep_substrings,
             search_engines=search_engines,
@@ -702,6 +716,7 @@ def __init__(  # ruff:ignore[too-many-arguments]
         max_num_concurrent_browsers=10,
         max_num_concurrent_website_searches=10,
         max_num_concurrent_jurisdictions=25,
+        website_crawl_timeout_seconds=3600,
         url_ignore_substrings=None,
         url_keep_substrings=None,
         known_local_docs=None,
@@ -819,6 +834,11 @@ def __init__(  # ruff:ignore[too-many-arguments]
             Maximum number of jurisdictions to process concurrently.
             Limiting this can help manage memory usage when dealing with
             a large number of documents. By default, ``25``.
+        website_crawl_timeout_seconds : int, default=3600
+            Maximum number of seconds to allow for a website crawl to
+            complete before timing out. If the crawl exceeds this time,
+            it will be terminated and no documents will be returned for
+            the crawl step for the jurisdiction. By default, ``3600``
         url_ignore_substrings : list of str, optional
             A list of substrings that, if found in any URL, will cause
             the URL to be excluded from consideration. This can be used
@@ -996,6 +1016,7 @@ def __init__(  # ruff:ignore[too-many-arguments]
                 max_num_concurrent_website_searches
             ),
             max_num_concurrent_jurisdictions=max_num_concurrent_jurisdictions,
+            website_crawl_timeout_seconds=website_crawl_timeout_seconds,
             url_ignore_substrings=url_ignore_substrings,
             url_keep_substrings=url_keep_substrings,
             known_local_docs=known_local_docs,

compass/pipeline/runtime.py

@@ -116,14 +116,21 @@ def browser_semaphore(self):
         )
 
     @cached_property
-    def crawl_semaphore(self):
-        """Crawl concurrency limiter"""
+    def _crawl_semaphore(self):
+        """Crawl concurrency limiter or None"""
         if not self.search_params.max_num_concurrent_website_searches:
             return None
         return asyncio.Semaphore(
             self.search_params.max_num_concurrent_website_searches
         )
 
+    @property
+    def crawl_semaphore(self):
+        """Crawl concurrency limiter"""
+        if self._crawl_semaphore is None:
+            return AsyncExitStack()
+        return self._crawl_semaphore
+
     @cached_property
     def search_engine_semaphore(self):
         """Search engine concurrency limiter"""
@@ -172,6 +179,9 @@ def local_file_loader_kwargs(self):
             "html_read_kwargs": self.file_loader_kwargs.get(
                 "html_read_kwargs"
             ),
+            "pdf_pipeline_options": self.file_loader_kwargs.get(
+                "pdf_pipeline_options"
+            ),
         }
         if self.search_params.pytesseract_exe_fp is not None:
             self._setup_pytesseract()
@@ -240,7 +250,9 @@ def _base_services(self):
             )
 
         if self.search_params.pytesseract_exe_fp is not None:
-            services.append(OCRPDFLoader(max_workers=1))
+            kwargs = deepcopy(runtime_settings.ppe_kwargs or {})
+            kwargs["max_workers"] = 1
+            services.append(OCRPDFLoader(**kwargs))
         return services
 
     @cached_property

compass/plugin/one_shot/base.py

@@ -590,14 +590,12 @@ def _out_cols_from_config(config):
     )
 
     source_col_ind = next(
-        (ind for ind, col in enumerate(cols) if col.name == "source"),
-        None,
+        (ind for ind, col in enumerate(cols) if col.name == "source"), None
     )
-    year_col = OutputColumn("year")
     if source_col_ind is None:
-        cols.append(year_col)
+        cols.extend((OutputColumn("year"), OutputColumn("source")))
     else:
-        cols.insert(source_col_ind, year_col)
+        cols.insert(source_col_ind, OutputColumn("year"))
 
     cols.append(
         OutputColumn(

compass/scripts/download.py

@@ -272,7 +272,6 @@ async def download_jurisdiction_ordinances_from_website(
     browser_config_kwargs=None,
     crawler_config_kwargs=None,
     max_urls=100,
-    crawl_semaphore=None,
     pb_jurisdiction_name=None,
     return_c4ai_results=False,
 ):
@@ -307,10 +306,6 @@ async def download_jurisdiction_ordinances_from_website(
     max_urls : int, optional
         Max number of URLs to check from the website before terminating
         the search. By default, ``100``.
-    crawl_semaphore : :class:`asyncio.Semaphore`, optional
-        Semaphore instance that can be used to limit the number of
-        website searches happening concurrently. If ``None``, no limits
-        are applied. By default, ``None``.
     pb_jurisdiction_name : str, optional
         Optional jurisdiction name to use to update progress bar, if
         it's being used. By default, ``None``.
@@ -337,9 +332,6 @@ async def download_jurisdiction_ordinances_from_website(
     to be running.
     """
 
-    if crawl_semaphore is None:
-        crawl_semaphore = AsyncExitStack()
-
     async def _doc_heuristic(doc):  # ruff:ignore[unused-async]
         """Heuristic check for wind ordinance documents"""
         is_valid_document = heuristic.check(doc.text.lower())
@@ -393,7 +385,7 @@ async def _crawl_hook(*__, **___):  # ruff:ignore[unused-async]
         cpb = AsyncExitStack()
         ch = None
 
-    async with crawl_semaphore, cpb:
+    async with cpb:
         docs_or_pair = await crawler.run(
             website,
             on_result_hook=ch,
@@ -416,7 +408,6 @@ async def download_jurisdiction_ordinances_from_website_compass_crawl(
     already_visited=None,
     num_link_scores_to_check_per_page=4,
     max_urls=100,
-    crawl_semaphore=None,
     pb_jurisdiction_name=None,
 ):
     """Download ord documents from a website using the COMPASS crawler
@@ -452,10 +443,6 @@ async def download_jurisdiction_ordinances_from_website_compass_crawl(
     max_urls : int, default=100
         Max number of URLs to check from the website before terminating
         the search. By default, ``100``.
-    crawl_semaphore : :class:`asyncio.Semaphore`, optional
-        Semaphore instance that can be used to limit the number of
-        website crawls happening concurrently. If ``None``, no limits
-        are applied. By default, ``None``.
     pb_jurisdiction_name : str, optional
         Optional jurisdiction name to use to update progress bar, if
         it's being used. By default, ``None``.
@@ -472,8 +459,6 @@ async def download_jurisdiction_ordinances_from_website_compass_crawl(
     Requires :class:`~compass.services.threaded.TempFileCache` service
     to be running.
     """
-    if crawl_semaphore is None:
-        crawl_semaphore = AsyncExitStack()
 
     async def _doc_heuristic(doc):  # ruff:ignore[unused-async]
         """Heuristic check for wind ordinance documents"""
@@ -515,7 +500,7 @@ async def _crawl_hook(*__, **___):  # ruff:ignore[unused-async]
         cpb = AsyncExitStack()
         ch = None
 
-    async with crawl_semaphore, cpb:
+    async with cpb:
         return await crawler.run(website, on_new_page_visit_hook=ch)