Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions astrbot/core/astr_main_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,8 @@
BraveWebSearchTool,
FirecrawlExtractWebPageTool,
FirecrawlWebSearchTool,
KeenableExtractWebPageTool,
KeenableWebSearchTool,
TavilyExtractWebPageTool,
TavilyWebSearchTool,
normalize_legacy_web_search_config,
Expand Down Expand Up @@ -1194,6 +1196,9 @@ async def _apply_web_search_tools(
req.func_tool.add_tool(tool_mgr.get_builtin_tool(FirecrawlExtractWebPageTool))
elif provider == "baidu_ai_search":
req.func_tool.add_tool(tool_mgr.get_builtin_tool(BaiduWebSearchTool))
elif provider == "keenable":
req.func_tool.add_tool(tool_mgr.get_builtin_tool(KeenableWebSearchTool))
req.func_tool.add_tool(tool_mgr.get_builtin_tool(KeenableExtractWebPageTool))


def _apply_web_search_citation_prompt(
Expand Down
12 changes: 12 additions & 0 deletions astrbot/core/config/default.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,7 @@
"websearch_brave_key": [],
"websearch_baidu_app_builder_key": "",
"websearch_firecrawl_key": [],
"websearch_keenable_key": [],
"web_search_link": False,
"display_reasoning_text": False,
"identifier": False,
Expand Down Expand Up @@ -3284,6 +3285,7 @@
"bocha",
"brave",
"firecrawl",
"keenable",
],
"condition": {
"provider_settings.web_search": True,
Expand Down Expand Up @@ -3338,6 +3340,16 @@
"provider_settings.web_search": True,
},
},
"provider_settings.websearch_keenable_key": {
"description": "Keenable API Key",
"type": "list",
"items": {"type": "string"},
"hint": "可添加多个 Key 进行轮询。获取地址:https://keenable.ai",
"condition": {
"provider_settings.websearch_provider": "keenable",
"provider_settings.web_search": True,
},
},
"provider_settings.web_search_link": {
"description": "显示来源引用",
"type": "bool",
Expand Down
170 changes: 170 additions & 0 deletions astrbot/core/tools/web_search_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@
"web_search_brave",
"web_search_firecrawl",
"firecrawl_extract_web_page",
"web_search_keenable",
"keenable_extract_web_page",
]
_TAVILY_WEB_SEARCH_TOOL_CONFIG = {
"provider_settings.web_search": True,
Expand All @@ -42,6 +44,10 @@
"provider_settings.web_search": True,
"provider_settings.websearch_provider": "baidu_ai_search",
}
_KEENABLE_WEB_SEARCH_TOOL_CONFIG = {
"provider_settings.web_search": True,
"provider_settings.websearch_provider": "keenable",
}


@std_dataclass
Expand Down Expand Up @@ -76,6 +82,7 @@ async def get(self, provider_settings: dict) -> str:
_BOCHA_KEY_ROTATOR = _KeyRotator("websearch_bocha_key", "BoCha")
_BRAVE_KEY_ROTATOR = _KeyRotator("websearch_brave_key", "Brave")
_FIRECRAWL_KEY_ROTATOR = _KeyRotator("websearch_firecrawl_key", "Firecrawl")
_KEENABLE_KEY_ROTATOR = _KeyRotator("websearch_keenable_key", "Keenable")


def normalize_legacy_web_search_config(cfg) -> None:
Expand All @@ -99,6 +106,7 @@ def normalize_legacy_web_search_config(cfg) -> None:
"websearch_bocha_key",
"websearch_brave_key",
"websearch_firecrawl_key",
"websearch_keenable_key",
):
value = provider_settings.get(setting_name)
if isinstance(value, str):
Expand Down Expand Up @@ -370,6 +378,61 @@ async def _baidu_search(
]


async def _keenable_search(
provider_settings: dict,
payload: dict,
) -> list[SearchResult]:
api_key = await _KEENABLE_KEY_ROTATOR.get(provider_settings)
header = {
"X-API-Key": api_key,
"X-Keenable-Title": "astrbot",
"Content-Type": "application/json",
}
async with aiohttp.ClientSession(trust_env=True) as session:
async with session.post(
"https://api.keenable.ai/v1/search",
json=payload,
headers=header,
) as response:
if response.status != 200:
reason = await response.text()
raise Exception(
f"Keenable web search failed: {reason}, status: {response.status}",
)
data = await response.json()
return [
SearchResult(
title=item.get("title", ""),
url=item.get("url", ""),
snippet=item.get("snippet") or item.get("description") or "",
)
for item in (data.get("results") or [])
if item and item.get("url")
]
Comment on lines +403 to +411

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Using data.get("results", []) can lead to a TypeError: 'NoneType' object is not iterable if the API returns {"results": null} (which parses to None in Python). Since dict.get() only returns the default value when the key is absent, it will return None if the key is present but has a null value.

