Allow URL whitelist

elm/version.py

@@ -2,4 +2,4 @@
 ELM version number
 """
 
-__version__ = "0.0.42"
+__version__ = "0.0.43"

elm/web/search/run.py

@@ -7,6 +7,7 @@
 from collections import namedtuple
 from itertools import zip_longest, chain
 from contextlib import AsyncExitStack
+from warnings import warn
 
 from elm.web.file_loader import AsyncWebFileLoader
 from elm.web.search.bing import PlaywrightBingLinkSearch
@@ -59,7 +60,8 @@
 
 
 async def web_search_links_as_docs(queries, search_engines=_DEFAULT_SE,
-                                   num_urls=None, ignore_url_parts=None,
+                                   num_urls=None, url_ignore_substrings=None,
+                                   url_keep_substrings=None,
                                    search_semaphore=None,
                                    browser_semaphore=None, task_name=None,
                                    use_fallback_per_query=True,
@@ -92,10 +94,15 @@ async def web_search_links_as_docs(queries, search_engines=_DEFAULT_SE,
         number is less than ``len(queries)``, some of your queries may
         not contribute to the final output. By default, ``None``, which
         sets ``num_urls = 3 * len(queries)``.
-    ignore_url_parts : iterable of str, optional
+    url_ignore_substrings : iterable of str, optional
         Optional URL components to blacklist. For example, supplying
-        `ignore_url_parts={"wikipedia.org"}` will ignore all URLs that
-        contain "wikipedia.org". By default, ``None``.
+        `url_ignore_substrings={"wikipedia.org"}` will ignore all URLs
+        that contain "wikipedia.org". Substrings are applied
+        case-insensitively. By default, ``None``.
+    url_keep_substrings : list of str, optional
+        URL substrings that should be included in search results even if
+        they match an ignore substring. Substrings are applied
+        case-insensitively. By default, ``None``.
     search_semaphore : :class:`asyncio.Semaphore`, optional
         Semaphore instance that can be used to limit the number of
         playwright browsers used to submit search engine queries open
@@ -164,9 +171,11 @@ async def web_search_links_as_docs(queries, search_engines=_DEFAULT_SE,
         search_semaphore = browser_semaphore
 
     fpq = use_fallback_per_query
+    ignore, kwargs = _handle_old_ignore_key(url_ignore_substrings, kwargs)
     urls = await search_with_fallback(queries, search_engines=search_engines,
                                       num_urls=num_urls,
-                                      ignore_url_parts=ignore_url_parts,
+                                      url_ignore_substrings=ignore,
+                                      url_keep_substrings=url_keep_substrings,
                                       browser_semaphore=search_semaphore,
                                       task_name=task_name,
                                       use_fallback_per_query=fpq, **kwargs)
@@ -183,7 +192,8 @@ async def web_search_links_as_docs(queries, search_engines=_DEFAULT_SE,
 
 
 async def search_with_fallback(queries, search_engines=_DEFAULT_SE,
-                               num_urls=None, ignore_url_parts=None,
+                               num_urls=None, url_ignore_substrings=None,
+                               url_keep_substrings=None,
                                browser_semaphore=None, task_name=None,
                                use_fallback_per_query=True, **kwargs):
     """Retrieve search query URLs using multiple search engines if needed
@@ -211,10 +221,15 @@ async def search_with_fallback(queries, search_engines=_DEFAULT_SE,
         number is less than ``len(queries)``, some of your queries may
         not contribute to the final output. By default, ``None``, which
         sets ``num_urls = 3 * len(queries)``.
-    ignore_url_parts : iterable of str, optional
+    url_ignore_substrings : iterable of str, optional
         Optional URL components to blacklist. For example, supplying
-        `ignore_url_parts={"wikipedia.org"}` will ignore all URLs that
-        contain "wikipedia.org". By default, ``None``.
+        `url_ignore_substrings={"wikipedia.org"}` will ignore all URLs
+        that contain "wikipedia.org". Substrings are applied
+        case-insensitively. By default, ``None``.
+    url_keep_substrings : list of str, optional
+        URL substrings that should be included in search results even if
+        they match an ignore substring. Substrings are applied
+        case-insensitively. By default, ``None``.
     browser_semaphore : :class:`asyncio.Semaphore`, optional
         Semaphore instance that can be used to limit the number of
         playwright browsers open concurrently. If ``None``, no limits
@@ -274,17 +289,19 @@ async def search_with_fallback(queries, search_engines=_DEFAULT_SE,
         logger.error(msg)
         raise ELMInputError(msg)
 
+    ignore, kwargs = _handle_old_ignore_key(url_ignore_substrings, kwargs)
     if use_fallback_per_query:
         urls = await _multi_se_search(search_engines, queries, num_urls,
-                                      ignore_url_parts, browser_semaphore,
-                                      task_name, kwargs)
+                                      ignore, url_keep_substrings,
+                                      browser_semaphore, task_name, kwargs)
         if urls:
             return urls
     else:
         for se_name in search_engines:
             urls = await _single_se_search(se_name, queries, num_urls,
-                                           ignore_url_parts, browser_semaphore,
-                                           task_name, kwargs, raw=False)
+                                           ignore, url_keep_substrings,
+                                           browser_semaphore, task_name,
+                                           kwargs, raw=False)
             if urls:
                 return urls
 
