diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index d432dcfbf1..7d2cdc9ff1 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -50,6 +50,8 @@ logger = get_logger(__name__) _EXEC_TIMEOUT = 300 # 5 minutes _RAG_SEARCH_SLOT = threading.BoundedSemaphore(1) +# Candidate multiplier when a website policy will filter the results after the search. +_POLICY_OVERFETCH = 4 _DISABLE_DNS_PINNING_ENV = "UNSLOTH_STUDIO_DISABLE_DNS_PINNING" # Splits the UI source-map from the result; loops strip it (like __IMAGES__). @@ -4640,13 +4642,19 @@ def _web_search( from .web_access_policy import check_url_access, scope_search_query effective_query = scope_search_query(query, website_policy) - results = DDGS(timeout = timeout).text(effective_query, max_results = max_results) + # The policy filters below, so ask for a deeper pool when one is set: otherwise a page + # whose top hits are all disallowed yields nothing even when valid results rank just + # under them, wasting a research step. + wanted = max_results * _POLICY_OVERFETCH if website_policy else max_results + results = DDGS(timeout = timeout).text(effective_query, max_results = wanted) if cancel_event is not None and cancel_event.is_set(): return "Search cancelled." if not results: return "No results found." parts = [] for r in results: + if len(parts) >= max_results: + break href = str(r.get("href") or "").strip() allowed, _reason, _hostname = check_url_access(href, website_policy) if not allowed: diff --git a/studio/backend/core/inference/web_access_policy.py b/studio/backend/core/inference/web_access_policy.py index 21134f05eb..52becd450d 100644 --- a/studio/backend/core/inference/web_access_policy.py +++ b/studio/backend/core/inference/web_access_policy.py @@ -7,11 +7,14 @@ from __future__ import annotations import ipaddress import re +import zlib from typing import Any from urllib.parse import urlsplit _DOMAIN_LABEL = re.compile(r"^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?$") _MAX_DOMAINS_PER_LIST = 100 +# Most search engines stop honouring site: past a handful of OR terms. +_SITE_FILTER_LIMIT = 8 def normalize_domain(value: Any) -> str: @@ -139,5 +142,12 @@ def scope_search_query(query: str, policy: dict[str, Any] | None) -> str: return query # Cap the site: filter (search engines limit OR operators) instead of dropping scoping # entirely for large allow lists, which returned unrelated results that all got filtered out. - site_filter = " OR ".join(f"site:{domain}" for domain in allowed[:8]) + # Rotate the window by query so every allowed domain is reachable across a multi-step run; + # a fixed head made domains past the cap permanently undiscoverable. Keyed on the query so + # the same search is always scoped the same way. + window = allowed + if len(allowed) > _SITE_FILTER_LIMIT: + offset = zlib.crc32(query.encode("utf-8")) % len(allowed) + window = (allowed + allowed)[offset : offset + _SITE_FILTER_LIMIT] + site_filter = " OR ".join(f"site:{domain}" for domain in window) return f"{query} ({site_filter})" diff --git a/studio/backend/tests/test_web_access_policy.py b/studio/backend/tests/test_web_access_policy.py index 6f05c3700f..57dbf9a7d9 100644 --- a/studio/backend/tests/test_web_access_policy.py +++ b/studio/backend/tests/test_web_access_policy.py @@ -121,12 +121,72 @@ def test_web_search_filters_results_before_model_exposure(monkeypatch): monkeypatch.setitem(sys.modules, "ddgs", SimpleNamespace(DDGS = FakeDDGS)) result = tools._web_search("latest paper", website_policy = ARXIV_ONLY) - assert queries == [("latest paper (site:arxiv.org)", 5)] + # A policy filters after the search, so a deeper candidate pool is requested. + assert queries == [("latest paper (site:arxiv.org)", 5 * tools._POLICY_OVERFETCH)] assert "https://arxiv.org/abs/1" in result assert "example.com" not in result assert "arxiv.org.evil.test" not in result +def test_web_search_refills_past_disallowed_results(monkeypatch): + # Without over-fetching, a page whose top hits are all blocked returned nothing even though + # valid results ranked just below them, wasting a research step. + blocked_then_allowed = [ + {"title": "Bad", "href": f"https://example.com/{i}", "body": "Blocked"} for i in range(5) + ] + [{"title": "Good", "href": f"https://arxiv.org/abs/{i}", "body": "Allowed"} for i in range(5)] + + class FakeDDGS: + def __init__(self, **_kwargs): + pass + + def text(self, query, max_results = 5): + return blocked_then_allowed[:max_results] + + monkeypatch.setitem(sys.modules, "ddgs", SimpleNamespace(DDGS = FakeDDGS)) + result = tools._web_search("q", website_policy = {"blockedDomains": ["example.com"]}) + + assert "arxiv.org/abs/0" in result + assert "example.com" not in result + # Still capped at max_results allowed entries, not the whole deeper pool. + assert result.count("Title: ") == 5 + + +def test_web_search_without_a_policy_does_not_overfetch(monkeypatch): + queries = [] + + class FakeDDGS: + def __init__(self, **_kwargs): + pass + + def text(self, query, max_results = 5): + queries.append((query, max_results)) + return [{"title": "T", "href": "https://a.example/1", "body": "B"}] + + monkeypatch.setitem(sys.modules, "ddgs", SimpleNamespace(DDGS = FakeDDGS)) + tools._web_search("q", website_policy = None) + assert queries == [("q", 5)] + + +def test_scope_search_query_reaches_every_allowed_domain(): + # The site: filter is capped because engines stop honouring long OR chains, but a fixed + # head made domains past the cap permanently undiscoverable. + domains = [f"d{i}.example" for i in range(20)] + policy = {"allowedDomains": domains} + covered = set() + for i in range(200): + scoped = scope_search_query(f"query {i}", policy) + hits = [d for d in domains if f"site:{d}" in scoped] + assert len(hits) == 8 + covered.update(hits) + assert covered == set(domains) + # Deterministic: the same query always scopes the same way. + assert scope_search_query("stable", policy) == scope_search_query("stable", policy) + # At or under the cap every domain is always included. + small = [f"s{i}.example" for i in range(8)] + scoped = scope_search_query("q", {"allowedDomains": small}) + assert all(f"site:{d}" in scoped for d in small) + + def test_web_search_flattens_source_framing_in_untrusted_metadata(monkeypatch): class FakeDDGS: def __init__(self, **_kwargs):