Using data.get("results") or [] is a safer and more robust pattern that handles both missing keys and null/None values gracefully. Additionally, we should check if item is not None before calling item.get("url") to prevent potential AttributeErrors.

Suggested change
return [
SearchResult(
title=item.get("title", ""),
url=item.get("url", ""),
snippet=item.get("snippet") or item.get("description") or "",
)
for item in data.get("results", [])
if item.get("url")
]
return [
SearchResult(
title=item.get("title", ""),
url=item.get("url", ""),
snippet=item.get("snippet") or item.get("description") or "",
)
for item in (data.get("results") or [])
if item and item.get("url")
]



async def _keenable_fetch(provider_settings: dict, params: dict) -> dict:
api_key = await _KEENABLE_KEY_ROTATOR.get(provider_settings)
header = {"X-API-Key": api_key, "X-Keenable-Title": "astrbot"}
async with aiohttp.ClientSession(trust_env=True) as session:
async with session.get(
"https://api.keenable.ai/v1/fetch",
params=params,
headers=header,
) as response:
if response.status != 200:
reason = await response.text()
raise Exception(
f"Keenable web fetch failed: {reason}, status: {response.status}",
)
data = await response.json()
if not data.get("content"):
raise ValueError(
"Error: Keenable web fetcher does not return any results."
)
return data


@builtin_tool(config=_TAVILY_WEB_SEARCH_TOOL_CONFIG)
@pydantic_dataclass
class TavilyWebSearchTool(FunctionTool[AstrAgentContext]):
Expand Down Expand Up @@ -803,10 +866,117 @@ async def call(self, context, **kwargs) -> ToolExecResult:
return _search_result_payload(results)


@builtin_tool(config=_KEENABLE_WEB_SEARCH_TOOL_CONFIG)
@pydantic_dataclass
class KeenableWebSearchTool(FunctionTool[AstrAgentContext]):
name: str = "web_search_keenable"
description: str = (
"A web search tool based on Keenable Search API, used to retrieve web "
"pages related to the user's query."
)
parameters: dict = Field(
default_factory=lambda: {
"type": "object",
"properties": {
"query": {"type": "string", "description": "Required. Search query."},
"site": {
"type": "string",
"description": 'Optional. Restrict results to a specific site, for example "techcrunch.com".',
},
"published_after": {
"type": "string",
"description": 'Optional. Only include pages published at or after this time. Accepts "YYYY-MM-DD", an ISO 8601 datetime, or a relative delta like "7d", "3mo".',
},
"published_before": {
"type": "string",
"description": "Optional. Only include pages published at or before this time. Same formats as published_after.",
},
"acquired_after": {
"type": "string",
"description": "Optional. Only include pages indexed at or after this time. Same formats as published_after.",
},
"acquired_before": {
"type": "string",
"description": "Optional. Only include pages indexed at or before this time. Same formats as published_after.",
},
},
"required": ["query"],
}
)

async def call(self, context, **kwargs) -> ToolExecResult:
_, provider_settings, _ = _get_runtime(context)
if not provider_settings.get("websearch_keenable_key", []):
return "Error: Keenable API key is not configured in AstrBot."

payload = {"query": kwargs["query"]}
for key in (
"site",
"published_after",
"published_before",
"acquired_after",
"acquired_before",
):
if kwargs.get(key):
payload[key] = kwargs[key]

results = await _keenable_search(provider_settings, payload)
if not results:
return "Error: Keenable web searcher does not return any results."
return _search_result_payload(results)


@builtin_tool(config=_KEENABLE_WEB_SEARCH_TOOL_CONFIG)
@pydantic_dataclass
class KeenableExtractWebPageTool(FunctionTool[AstrAgentContext]):
name: str = "keenable_extract_web_page"
description: str = (
"Extract the content of a web page using Keenable. "
"Only URLs indexed by Keenable are supported."
)
parameters: dict = Field(
default_factory=lambda: {
"type": "object",
"properties": {
"url": {
"type": "string",
"description": "Required. A URL to extract content from.",
},
"max_chars": {
"type": "integer",
"description": "Optional. Maximum number of characters of content to return. Default is 50000.",
},
},
"required": ["url"],
}
)

async def call(self, context, **kwargs) -> ToolExecResult:
_, provider_settings, _ = _get_runtime(context)
if not provider_settings.get("websearch_keenable_key", []):
return "Error: Keenable API key is not configured in AstrBot."

url = str(kwargs.get("url", "")).strip()
if not url:
return "Error: url must be a non-empty string."

params = {"url": url}
if kwargs.get("max_chars"):
params["max_chars"] = kwargs["max_chars"]

