From 7ff4057fc9456b1dd8b1e336f45b076b1cb81a1b Mon Sep 17 00:00:00 2001 From: Jacobinwwey Date: Tue, 31 Mar 2026 02:06:51 -0700 Subject: [PATCH] feat(web-search): add provider routing aliases and explicit-engine selection while preserving default engine behavior (#6442) * feat(web-search): add provider routing and ddg-first default chain * fix(web-search): keep dev default engine order and demote ddg fallback * fix(web-search): harden selector/query handling and unify engine naming * refactor(web-search): streamline provider wiring and harden google query encoding * fix(web-search): return empty for non-positive filtered result limits * refactor(web-search): share provider constants and validate engine registry * refactor(web-search): unify provider normalization and engine registry * fix(web-search): reduce comet link noise with stricter url filtering * style(web-search): apply ruff formatting --- .../web_searcher/engines/__init__.py | 34 ++++ .../web_searcher/engines/bing.py | 5 +- .../web_searcher/engines/comet.py | 68 +++++++ .../web_searcher/engines/duckduckgo.py | 44 +++++ .../web_searcher/engines/google.py | 52 ++++++ .../web_searcher/engines/sogo.py | 2 + astrbot/builtin_stars/web_searcher/main.py | 90 +++++++-- .../web_searcher/provider_constants.py | 24 +++ .../web_searcher/provider_routing.py | 132 +++++++++++++ astrbot/core/config/default.py | 8 +- tests/unit/test_web_search_engines.py | 173 ++++++++++++++++++ .../unit/test_web_search_provider_routing.py | 110 +++++++++++ 12 files changed, 719 insertions(+), 23 deletions(-) create mode 100644 astrbot/builtin_stars/web_searcher/engines/comet.py create mode 100644 astrbot/builtin_stars/web_searcher/engines/duckduckgo.py create mode 100644 astrbot/builtin_stars/web_searcher/engines/google.py create mode 100644 astrbot/builtin_stars/web_searcher/provider_constants.py create mode 100644 astrbot/builtin_stars/web_searcher/provider_routing.py create mode 100644 tests/unit/test_web_search_engines.py create mode 100644 tests/unit/test_web_search_provider_routing.py diff --git a/astrbot/builtin_stars/web_searcher/engines/__init__.py b/astrbot/builtin_stars/web_searcher/engines/__init__.py index b9041f282..d1a0ada15 100644 --- a/astrbot/builtin_stars/web_searcher/engines/__init__.py +++ b/astrbot/builtin_stars/web_searcher/engines/__init__.py @@ -1,5 +1,6 @@ import random import urllib.parse +from collections.abc import Callable from dataclasses import dataclass from aiohttp import ClientSession @@ -95,6 +96,12 @@ class SearchEngine: soup = BeautifulSoup(resp, "html.parser") links = soup.select(self._set_selector("links")) results = [] + try: + text_selector = self._set_selector("text") + except (KeyError, NotImplementedError): + # Keep backward compatibility with engines that only expose + # title/url/link selectors and do not provide snippets. + text_selector = "" for link in links: # Safely get the title text (select_one may return None) title_elem = link.select_one(self._set_selector("title")) @@ -104,9 +111,36 @@ class SearchEngine: url_tag = link.select_one(self._set_selector("url")) snippet = "" + if text_selector: + text_elem = link.select_one(text_selector) + if text_elem is not None: + snippet = self.tidy_text(text_elem.get_text()) if title and url_tag: url = self._get_url(url_tag) + if not url: + continue + if url.startswith("//"): + url = f"https:{url}" results.append(SearchResult(title=title, url=url, snippet=snippet)) return results[:num_results] if len(results) > num_results else results except Exception as e: raise e + + async def _search_with_result_filter( + self, + query: str, + num_results: int, + predicate: Callable[[SearchResult], bool], + ) -> list[SearchResult]: + if num_results <= 0: + return [] + + rough_results = await SearchEngine.search(self, query, max(num_results * 2, 10)) + final_results: list[SearchResult] = [] + for result in rough_results: + if not predicate(result): + continue + final_results.append(result) + if len(final_results) >= num_results: + break + return final_results diff --git a/astrbot/builtin_stars/web_searcher/engines/bing.py b/astrbot/builtin_stars/web_searcher/engines/bing.py index 7565e5df3..072000faf 100644 --- a/astrbot/builtin_stars/web_searcher/engines/bing.py +++ b/astrbot/builtin_stars/web_searcher/engines/bing.py @@ -2,9 +2,12 @@ from . import USER_AGENT_BING, SearchEngine class Bing(SearchEngine): + NAME = "bing" + def __init__(self) -> None: super().__init__() - self.base_urls = ["https://cn.bing.com", "https://www.bing.com"] + # Prefer international Bing first, keep cn endpoint as compatibility fallback. + self.base_urls = ["https://www.bing.com", "https://cn.bing.com"] self.headers.update({"User-Agent": USER_AGENT_BING}) def _set_selector(self, selector: str): diff --git a/astrbot/builtin_stars/web_searcher/engines/comet.py b/astrbot/builtin_stars/web_searcher/engines/comet.py new file mode 100644 index 000000000..7b276a4ca --- /dev/null +++ b/astrbot/builtin_stars/web_searcher/engines/comet.py @@ -0,0 +1,68 @@ +from typing import cast +from urllib.parse import unquote, urlencode, urlparse + +from bs4 import Tag + +from . import SearchEngine, SearchResult + + +class Comet(SearchEngine): + """Best-effort search via public Perplexity/Comet page. + + Note: + - This endpoint is often protected by anti-bot challenges. + - We intentionally treat failures as non-fatal and rely on fallback engines. + """ + + NAME = "comet" + + def __init__(self) -> None: + super().__init__() + self.base_url = "https://www.perplexity.ai" + + def _set_selector(self, selector: str): + selectors = { + "url": "a[href^='http'], a[href^='//']", + "title": "main h1, main h2, main h3, h3, h2", + "text": "main article, main div[role='article'], main section, main p, p", + "links": ( + "main article, main div[role='article'], main li, main div.result, " + "article, div[role='article'], li, div.result" + ), + "next": "", + } + return selectors[selector] + + async def _get_next_page(self, query: str) -> str: + url = f"{self.base_url}/search?{urlencode({'q': unquote(query)})}" + return await self._get_html(url, None) + + def _get_url(self, tag: Tag) -> str: + href = cast(str, tag.get("href") or "") + if href.startswith("//"): + return f"https:{href}" + return href + + @staticmethod + def _is_valid_result_url(url: str) -> bool: + lowered = (url or "").strip().lower() + if not lowered: + return False + if lowered.startswith(("#", "javascript:", "mailto:")): + return False + if not lowered.startswith(("http://", "https://")): + return False + + netloc = urlparse(lowered).netloc + if not netloc: + return False + if netloc.endswith("perplexity.ai"): + return False + return True + + async def search(self, query: str, num_results: int) -> list[SearchResult]: + return await self._search_with_result_filter( + query=query, + num_results=num_results, + predicate=lambda result: self._is_valid_result_url(result.url), + ) diff --git a/astrbot/builtin_stars/web_searcher/engines/duckduckgo.py b/astrbot/builtin_stars/web_searcher/engines/duckduckgo.py new file mode 100644 index 000000000..c82da2a61 --- /dev/null +++ b/astrbot/builtin_stars/web_searcher/engines/duckduckgo.py @@ -0,0 +1,44 @@ +import urllib.parse +from typing import cast + +from bs4 import Tag + +from . import SearchEngine, SearchResult + + +class DuckDuckGo(SearchEngine): + NAME = "duckduckgo" + + def __init__(self) -> None: + super().__init__() + self.base_url = "https://html.duckduckgo.com/html" + + def _set_selector(self, selector: str): + selectors = { + "url": "a.result__a, h2 a", + "title": "a.result__a, h2", + "text": "a.result__snippet, div.result__snippet", + "links": "div.result, div.web-result", + "next": "a.result--more__btn", + } + return selectors[selector] + + async def _get_next_page(self, query: str) -> str: + params = {"q": urllib.parse.unquote(query), "kl": "us-en"} + url = f"{self.base_url}/?{urllib.parse.urlencode(params)}" + return await self._get_html(url, None) + + def _get_url(self, tag: Tag) -> str: + href = cast(str, tag.get("href") or "") + if "duckduckgo.com/l/?" in href: + parsed = urllib.parse.urlparse(href) + target = urllib.parse.parse_qs(parsed.query).get("uddg", [""])[0] + return urllib.parse.unquote(target) + return href + + async def search(self, query: str, num_results: int) -> list[SearchResult]: + return await self._search_with_result_filter( + query=query, + num_results=num_results, + predicate=lambda result: result.url.startswith("http"), + ) diff --git a/astrbot/builtin_stars/web_searcher/engines/google.py b/astrbot/builtin_stars/web_searcher/engines/google.py new file mode 100644 index 000000000..96a3424e3 --- /dev/null +++ b/astrbot/builtin_stars/web_searcher/engines/google.py @@ -0,0 +1,52 @@ +import urllib.parse +from typing import cast + +from bs4 import Tag + +from . import SearchEngine, SearchResult + + +class Google(SearchEngine): + NAME = "google" + + def __init__(self) -> None: + super().__init__() + self.base_url = "https://www.google.com" + + def _set_selector(self, selector: str): + selectors = { + "url": "a[href]", + "title": "h3", + "text": "div.VwiC3b, span.aCOpRe", + "links": "div#search div.g, div#search div.MjjYud", + "next": "a#pnnext", + } + return selectors[selector] + + async def _get_next_page(self, query: str) -> str: + params = { + "q": urllib.parse.unquote(query), + "hl": "en", + "gl": "us", + "pws": "0", + "num": "10", + } + url = f"{self.base_url}/search?{urllib.parse.urlencode(params)}" + return await self._get_html(url, None) + + def _get_url(self, tag: Tag) -> str: + href = cast(str, tag.get("href") or "") + if href.startswith("/url?"): + parsed = urllib.parse.urlparse(href) + q = urllib.parse.parse_qs(parsed.query).get("q", [""])[0] + return urllib.parse.unquote(q) + return href + + async def search(self, query: str, num_results: int) -> list[SearchResult]: + return await self._search_with_result_filter( + query=query, + num_results=num_results, + predicate=lambda result: ( + result.url.startswith("http") and "google.com/search?" not in result.url + ), + ) diff --git a/astrbot/builtin_stars/web_searcher/engines/sogo.py b/astrbot/builtin_stars/web_searcher/engines/sogo.py index f490f1106..b5026c9e8 100644 --- a/astrbot/builtin_stars/web_searcher/engines/sogo.py +++ b/astrbot/builtin_stars/web_searcher/engines/sogo.py @@ -8,6 +8,8 @@ from . import USER_AGENTS, SearchEngine, SearchResult class Sogo(SearchEngine): + NAME = "sogo" + def __init__(self) -> None: super().__init__() self.base_url = "https://www.sogou.com" diff --git a/astrbot/builtin_stars/web_searcher/main.py b/astrbot/builtin_stars/web_searcher/main.py index 077f78792..72f24962b 100644 --- a/astrbot/builtin_stars/web_searcher/main.py +++ b/astrbot/builtin_stars/web_searcher/main.py @@ -15,7 +15,17 @@ from astrbot.core.provider.func_tool_manager import FunctionToolManager from .engines import HEADERS, USER_AGENTS, SearchResult from .engines.bing import Bing +from .engines.comet import Comet +from .engines.duckduckgo import DuckDuckGo +from .engines.google import Google from .engines.sogo import Sogo +from .provider_routing import ( + DEFAULT_WEB_SEARCH_PROVIDER, + build_default_engine_order, + normalize_websearch_provider, + normalize_websearch_provider_for_tools, + validate_default_engine_registry, +) class Main(star.Star): @@ -58,8 +68,22 @@ class Main(star.Star): provider_settings["websearch_bocha_key"] = [] cfg.save_config() + self.google_search = Google() self.bing_search = Bing() + self.ddg_search = DuckDuckGo() + self.comet_search = Comet() self.sogo_search = Sogo() + self.default_search_engines = { + engine.NAME: engine + for engine in ( + self.google_search, + self.bing_search, + self.ddg_search, + self.comet_search, + self.sogo_search, + ) + } + validate_default_engine_registry(self.default_search_engines) self.baidu_initialized = False async def _tidy_text(self, text: str) -> str: @@ -106,23 +130,29 @@ class Main(star.Star): self, query, num_results: int = 5, + preferred_provider: str = DEFAULT_WEB_SEARCH_PROVIDER, ) -> list[SearchResult]: - results = [] - try: - results = await self.bing_search.search(query, num_results) - except Exception as e: - logger.error(f"bing search error: {e}, try the next one...") - if len(results) == 0: - logger.debug("search bing failed") + for engine_name in build_default_engine_order(preferred_provider): + engine = self.default_search_engines.get(engine_name) + if not engine: + continue try: - results = await self.sogo_search.search(query, num_results) + results = await engine.search(query, num_results) except Exception as e: - logger.error(f"sogo search error: {e}") - if len(results) == 0: - logger.debug("search sogo failed") - return [] + logger.error( + f"{engine_name} search error: {e}, try the next one...", + ) + continue - return results + if results: + logger.info( + f"web_searcher - provider `{engine_name}` success: {len(results)} results", + ) + return results + + logger.debug(f"search {engine_name} returned no results") + + return [] async def _get_tavily_key(self, cfg: AstrBotConfig) -> str: """并发安全的从列表中获取并轮换Tavily API密钥｡""" @@ -214,8 +244,17 @@ class Main(star.Star): logger.info(f"web_searcher - search_from_search_engine: {query}") cfg = self.context.get_config(umo=event.unified_msg_origin) websearch_link = cfg["provider_settings"].get("web_search_link", False) - - results = await self._web_search_default(query, max_results) + preferred_provider = normalize_websearch_provider( + cfg.get("provider_settings", {}).get( + "websearch_provider", + DEFAULT_WEB_SEARCH_PROVIDER, + ), + ) + results = await self._web_search_default( + query, + max_results, + preferred_provider=preferred_provider, + ) if not results: return "Error: web searcher does not return any results." @@ -548,7 +587,13 @@ class Main(star.Star): cfg = self.context.get_config(umo=event.unified_msg_origin) prov_settings = cfg.get("provider_settings", {}) websearch_enable = prov_settings.get("web_search", False) - provider = prov_settings.get("websearch_provider", "default") + raw_provider = prov_settings.get( + "websearch_provider", + DEFAULT_WEB_SEARCH_PROVIDER, + ) + branch_provider, is_known_provider = normalize_websearch_provider_for_tools( + raw_provider + ) tool_set = req.func_tool if isinstance(tool_set, FunctionToolManager): @@ -565,7 +610,12 @@ class Main(star.Star): return func_tool_mgr = self.context.get_llm_tool_manager() - if provider == "default": + if branch_provider == "default": + if not is_known_provider: + logger.warning( + "Unsupported websearch_provider `%s`, fallback to default search tool branch.", + raw_provider, + ) web_search_t = func_tool_mgr.get_func("web_search") fetch_url_t = func_tool_mgr.get_func("fetch_url") if web_search_t and web_search_t.active: @@ -576,7 +626,7 @@ class Main(star.Star): tool_set.remove_tool("tavily_extract_web_page") tool_set.remove_tool("AIsearch") tool_set.remove_tool("web_search_bocha") - elif provider == "tavily": + elif branch_provider == "tavily": web_search_tavily = func_tool_mgr.get_func("web_search_tavily") tavily_extract_web_page = func_tool_mgr.get_func("tavily_extract_web_page") if web_search_tavily and web_search_tavily.active: @@ -587,7 +637,7 @@ class Main(star.Star): tool_set.remove_tool("fetch_url") tool_set.remove_tool("AIsearch") tool_set.remove_tool("web_search_bocha") - elif provider == "baidu_ai_search": + elif branch_provider == "baidu_ai_search": try: await self.ensure_baidu_ai_search_mcp(event.unified_msg_origin) aisearch_tool = func_tool_mgr.get_func("AIsearch") @@ -600,7 +650,7 @@ class Main(star.Star): tool_set.remove_tool("web_search_bocha") except Exception as e: logger.error(f"Cannot Initialize Baidu AI Search MCP Server: {e}") - elif provider == "bocha": + elif branch_provider == "bocha": web_search_bocha = func_tool_mgr.get_func("web_search_bocha") if web_search_bocha and web_search_bocha.active: tool_set.add_tool(web_search_bocha) # type: ignore[arg-type] diff --git a/astrbot/builtin_stars/web_searcher/provider_constants.py b/astrbot/builtin_stars/web_searcher/provider_constants.py new file mode 100644 index 000000000..249716be6 --- /dev/null +++ b/astrbot/builtin_stars/web_searcher/provider_constants.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +DEFAULT_WEB_SEARCH_PROVIDER = "default" + +# Canonical provider ids shown in config UI options. +WEB_SEARCH_PROVIDER_OPTIONS: tuple[str, ...] = ( + DEFAULT_WEB_SEARCH_PROVIDER, + "duckduckgo", + "google", + "bing", + "comet", + "sogo", + "tavily", + "baidu_ai_search", + "bocha", +) + +# Provider ids that select non-default tool branches directly. +WEB_SEARCH_TOOL_BRANCH_PROVIDERS: tuple[str, ...] = ( + DEFAULT_WEB_SEARCH_PROVIDER, + "tavily", + "baidu_ai_search", + "bocha", +) diff --git a/astrbot/builtin_stars/web_searcher/provider_routing.py b/astrbot/builtin_stars/web_searcher/provider_routing.py new file mode 100644 index 000000000..744cf43f2 --- /dev/null +++ b/astrbot/builtin_stars/web_searcher/provider_routing.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass + +from .engines.bing import Bing +from .engines.comet import Comet +from .engines.duckduckgo import DuckDuckGo +from .engines.google import Google +from .engines.sogo import Sogo +from .provider_constants import ( + DEFAULT_WEB_SEARCH_PROVIDER, + WEB_SEARCH_PROVIDER_OPTIONS, + WEB_SEARCH_TOOL_BRANCH_PROVIDERS, +) + +ENGINE_REGISTRY: tuple[tuple[str, type[object], bool], ...] = ( + (Bing.NAME, Bing, True), + (Sogo.NAME, Sogo, True), + # Compatibility first: DDG should stay as fallback and cannot become primary. + (DuckDuckGo.NAME, DuckDuckGo, False), + (Google.NAME, Google, True), + (Comet.NAME, Comet, True), +) + +DEFAULT_ENGINE_ORDER: tuple[str, ...] = tuple(name for name, _, _ in ENGINE_REGISTRY) + +_ENGINE_PROVIDER_SET = {name for name, _, _ in ENGINE_REGISTRY} +_ENGINE_CAN_BE_PRIMARY = { + name: can_be_primary for name, _, can_be_primary in ENGINE_REGISTRY +} +_TOOL_BRANCH_PROVIDER_SET = set(WEB_SEARCH_TOOL_BRANCH_PROVIDERS) +_CANONICAL_PROVIDER_SET = _ENGINE_PROVIDER_SET | _TOOL_BRANCH_PROVIDER_SET + +if not _CANONICAL_PROVIDER_SET.issubset(set(WEB_SEARCH_PROVIDER_OPTIONS)): + raise RuntimeError( + "web search provider options and routing providers are out of sync: " + f"canonical={sorted(_CANONICAL_PROVIDER_SET)} options={list(WEB_SEARCH_PROVIDER_OPTIONS)}", + ) + +_WEB_SEARCH_PROVIDER_ALIASES: dict[str, str] = { + "": DEFAULT_WEB_SEARCH_PROVIDER, + "default": DEFAULT_WEB_SEARCH_PROVIDER, + "native": DEFAULT_WEB_SEARCH_PROVIDER, +} +_WEB_SEARCH_PROVIDER_ALIASES.update({name: name for name in _CANONICAL_PROVIDER_SET}) +_WEB_SEARCH_PROVIDER_ALIASES.update( + { + "duckduck_go": DuckDuckGo.NAME, + "duckduck-go": DuckDuckGo.NAME, + "ddg": DuckDuckGo.NAME, + "baidu_ai": "baidu_ai_search", + "baidu": "baidu_ai_search", + "bochaai": "bocha", + # ZeroClaw compatibility: AstrBot has no Brave provider yet, so downgrade to default. + "brave": DEFAULT_WEB_SEARCH_PROVIDER, + } +) + + +@dataclass(frozen=True) +class NormalizedProvider: + canonical: str + tool_branch: str + is_known: bool + + +def _normalize_raw_provider(provider: object) -> str: + return str(provider or "").strip().lower().replace(" ", "") + + +def normalize_websearch(provider: object) -> NormalizedProvider: + raw = _normalize_raw_provider(provider) + alias = _WEB_SEARCH_PROVIDER_ALIASES.get(raw, raw) + canonical = alias or DEFAULT_WEB_SEARCH_PROVIDER + + is_engine = canonical in _ENGINE_PROVIDER_SET + is_tool_branch = canonical in _TOOL_BRANCH_PROVIDER_SET + is_known = is_engine or is_tool_branch + tool_branch = canonical if is_tool_branch else DEFAULT_WEB_SEARCH_PROVIDER + + return NormalizedProvider( + canonical=canonical, + tool_branch=tool_branch, + is_known=is_known, + ) + + +def normalize_websearch_provider(provider: object) -> str: + return normalize_websearch(provider).canonical + + +def normalize_websearch_provider_for_tools(provider: object) -> tuple[str, bool]: + normalized = normalize_websearch(provider) + return normalized.tool_branch, normalized.is_known + + +def resolve_tool_branch_provider(provider: object) -> str: + return normalize_websearch(provider).tool_branch + + +def build_default_engine_order(provider: object) -> tuple[str, ...]: + normalized = normalize_websearch(provider) + engine_name = normalized.canonical + + if engine_name not in _ENGINE_PROVIDER_SET: + return DEFAULT_ENGINE_ORDER + + if not _ENGINE_CAN_BE_PRIMARY.get(engine_name, False): + return DEFAULT_ENGINE_ORDER + + return ( + engine_name, + *tuple(name for name in DEFAULT_ENGINE_ORDER if name != engine_name), + ) + + +def is_known_websearch_provider(provider: object) -> bool: + return normalize_websearch(provider).is_known + + +def validate_default_engine_registry(engines_by_name: Mapping[str, object]) -> None: + expected_names = {name for name, _, _ in ENGINE_REGISTRY} + missing = [name for name in DEFAULT_ENGINE_ORDER if name not in engines_by_name] + extra = [name for name in engines_by_name if name not in expected_names] + if not missing and not extra: + return + + raise ValueError( + "default search engine registry mismatch. " + f"missing={missing}, extra={extra}, expected_order={list(DEFAULT_ENGINE_ORDER)}", + ) diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py index 6debeb1a3..c00dc563d 100644 --- a/astrbot/core/config/default.py +++ b/astrbot/core/config/default.py @@ -3,6 +3,10 @@ import os from typing import Any, TypedDict +from astrbot.builtin_stars.web_searcher.provider_constants import ( + DEFAULT_WEB_SEARCH_PROVIDER, + WEB_SEARCH_PROVIDER_OPTIONS, +) from astrbot.core.utils.astrbot_path import get_astrbot_data_path VERSION = "4.25.0" @@ -106,7 +110,7 @@ DEFAULT_CONFIG = { "provider_pool": ["*"], # "*" 表示使用所有可用的提供者 "wake_prefix": "", "web_search": False, - "websearch_provider": "default", + "websearch_provider": DEFAULT_WEB_SEARCH_PROVIDER, "websearch_tavily_key": [], "websearch_bocha_key": [], "websearch_baidu_app_builder_key": "", @@ -3080,7 +3084,7 @@ CONFIG_METADATA_3 = { "provider_settings.websearch_provider": { "description": "网页搜索提供商", "type": "string", - "options": ["default", "tavily", "baidu_ai_search", "bocha"], + "options": list(WEB_SEARCH_PROVIDER_OPTIONS), "condition": { "provider_settings.web_search": True, }, diff --git a/tests/unit/test_web_search_engines.py b/tests/unit/test_web_search_engines.py new file mode 100644 index 000000000..7217853a1 --- /dev/null +++ b/tests/unit/test_web_search_engines.py @@ -0,0 +1,173 @@ +import urllib.parse + +import pytest +from bs4 import BeautifulSoup, Tag + +from astrbot.builtin_stars.web_searcher.engines import SearchEngine, SearchResult +from astrbot.builtin_stars.web_searcher.engines.comet import Comet +from astrbot.builtin_stars.web_searcher.engines.duckduckgo import DuckDuckGo +from astrbot.builtin_stars.web_searcher.engines.google import Google + + +class EngineWithoutTextSelector(SearchEngine): + def _set_selector(self, selector: str) -> str: + selectors = { + "url": "a.title", + "title": "a.title", + "links": "div.item", + "next": "", + } + return selectors[selector] + + async def _get_next_page(self, query: str) -> str: + return """ +
+ Example Title +
+ """ + + def _get_url(self, tag: Tag) -> str: + return str(tag.get("href") or "") + + +@pytest.mark.asyncio +async def test_base_search_allows_engine_without_text_selector() -> None: + engine = EngineWithoutTextSelector() + results = await engine.search("hello world", 3) + + assert len(results) == 1 + assert results[0].title == "Example Title" + assert results[0].url == "https://example.com/a" + assert results[0].snippet == "" + + +@pytest.mark.asyncio +async def test_duckduckgo_get_next_page_urlencodes_query(monkeypatch: pytest.MonkeyPatch) -> None: + engine = DuckDuckGo() + captured: dict[str, str] = {} + + async def fake_get_html(url: str, data: dict | None = None) -> str: + captured["url"] = url + return "" + + monkeypatch.setattr(engine, "_get_html", fake_get_html) + await engine._get_next_page("hello%20world%2Bv2") + + parsed = urllib.parse.urlparse(captured["url"]) + params = urllib.parse.parse_qs(parsed.query) + assert params["q"] == ["hello world+v2"] + assert params["kl"] == ["us-en"] + + +@pytest.mark.asyncio +async def test_comet_get_next_page_urlencodes_query(monkeypatch: pytest.MonkeyPatch) -> None: + engine = Comet() + captured: dict[str, str] = {} + + async def fake_get_html(url: str, data: dict | None = None) -> str: + captured["url"] = url + return "" + + monkeypatch.setattr(engine, "_get_html", fake_get_html) + await engine._get_next_page("astrbot%20rtk%20scrapling%2Btest") + + parsed = urllib.parse.urlparse(captured["url"]) + params = urllib.parse.parse_qs(parsed.query) + assert params["q"] == ["astrbot rtk scrapling+test"] + + +@pytest.mark.asyncio +async def test_google_get_next_page_urlencodes_query(monkeypatch: pytest.MonkeyPatch) -> None: + engine = Google() + captured: dict[str, str] = {} + + async def fake_get_html(url: str, data: dict | None = None) -> str: + captured["url"] = url + return "" + + monkeypatch.setattr(engine, "_get_html", fake_get_html) + await engine._get_next_page("hello%20world%2Bv2") + + parsed = urllib.parse.urlparse(captured["url"]) + params = urllib.parse.parse_qs(parsed.query) + assert params["q"] == ["hello world+v2"] + assert params["hl"] == ["en"] + assert params["gl"] == ["us"] + + +def test_comet_get_url_prefers_href_and_supports_scheme_relative() -> None: + engine = Comet() + http_tag = BeautifulSoup( + 'Title', + "html.parser", + ).find("a") + rel_tag = BeautifulSoup('Title', "html.parser").find( + "a" + ) + assert http_tag is not None + assert rel_tag is not None + + assert engine._get_url(http_tag) == "https://example.com/result" + assert engine._get_url(rel_tag) == "https://example.org/path" + + +@pytest.mark.asyncio +async def test_comet_search_filters_internal_and_non_content_links( + monkeypatch: pytest.MonkeyPatch, +) -> None: + engine = Comet() + + async def fake_search( + self: SearchEngine, + query: str, + num_results: int, + ) -> list[SearchResult]: + return [ + SearchResult(title="internal", url="https://www.perplexity.ai/search?q=a", snippet=""), + SearchResult(title="anchor", url="#section", snippet=""), + SearchResult(title="js", url="javascript:void(0)", snippet=""), + SearchResult(title="mailto", url="mailto:a@example.com", snippet=""), + SearchResult(title="ok1", url="https://example.com/a", snippet=""), + SearchResult(title="ok2", url="http://example.org/b", snippet=""), + ] + + monkeypatch.setattr(SearchEngine, "search", fake_search) + + results = await engine.search("hello", 10) + assert [r.url for r in results] == [ + "https://example.com/a", + "http://example.org/b", + ] + + +@pytest.mark.asyncio +async def test_search_with_result_filter_skips_fetch_on_non_positive_num_results( + monkeypatch: pytest.MonkeyPatch, +) -> None: + engine = EngineWithoutTextSelector() + called = {"value": False} + + async def fake_search( + self: SearchEngine, + query: str, + num_results: int, + ) -> list[SearchResult]: + called["value"] = True + return [SearchResult(title="t", url="https://example.com", snippet="s")] + + monkeypatch.setattr(SearchEngine, "search", fake_search) + + results_zero = await engine._search_with_result_filter( + query="hello", + num_results=0, + predicate=lambda _: True, + ) + results_negative = await engine._search_with_result_filter( + query="hello", + num_results=-3, + predicate=lambda _: True, + ) + + assert results_zero == [] + assert results_negative == [] + assert called["value"] is False diff --git a/tests/unit/test_web_search_provider_routing.py b/tests/unit/test_web_search_provider_routing.py new file mode 100644 index 000000000..4292af26f --- /dev/null +++ b/tests/unit/test_web_search_provider_routing.py @@ -0,0 +1,110 @@ +import pytest + +from astrbot.builtin_stars.web_searcher.provider_routing import ( + DEFAULT_ENGINE_ORDER, + build_default_engine_order, + normalize_websearch, + normalize_websearch_provider, + normalize_websearch_provider_for_tools, + resolve_tool_branch_provider, + validate_default_engine_registry, +) + + +def test_normalize_websearch_provider_aliases() -> None: + assert normalize_websearch_provider("duckduckgo") == "duckduckgo" + assert normalize_websearch_provider("ddg") == "duckduckgo" + assert normalize_websearch_provider("duckduck-go") == "duckduckgo" + assert normalize_websearch_provider("baidu") == "baidu_ai_search" + assert normalize_websearch_provider("bochaai") == "bocha" + assert normalize_websearch_provider("brave") == "default" + + +def test_normalize_websearch_returns_unified_descriptor() -> None: + ddg = normalize_websearch("ddg") + assert ddg.canonical == "duckduckgo" + assert ddg.tool_branch == "default" + assert ddg.is_known is True + + tavily = normalize_websearch("tavily") + assert tavily.canonical == "tavily" + assert tavily.tool_branch == "tavily" + assert tavily.is_known is True + + unknown = normalize_websearch("unknown_provider") + assert unknown.canonical == "unknown_provider" + assert unknown.tool_branch == "default" + assert unknown.is_known is False + + +def test_resolve_tool_branch_provider_uses_default_branch_for_engine_aliases() -> None: + assert resolve_tool_branch_provider("duckduckgo") == "default" + assert resolve_tool_branch_provider("google") == "default" + assert resolve_tool_branch_provider("ddg") == "default" + assert resolve_tool_branch_provider("tavily") == "tavily" + assert resolve_tool_branch_provider("baidu_ai_search") == "baidu_ai_search" + assert resolve_tool_branch_provider("bocha") == "bocha" + # Unknown provider should fall back to default branch instead of leaving mixed tool set. + assert resolve_tool_branch_provider("unknown_provider") == "default" + + +def test_build_default_engine_order_keeps_dev_compatible_default_chain() -> None: + assert DEFAULT_ENGINE_ORDER[:2] == ("bing", "sogo") + + order = build_default_engine_order("duckduckgo") + assert order == DEFAULT_ENGINE_ORDER + assert order[0] == "bing" + + order = build_default_engine_order("bing") + assert order[0] == "bing" + assert set(order) == set(DEFAULT_ENGINE_ORDER) + + order = build_default_engine_order("google") + assert order[0] == "google" + assert set(order) == set(DEFAULT_ENGINE_ORDER) + + order = build_default_engine_order("tavily") + assert order == DEFAULT_ENGINE_ORDER + + +def test_build_default_engine_order_covers_all_engines_unknown_and_no_duplicates() -> None: + order = build_default_engine_order("sogo") + assert order[0] == "sogo" + assert set(order) == set(DEFAULT_ENGINE_ORDER) + + order = build_default_engine_order("comet") + assert order[0] == "comet" + assert set(order) == set(DEFAULT_ENGINE_ORDER) + + order = build_default_engine_order("xxx") + assert order == DEFAULT_ENGINE_ORDER + + order = build_default_engine_order("bing") + assert len(order) == len(set(order)) + + +def test_normalize_websearch_provider_for_tools_returns_branch_and_known() -> None: + assert normalize_websearch_provider_for_tools("tavily") == ("tavily", True) + assert normalize_websearch_provider_for_tools("ddg") == ("default", True) + assert normalize_websearch_provider_for_tools("unknown_provider") == ( + "default", + False, + ) + + +def test_validate_default_engine_registry_allows_expected_names() -> None: + engines_by_name = {name: object() for name in DEFAULT_ENGINE_ORDER} + validate_default_engine_registry(engines_by_name) + + +def test_validate_default_engine_registry_rejects_missing_or_extra() -> None: + missing_last = DEFAULT_ENGINE_ORDER[:-1] + with_missing = {name: object() for name in missing_last} + + with_extra = {name: object() for name in DEFAULT_ENGINE_ORDER} + with_extra["unknown_engine"] = object() + + with pytest.raises(ValueError): + validate_default_engine_registry(with_missing) + with pytest.raises(ValueError): + validate_default_engine_registry(with_extra)