diff --git a/elm/version.py b/elm/version.py index 7bf6110e..6fb1b13a 100644 --- a/elm/version.py +++ b/elm/version.py @@ -2,4 +2,4 @@ ELM version number """ -__version__ = "0.0.42" +__version__ = "0.0.43" diff --git a/elm/web/search/run.py b/elm/web/search/run.py index e8f6398a..67afd3a2 100644 --- a/elm/web/search/run.py +++ b/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,19 +567,28 @@ 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 @@ -559,7 +596,7 @@ 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 \ No newline at end of file diff --git a/examples/web_information_retrieval/example_search_retrieval_wiki.ipynb b/examples/web_information_retrieval/example_search_retrieval_wiki.ipynb index 88b9ea19..836e837f 100644 --- a/examples/web_information_retrieval/example_search_retrieval_wiki.ipynb +++ b/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", ")" ] }, diff --git a/tests/web/search/test_web_search_run.py b/tests/web/search/test_web_search_run.py index 01fd1aa2..af154796 100644 --- a/tests/web/search/test_web_search_run.py +++ b/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()