@@ -294,8 +311,9 @@ async def search_with_fallback(queries, search_engines=_DEFAULT_SE,
 
 
 async def search_all_se(queries, search_engines=_DEFAULT_SE,
-                        num_urls=None, ignore_url_parts=None,
-                        browser_semaphore=None, task_name=None, **kwargs):
+                        num_urls=None, url_ignore_substrings=None,
+                        url_keep_substrings=None, browser_semaphore=None,
+                        task_name=None, **kwargs):
     """Retrieve search query URLs using multiple search engines if needed
 
     Parameters
@@ -321,10 +339,15 @@ async def search_all_se(queries, search_engines=_DEFAULT_SE,
         number is less than ``len(queries)``, some of your queries may
         not contribute to the final output. By default, ``None``, which
         sets ``num_urls = 3 * len(queries)``.
-    ignore_url_parts : iterable of str, optional
+    url_ignore_substrings : iterable of str, optional
         Optional URL components to blacklist. For example, supplying
-        `ignore_url_parts={"wikipedia.org"}` will ignore all URLs that
-        contain "wikipedia.org". By default, ``None``.
+        `url_ignore_substrings={"wikipedia.org"}` will ignore all URLs
+        that contain "wikipedia.org". Substrings are applied
+        case-insensitively. By default, ``None``.
+    url_keep_substrings : list of str, optional
+        URL substrings that should be included in search results even if
+        they match an ignore substring. Substrings are applied
+        case-insensitively. By default, ``None``.
     browser_semaphore : :class:`asyncio.Semaphore`, optional
         Semaphore instance that can be used to limit the number of
         playwright browsers open concurrently. If ``None``, no limits
@@ -387,9 +410,10 @@ async def search_all_se(queries, search_engines=_DEFAULT_SE,
         logger.error(msg)
         raise ELMInputError(msg)
 
+    ignore, kwargs = _handle_old_ignore_key(url_ignore_substrings, kwargs)
     searchers = [asyncio.create_task(
-        _single_se_search(se_name, queries, num_urls,
-                          ignore_url_parts, browser_semaphore,
+        _single_se_search(se_name, queries, num_urls, ignore,
+                          url_keep_substrings, browser_semaphore,
                           task_name, kwargs, raw=True),
         name=task_name) for se_name in search_engines]
 
@@ -432,8 +456,9 @@ async def load_docs(sources, file_loader):
     return docs
 
 
-async def _single_se_search(se_name, queries, num_urls, ignore_url_parts,
-                            browser_sem, task_name, kwargs, raw=False):
+async def _single_se_search(se_name, queries, num_urls, url_ignore_substrings,
+                            url_keep_substrings, browser_sem, task_name,
+                            kwargs, raw=False):
     """Search for links using a single search engine"""
     _validate_se_name(se_name)
     logger.debug("Searching web using %r", se_name)
@@ -442,11 +467,13 @@ async def _single_se_search(se_name, queries, num_urls, ignore_url_parts,
     if raw:
         return [link[0] for link in links]
     return _down_select_urls(links, num_urls=num_urls,
-                             ignore_url_parts=ignore_url_parts)
+                             url_ignore_substrings=url_ignore_substrings,
+                             url_keep_substrings=url_keep_substrings)
 
 
 async def _multi_se_search(search_engines, queries, num_urls,
-                           ignore_url_parts, browser_sem, task_name, kwargs):
+                           url_ignore_substrings, url_keep_substrings,
+                           browser_sem, task_name, kwargs):
     """Search for links using one or more search engines as fallback"""
     outputs = {q: None for q in queries}
     remaining_queries = list(queries)
@@ -474,7 +501,8 @@ async def _multi_se_search(search_engines, queries, num_urls,
     links = [link or [[]] for link in outputs.values()]
 
     return _down_select_urls(links, num_urls=num_urls,
-                             ignore_url_parts=ignore_url_parts)
+                             url_ignore_substrings=url_ignore_substrings,
+                             url_keep_substrings=url_keep_substrings)
 
 
 async def _run_search(se_name, queries, browser_sem, task_name, kwargs, raw):
@@ -539,27 +567,36 @@ def _init_se(se_name, kwargs):
     return se_class(**init_kwargs), uses_browser
 
 
-def _down_select_urls(search_results, num_urls=5, ignore_url_parts=None):
+def _down_select_urls(search_results, num_urls=5, url_ignore_substrings=None,
+                      url_keep_substrings=None):
     """Select the top N URLs"""
-    ignore_url_parts = _as_set(ignore_url_parts)
+    url_ignore_substrings = _as_set(url_ignore_substrings)
+    url_keep_substrings = _as_set(url_keep_substrings)
     all_urls = chain.from_iterable(zip_longest(*[results[0]
                                                  for results
                                                  in search_results]))
     urls = set()
     for url in all_urls:
-        if not url or any(substr in url for substr in ignore_url_parts):
+        if not url:
             continue
+
+        is_whitelisted = any(substr in url for substr in url_keep_substrings)
+        is_blacklisted = any(substr in url for substr in url_ignore_substrings)
+        if not is_whitelisted and is_blacklisted:
+            continue
+
         urls.add(url)
         if len(urls) == num_urls:
             break
+
     return urls
 
 
 def _as_set(user_input):
     """Convert user input (possibly None or str) to set of strings"""
     if isinstance(user_input, str):
         user_input = {user_input}
-    return set(user_input or [])
+    return {substr.casefold() for substr in (user_input or [])}
 
 
 def _validate_se_name(se_name):
@@ -569,3 +606,23 @@ def _validate_se_name(se_name):
                f"Got {se_name=}")
         logger.error(msg)
         raise ELMKeyError(msg)
+
+
+def _handle_old_ignore_key(url_ignore_substrings, kwargs):
+    """Handle old input gracefully"""
+    old_ignore_key = kwargs.pop("ignore_url_parts", None)
+    if old_ignore_key is None:
+        return url_ignore_substrings, kwargs
+
+    msg = ("`ignore_url_parts` is deprecated. Please use "
+           "`url_ignore_substrings` instead.")
+    warn(msg, DeprecationWarning)
+    if url_ignore_substrings is not None:
+        msg = ("Got both `ignore_url_parts` and `url_ignore_substrings`. "
+               "Using `url_ignore_substrings` and ignoring "
+               "`ignore_url_parts`.")
+        warn(msg, UserWarning)
+    else:
+        url_ignore_substrings = old_ignore_key
+
+    return url_ignore_substrings, kwargs

examples/web_information_retrieval/example_search_retrieval_wiki.ipynb

@@ -61,7 +61,7 @@
     "docs = await web_search_links_as_docs(\n",
     "    QUERIES,\n",
     "    pdf_read_kwargs={\"verbose\": False},\n",
-    "    ignore_url_parts={\"openei.org\"},\n",
+    "    url_ignore_substrings={\"openei.org\"},\n",
     ")"
    ]
   },

tests/web/search/test_web_search_run.py

@@ -16,7 +16,8 @@
 async def test_single_se_search_name_dne():
     """Test error for unknown search engine"""
     with pytest.raises(ELMKeyError) as err:
-        await _single_se_search("DNE", None, None, None, None, None, None)
+        await _single_se_search("DNE", None, None, None, None, None, None,
+                                None)
 
     assert "'se_name' must be one of" in str(err)
 
@@ -41,6 +42,25 @@ def test_down_select_urls_one_empty():
     assert _down_select_urls([[[]], [['bc', 'cd']]]) == {'bc', 'cd'}
 
 
+def test_down_select_urls_keep_substrings_override_ignore():
+    """Test keep substrings override ignored URLs"""
+    results = [[[
+        "https://blocked.com/keep-me",
+        "https://blocked.com/drop-me",
+    ]], [["https://allowed.com/keep"]]]
+
+    urls = _down_select_urls(
+        results,
+        url_ignore_substrings={"blocked.com"},
+        url_keep_substrings={"keep-me"},
+    )
+
+    assert urls == {
+        "https://blocked.com/keep-me",
+        "https://allowed.com/keep",
+    }
+
+
 def test_init_se():
     """Test initializing a playwright search engine"""
     test_kwargs = {"pw_launch_kwargs": {"test": 1}}
@@ -72,7 +92,7 @@ async def test_single_se_search_bad_build():
     """Test that bad init of SE gives no results"""
     test_kwargs = {"google_cse_api_kwargs": {"dne_arg": "test_key"}}
     results = await _single_se_search("APIGoogleCSESearch", [""], None, None,
-                                      None, None, test_kwargs)
+                                      None, None, None, test_kwargs)
     assert results == set()