result = await _keenable_fetch(provider_settings, params)
content = result.get("content", "")
result_url = result.get("url") or url
ret = f"URL: {result_url}\nContent: {content}" if content else ""
return ret or "Error: Keenable web fetcher does not return any results."


__all__ = [
"BaiduWebSearchTool",
"BochaWebSearchTool",
"BraveWebSearchTool",
"KeenableExtractWebPageTool",
"KeenableWebSearchTool",
"TavilyExtractWebPageTool",
"TavilyWebSearchTool",
"WEB_SEARCH_TOOL_NAMES",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,10 @@
"description": "Baidu Qianfan Smart Cloud APP Builder API Key",
"hint": "Reference: [https://console.bce.baidu.com/iam/#/iam/apikey/list](https://console.bce.baidu.com/iam/#/iam/apikey/list)"
},
"websearch_keenable_key": {
"description": "Keenable API Key",
"hint": "Multiple keys can be added for rotation. Get one at [https://keenable.ai](https://keenable.ai)."
},
"web_search_link": {
"description": "Display Source Citations"
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,10 @@
"description": "API-ключ Baidu Qianfan APP Builder",
"hint": "Ссылка: [https://console.bce.baidu.com/iam/#/iam/apikey/list](https://console.bce.baidu.com/iam/#/iam/apikey/list)"
},
"websearch_keenable_key": {
"description": "API-ключ Keenable",
"hint": "Можно добавить несколько ключей для ротации. Получить ключ: [https://keenable.ai](https://keenable.ai)."
},
"web_search_link": {
"description": "Показывать ссылки на источники"
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,10 @@
"description": "百度千帆智能云 APP Builder API Key",
"hint": "参考:[https://console.bce.baidu.com/iam/#/iam/apikey/list](https://console.bce.baidu.com/iam/#/iam/apikey/list)"
},
"websearch_keenable_key": {
"description": "Keenable API Key",
"hint": "可添加多个 Key 进行轮询。获取地址:[https://keenable.ai](https://keenable.ai)"
},
"web_search_link": {
"description": "显示来源引用"
}
Expand Down
31 changes: 31 additions & 0 deletions tests/unit/test_astr_main_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -510,6 +510,37 @@ async def test_apply_web_search_tools_adds_firecrawl_search_and_extract_tools(
assert req.func_tool.get_tool("web_search_firecrawl") is search_tool
assert req.func_tool.get_tool("firecrawl_extract_web_page") is extract_tool

@pytest.mark.asyncio
async def test_apply_web_search_tools_adds_keenable_search_and_extract_tools(
self, mock_event, mock_context
):
"""Test Keenable web search injects search and extract tools."""
module = ama
req = ProviderRequest()
mock_context.get_config.return_value = {
"provider_settings": {
"web_search": True,
"websearch_provider": "keenable",
}
}
search_tool = MagicMock(spec=FunctionTool)
search_tool.name = "web_search_keenable"
extract_tool = MagicMock(spec=FunctionTool)
extract_tool.name = "keenable_extract_web_page"
tool_mgr = MagicMock()
tool_mgr.get_builtin_tool.side_effect = [search_tool, extract_tool]
mock_context.get_llm_tool_manager.return_value = tool_mgr

await module._apply_web_search_tools(mock_event, req, mock_context)

assert tool_mgr.get_builtin_tool.call_args_list == [
((module.KeenableWebSearchTool,),),
((module.KeenableExtractWebPageTool,),),
]
assert req.func_tool is not None
assert req.func_tool.get_tool("web_search_keenable") is search_tool
assert req.func_tool.get_tool("keenable_extract_web_page") is extract_tool

def test_apply_web_search_citation_prompt_for_webchat(self, mock_event):
module = ama
req = ProviderRequest(system_prompt="base")
Expand Down
14 changes: 14 additions & 0 deletions tests/unit/test_func_tool_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
from astrbot.core.tools.web_search_tools import (
FirecrawlExtractWebPageTool,
FirecrawlWebSearchTool,
KeenableExtractWebPageTool,
KeenableWebSearchTool,
)


Expand Down Expand Up @@ -345,3 +347,15 @@ def test_firecrawl_tools_are_registered_as_builtin_tools():
assert extract_tool.name == "firecrawl_extract_web_page"
assert manager.is_builtin_tool("web_search_firecrawl") is True
assert manager.is_builtin_tool("firecrawl_extract_web_page") is True


def test_keenable_tools_are_registered_as_builtin_tools():
manager = FunctionToolManager()

search_tool = manager.get_builtin_tool(KeenableWebSearchTool)
extract_tool = manager.get_builtin_tool(KeenableExtractWebPageTool)

assert search_tool.name == "web_search_keenable"
assert extract_tool.name == "keenable_extract_web_page"
assert manager.is_builtin_tool("web_search_keenable") is True
assert manager.is_builtin_tool("keenable_extract_web_page") is True
Loading