mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
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
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
68
astrbot/builtin_stars/web_searcher/engines/comet.py
Normal file
68
astrbot/builtin_stars/web_searcher/engines/comet.py
Normal file
@@ -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),
|
||||
)
|
||||
44
astrbot/builtin_stars/web_searcher/engines/duckduckgo.py
Normal file
44
astrbot/builtin_stars/web_searcher/engines/duckduckgo.py
Normal file
@@ -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"),
|
||||
)
|
||||
52
astrbot/builtin_stars/web_searcher/engines/google.py
Normal file
52
astrbot/builtin_stars/web_searcher/engines/google.py
Normal file
@@ -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
|
||||
),
|
||||
)
|
||||
@@ -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"
|
||||
|
||||
@@ -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]
|
||||
|
||||
24
astrbot/builtin_stars/web_searcher/provider_constants.py
Normal file
24
astrbot/builtin_stars/web_searcher/provider_constants.py
Normal file
@@ -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",
|
||||
)
|
||||
132
astrbot/builtin_stars/web_searcher/provider_routing.py
Normal file
132
astrbot/builtin_stars/web_searcher/provider_routing.py
Normal file
@@ -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)}",
|
||||
)
|
||||
@@ -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,
|
||||
},
|
||||
|
||||
173
tests/unit/test_web_search_engines.py
Normal file
173
tests/unit/test_web_search_engines.py
Normal file
@@ -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 """
|
||||
<div class="item">
|
||||
<a class="title" href="https://example.com/a">Example Title</a>
|
||||
</div>
|
||||
"""
|
||||
|
||||
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(
|
||||
'<a href="https://example.com/result">Title</a>',
|
||||
"html.parser",
|
||||
).find("a")
|
||||
rel_tag = BeautifulSoup('<a href="//example.org/path">Title</a>', "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
|
||||
110
tests/unit/test_web_search_provider_routing.py
Normal file
110
tests/unit/test_web_search_provider_routing.py
Normal file
@@ -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)
|
||||
Reference in New Issue
Block a user