diff --git a/astrbot/_internal/protocols/mcp/tool.py b/astrbot/_internal/protocols/mcp/tool.py index 5e9eda1db..39e8646cf 100644 --- a/astrbot/_internal/protocols/mcp/tool.py +++ b/astrbot/_internal/protocols/mcp/tool.py @@ -6,11 +6,12 @@ from typing import TYPE_CHECKING, Any try: import mcp except (ModuleNotFoundError, ImportError): - mcp = None # type: ignore + mcp: Any = None from astrbot._internal.tools.base import FunctionTool if TYPE_CHECKING: + from mcp.types import Tool as MCPTool_T from astrbot._internal.protocols.mcp.client import McpClient @@ -19,8 +20,8 @@ class MCPTool(FunctionTool): def __init__( self, - mcp_tool: "mcp.types.Tool", - mcp_client: "McpClient", + mcp_tool: MCPTool_T, + mcp_client: McpClient, mcp_server_name: str, **kwargs: Any, ) -> None: diff --git a/astrbot/_internal/tools/base.py b/astrbot/_internal/tools/base.py index ce29ca9aa..deca00a49 100644 --- a/astrbot/_internal/tools/base.py +++ b/astrbot/_internal/tools/base.py @@ -27,6 +27,9 @@ class ToolSchema: parameters: ParametersType = field(default_factory=dict) """The parameters of the tool, in JSON Schema format.""" + active: bool = True + """Whether the tool is active.""" + @model_validator(mode="after") def validate_parameters(self) -> "ToolSchema": """Validate the parameters JSON schema.""" @@ -54,12 +57,6 @@ class FunctionTool(ToolSchema): causing the handler's __module__ to be functools """ - active: bool = True - """ - Whether the tool is active. This field is a special field for AstrBot. - You can ignore it when integrating with other frameworks. - """ - is_background_task: bool = False """ Declare this tool as a background task. Background tasks return immediately @@ -90,22 +87,22 @@ class ToolSet: as "namespace/tool_name" when calling. """ - def __init__(self, namespace: str, tools: list[FunctionTool] | None = None) -> None: + def __init__(self, namespace: str, tools: list[ToolSchema] | None = None) -> None: self.namespace = namespace - self._tools: dict[str, FunctionTool] = {} + self._tools: dict[str, ToolSchema] = {} if tools: for tool in tools: self.add(tool) - def add(self, tool: FunctionTool) -> None: + def add(self, tool: ToolSchema) -> None: """Add a tool to the set.""" self._tools[tool.name] = tool - def add_tool(self, tool: FunctionTool) -> None: + def add_tool(self, tool: ToolSchema) -> None: """Add a tool to the set (alias for add()).""" self.add(tool) - def remove(self, name: str) -> FunctionTool | None: + def remove(self, name: str) -> ToolSchema | None: """Remove and return a tool by name.""" return self._tools.pop(name, None) @@ -113,19 +110,19 @@ class ToolSet: """Remove a tool by its name.""" self._tools.pop(name, None) - def get(self, name: str) -> FunctionTool | None: + def get(self, name: str) -> ToolSchema | None: """Get a tool by name.""" return self._tools.get(name) - def get_tool(self, name: str) -> FunctionTool | None: + def get_tool(self, name: str) -> ToolSchema | None: """Get a tool by name (alias for get).""" return self.get(name) - def list_tools(self) -> list[FunctionTool]: + def list_tools(self) -> list[ToolSchema]: """List all tools in this set.""" return list(self._tools.values()) - def __iter__(self) -> Iterator[FunctionTool]: + def __iter__(self) -> Iterator[ToolSchema]: return iter(self._tools.values()) def __len__(self) -> int: @@ -195,7 +192,7 @@ class ToolSet: return ToolSet("default", param_tools) @property - def tools(self) -> list[FunctionTool]: + def tools(self) -> list[ToolSchema]: """List all tools in this set.""" return list(self._tools.values()) @@ -205,18 +202,20 @@ class ToolSet: """Convert tools to OpenAI API function calling schema format.""" result: list[dict[str, Any]] = [] for tool in self._tools.values(): - func_def: dict[str, Any] = { - "type": "function", - "function": {"name": tool.name}, - } + func_dict: dict[str, Any] = {"name": tool.name} if tool.description: - func_def["function"]["description"] = tool.description + func_dict["description"] = tool.description if tool.parameters is not None: if ( tool.parameters.get("properties") ) or not omit_empty_parameter_field: - func_def["function"]["parameters"] = tool.parameters + func_dict["parameters"] = tool.parameters + + func_def: dict[str, Any] = { + "type": "function", + "function": func_dict, + } result.append(func_def) return result diff --git a/astrbot/_internal/tools/registry.py b/astrbot/_internal/tools/registry.py index 314f25739..d6e29d640 100644 --- a/astrbot/_internal/tools/registry.py +++ b/astrbot/_internal/tools/registry.py @@ -5,7 +5,7 @@ from __future__ import annotations from typing import Any # Re-export from base -from astrbot._internal.tools.base import FunctionTool, ToolSet +from astrbot._internal.tools.base import FunctionTool, ToolSchema, ToolSet __all__ = [ "DEFAULT_MCP_CONFIG", @@ -48,19 +48,19 @@ class FunctionToolManager: """Central registry for all function tools.""" def __init__(self) -> None: - self._func_list: list[FunctionTool] = [] + self._func_list: list[ToolSchema] = [] @property - def func_list(self) -> list[FunctionTool]: + def func_list(self) -> list[ToolSchema]: """Get the list of function tools.""" return self._func_list @func_list.setter - def func_list(self, value: list[FunctionTool]) -> None: + def func_list(self, value: list[ToolSchema]) -> None: """Set the list of function tools.""" self._func_list = value - def add(self, tool: FunctionTool) -> None: + def add(self, tool: ToolSchema) -> None: """Add a tool to the registry.""" self._func_list.append(tool) @@ -72,9 +72,9 @@ class FunctionToolManager: return True return False - def get_func(self, name: str) -> FunctionTool | None: + def get_func(self, name: str) -> ToolSchema | None: """Get a tool by name. Returns the last active tool if multiple match.""" - last_match: FunctionTool | None = None + last_match: ToolSchema | None = None for f in reversed(self._func_list): if f.name == name: if getattr(f, "active", True): @@ -85,7 +85,7 @@ class FunctionToolManager: def get_full_tool_set(self) -> ToolSet: """Return a ToolSet with all active tools, deduplicated by name.""" - seen: dict[str, FunctionTool] = {} + seen: dict[str, ToolSchema] = {} for tool in reversed(self._func_list): if tool.name not in seen and getattr(tool, "active", True): seen[tool.name] = tool @@ -98,7 +98,7 @@ class FunctionToolManager: for tool in get_all_tools(): if self.get_func(tool.name) is None: - self.add(tool) + self.add(tool) # type: ignore[arg-type] # MCP-related stub methods for base class compatibility async def enable_mcp_server( @@ -205,7 +205,7 @@ class FuncCall(FunctionToolManager): """Remove a function tool by name (deprecated, use remove() instead).""" self.remove(name) - def get_func(self, name: str) -> FunctionTool | None: + def get_func(self, name: str) -> ToolSchema | None: """Get a function tool by name.""" return super().get_func(name) diff --git a/astrbot/api/tools.py b/astrbot/api/tools.py index b3ff77a38..b4af23b07 100644 --- a/astrbot/api/tools.py +++ b/astrbot/api/tools.py @@ -22,10 +22,10 @@ from functools import wraps from typing import Any # Import from _internal package (the canonical source) -from astrbot._internal.tools.base import FunctionTool, ToolSet +from astrbot._internal.tools.base import FunctionTool, ToolSchema, ToolSet from astrbot._internal.tools.registry import FunctionToolManager -__all__ = ["FunctionTool", "ToolRegistry", "ToolSet", "get_registry", "tool"] +__all__ = ["FunctionTool", "ToolRegistry", "ToolSet", "get_registry", "tool", "ToolSchema"] class ToolRegistry: @@ -62,11 +62,11 @@ class ToolRegistry: return True return False - def list_tools(self) -> list[FunctionTool]: + def list_tools(self) -> list[ToolSchema]: """List all registered tools.""" return self._manager.func_list.copy() - def get_tool(self, name: str) -> FunctionTool | None: + def get_tool(self, name: str) -> ToolSchema | None: """Get a tool by name.""" return self._manager.get_func(name) diff --git a/astrbot/builtin_stars/web_searcher/main.py b/astrbot/builtin_stars/web_searcher/main.py index 5de9424ba..077f78792 100644 --- a/astrbot/builtin_stars/web_searcher/main.py +++ b/astrbot/builtin_stars/web_searcher/main.py @@ -552,7 +552,7 @@ class Main(star.Star): tool_set = req.func_tool if isinstance(tool_set, FunctionToolManager): - req.func_tool = tool_set.get_full_tool_set() + req.func_tool = tool_set.get_full_tool_set() # type: ignore[assignment] tool_set = req.func_tool if not tool_set: @@ -569,9 +569,9 @@ class Main(star.Star): 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: - tool_set.add_tool(web_search_t) + tool_set.add_tool(web_search_t) # type: ignore[arg-type] if fetch_url_t and fetch_url_t.active: - tool_set.add_tool(fetch_url_t) + tool_set.add_tool(fetch_url_t) # type: ignore[arg-type] tool_set.remove_tool("web_search_tavily") tool_set.remove_tool("tavily_extract_web_page") tool_set.remove_tool("AIsearch") @@ -580,9 +580,9 @@ class Main(star.Star): 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: - tool_set.add_tool(web_search_tavily) + tool_set.add_tool(web_search_tavily) # type: ignore[arg-type] if tavily_extract_web_page and tavily_extract_web_page.active: - tool_set.add_tool(tavily_extract_web_page) + tool_set.add_tool(tavily_extract_web_page) # type: ignore[arg-type] tool_set.remove_tool("web_search") tool_set.remove_tool("fetch_url") tool_set.remove_tool("AIsearch") @@ -592,7 +592,7 @@ class Main(star.Star): await self.ensure_baidu_ai_search_mcp(event.unified_msg_origin) aisearch_tool = func_tool_mgr.get_func("AIsearch") if aisearch_tool and aisearch_tool.active: - tool_set.add_tool(aisearch_tool) + tool_set.add_tool(aisearch_tool) # type: ignore[arg-type] tool_set.remove_tool("web_search") tool_set.remove_tool("fetch_url") tool_set.remove_tool("web_search_tavily") @@ -603,7 +603,7 @@ class Main(star.Star): elif 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) + tool_set.add_tool(web_search_bocha) # type: ignore[arg-type] tool_set.remove_tool("web_search") tool_set.remove_tool("fetch_url") tool_set.remove_tool("AIsearch") diff --git a/astrbot/core/agent/mcp_client.py b/astrbot/core/agent/mcp_client.py index 03f40ecf2..0518bfa75 100644 --- a/astrbot/core/agent/mcp_client.py +++ b/astrbot/core/agent/mcp_client.py @@ -21,7 +21,7 @@ import sys import warnings from contextlib import AsyncExitStack from datetime import timedelta -from typing import Any, Generic +from typing import Any, Generic, cast from tenacity import ( before_sleep_log, @@ -265,7 +265,7 @@ class MCPClient: mcp.ClientSession( *streams, read_timeout_seconds=read_timeout, - logging_callback=logging_callback, + logging_callback=logging_callback, # type: ignore[arg-type] ), ) else: @@ -291,7 +291,7 @@ class MCPClient: read_stream=read_s, write_stream=write_s, read_timeout_seconds=read_timeout, - logging_callback=logging_callback, + logging_callback=logging_callback, # type: ignore[arg-type] ), ) @@ -317,11 +317,14 @@ class MCPClient: stdio_transport = await self.exit_stack.enter_async_context( mcp.stdio_client( server_params, - errlog=LogPipe( - level=logging.INFO, - logger=logger, - identifier=f"MCPServer-{name}", - callback=callback, + errlog=cast( + Any, + LogPipe( + level=logging.INFO, + logger=logger, + identifier=f"MCPServer-{name}", + callback=callback, + ), ), ), ) diff --git a/astrbot/core/agent/runners/coze/coze_agent_runner.py b/astrbot/core/agent/runners/coze/coze_agent_runner.py index 93c1f8970..055ec651e 100644 --- a/astrbot/core/agent/runners/coze/coze_agent_runner.py +++ b/astrbot/core/agent/runners/coze/coze_agent_runner.py @@ -1,7 +1,6 @@ import base64 import json import sys -from collections.abc import AsyncGenerator from typing import Any import astrbot.core.message.components as Comp @@ -86,7 +85,7 @@ class CozeAgentRunner(BaseAgentRunner[TContext]): self.file_id_cache: dict[str, dict[str, str]] = {} @override - async def step(self) -> AsyncGenerator[AgentResponse, None]: + async def step(self): """ 执行 Coze Agent 的一个步骤 """ @@ -122,9 +121,7 @@ class CozeAgentRunner(BaseAgentRunner[TContext]): await self.api_client.close() @override - async def step_until_done( - self, max_step: int - ) -> AsyncGenerator[AgentResponse, None]: + async def step_until_done(self, max_step: int): while not self.done(): async for resp in self.step(): yield resp diff --git a/astrbot/core/agent/runners/dashscope/dashscope_agent_runner.py b/astrbot/core/agent/runners/dashscope/dashscope_agent_runner.py index 766474c23..fef568d9b 100644 --- a/astrbot/core/agent/runners/dashscope/dashscope_agent_runner.py +++ b/astrbot/core/agent/runners/dashscope/dashscope_agent_runner.py @@ -131,9 +131,7 @@ class DashscopeAgentRunner(BaseAgentRunner[TContext]): ) @override - async def step_until_done( - self, max_step: int - ) -> AsyncGenerator[AgentResponse, None]: + async def step_until_done(self, max_step: int): while not self.done(): async for resp in self.step(): yield resp @@ -260,7 +258,7 @@ class DashscopeAgentRunner(BaseAgentRunner[TContext]): key="session_variables", default={}, ) - payload_vars.update(session_var) + payload_vars.update(session_var) # type: ignore[arg-type] if ( self.dashscope_app_type in ["agent", "dialog-workflow"] diff --git a/astrbot/core/agent/runners/deerflow/deerflow_agent_runner.py b/astrbot/core/agent/runners/deerflow/deerflow_agent_runner.py index 518d92681..01bb06680 100644 --- a/astrbot/core/agent/runners/deerflow/deerflow_agent_runner.py +++ b/astrbot/core/agent/runners/deerflow/deerflow_agent_runner.py @@ -4,7 +4,6 @@ import json import sys import typing as T from collections import deque -from collections.abc import AsyncGenerator from dataclasses import dataclass, field from typing import Any from uuid import uuid4 @@ -13,9 +12,9 @@ import astrbot.core.message.components as Comp from astrbot import logger from astrbot.core import sp from astrbot.core.agent.hooks import BaseAgentRunHooks -from astrbot.core.agent.response import AgentResponseData +from astrbot.core.agent.response import AgentResponse, AgentResponseData from astrbot.core.agent.run_context import ContextWrapper, TContext -from astrbot.core.agent.runners.base import AgentResponse, AgentState, BaseAgentRunner +from astrbot.core.agent.runners.base import AgentState, BaseAgentRunner from astrbot.core.agent.tool_executor import BaseFunctionToolExecutor from astrbot.core.message.message_event_result import MessageChain from astrbot.core.provider.entities import ( @@ -319,9 +318,7 @@ class DeerFlowAgentRunner(BaseAgentRunner[TContext]): yield await self._finish_with_error(err_msg) @override - async def step_until_done( - self, max_step: int - ) -> AsyncGenerator[AgentResponse, None]: + async def step_until_done(self, max_step: int): if max_step <= 0: raise ValueError("max_step must be greater than 0") diff --git a/astrbot/core/agent/runners/dify/dify_agent_runner.py b/astrbot/core/agent/runners/dify/dify_agent_runner.py index ff22ac299..b1ddc2531 100644 --- a/astrbot/core/agent/runners/dify/dify_agent_runner.py +++ b/astrbot/core/agent/runners/dify/dify_agent_runner.py @@ -114,9 +114,7 @@ class DifyAgentRunner(BaseAgentRunner[TContext]): await self.api_client.close() @override - async def step_until_done( - self, max_step: int - ) -> AsyncGenerator[AgentResponse, None]: + async def step_until_done(self, max_step: int): while not self.done(): async for resp in self.step(): yield resp @@ -174,7 +172,7 @@ class DifyAgentRunner(BaseAgentRunner[TContext]): key="session_variables", default={}, ) - payload_vars.update(session_var) + payload_vars.update(session_var) # type: ignore[arg-type] payload_vars["system_prompt"] = system_prompt # 处理不同的 API 类型 @@ -189,7 +187,7 @@ class DifyAgentRunner(BaseAgentRunner[TContext]): }, query=prompt, user=session_id, - conversation_id=conversation_id, + conversation_id=conversation_id or "", files=files_payload, request_timeout=self.timeout, ): diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index e57507980..bc457f865 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -6,7 +6,7 @@ import traceback from collections.abc import AsyncGenerator, AsyncIterator from contextlib import suppress from dataclasses import dataclass, field -from typing import Any, Literal, TypeVar +from typing import Any, Literal, TypeVar, cast from mcp.types import ( BlobResourceContents, @@ -609,9 +609,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): self.req.append_tool_calls_result(tool_calls_result) - async def step_until_done( - self, max_step: int - ) -> AsyncGenerator[AgentResponse, None]: + async def step_until_done(self, max_step: int): """Process steps until the agent is done.""" step_count = 0 max_step = min(max_step, 3) @@ -772,7 +770,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): ) _final_resp: CallToolResult | None = None - async for resp in self._iter_tool_executor_results(executor): + async for resp in self._iter_tool_executor_results(executor): # type: ignore[arg-type] if isinstance(resp, CallToolResult): res = resp _final_resp = resp diff --git a/astrbot/core/computer/tools/shell.py b/astrbot/core/computer/tools/shell.py index ba1cd8b5a..f0c4b1bdc 100644 --- a/astrbot/core/computer/tools/shell.py +++ b/astrbot/core/computer/tools/shell.py @@ -25,9 +25,8 @@ import subprocess from dataclasses import dataclass, field from typing import Any -from astrbot.api import FunctionTool from astrbot.core.agent.run_context import ContextWrapper -from astrbot.core.agent.tool import ToolExecResult +from astrbot.core.agent.tool import FunctionTool, ToolExecResult from astrbot.core.astr_agent_context import AstrAgentContext from .permissions import check_admin_permission diff --git a/tests/test_profile_aware_tools.py b/tests/test_profile_aware_tools.py index 7778bf305..03948aaf1 100644 --- a/tests/test_profile_aware_tools.py +++ b/tests/test_profile_aware_tools.py @@ -49,7 +49,7 @@ class TestShipyardNeoBooterCapabilities: caps = booter.capabilities assert isinstance(caps, tuple) with pytest.raises(AttributeError): - caps.append("mutated") # type: ignore[attr-defined] + caps.append("mutated") # ═══════════════════════════════════════════════════════════════