From a306c63a92472384522f3cf4074460e8df74f790 Mon Sep 17 00:00:00 2001 From: LIghtJUNction Date: Tue, 31 Mar 2026 20:16:47 +0800 Subject: [PATCH] refactor(agent): improve type safety in agent components message.py: - Use TypeGuard for type narrowing instead of isinstance checks - Improve type annotations for ContentPart validation - Add type annotations for content part registry mcp_client.py: - Improve type annotations and code quality runners (base, dashscope, deerflow, dify, tool_loop): - Add/improve type annotations - Clean up code structure tool.py & tool_image_cache.py: - Improve type annotations --- astrbot/core/agent/mcp_client.py | 60 +++++++++++-------- astrbot/core/agent/message.py | 41 +++++++------ astrbot/core/agent/runners/base.py | 10 ++-- .../dashscope/dashscope_agent_runner.py | 49 +++++++++------ .../runners/deerflow/deerflow_agent_runner.py | 1 + .../runners/deerflow/deerflow_api_client.py | 4 +- .../agent/runners/dify/dify_agent_runner.py | 19 +++--- .../agent/runners/dify/dify_api_client.py | 40 ++++++++----- .../agent/runners/tool_loop_agent_runner.py | 15 ++--- astrbot/core/agent/tool.py | 36 +++++++---- astrbot/core/agent/tool_image_cache.py | 18 +++--- 11 files changed, 172 insertions(+), 121 deletions(-) diff --git a/astrbot/core/agent/mcp_client.py b/astrbot/core/agent/mcp_client.py index 0518bfa75..71135f271 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, cast +from typing import Any, Generic, TextIO from tenacity import ( before_sleep_log, @@ -226,14 +226,13 @@ class MCPClient: cfg = _prepare_config(mcp_server_config.copy()) - def logging_callback( - msg: str | mcp.types.LoggingMessageNotificationParams, + async def logging_callback( + params: mcp.types.LoggingMessageNotificationParams, ) -> None: # Handle MCP service error logs - if isinstance(msg, mcp.types.LoggingMessageNotificationParams): - if msg.level in ("warning", "error", "critical", "alert", "emergency"): - log_msg = f"[{msg.level.upper()}] {msg.data!s}" - self.server_errlogs.append(log_msg) + if params.level in ("warning", "error", "critical", "alert", "emergency"): + log_msg = f"[{params.level.upper()}] {params.data!s}" + self.server_errlogs.append(log_msg) if "url" in cfg: success, error_msg = await _quick_test_mcp_connection(cfg) @@ -255,19 +254,21 @@ class MCPClient: timeout=cfg.get("timeout", 5), sse_read_timeout=cfg.get("sse_read_timeout", 60 * 5), ) - streams = await self.exit_stack.enter_async_context( + read_stream, write_stream = await self.exit_stack.enter_async_context( self._streams_context, ) # Create a new client session read_timeout = timedelta(seconds=cfg.get("session_read_timeout", 60)) - self.session = await self.exit_stack.enter_async_context( + session = await self.exit_stack.enter_async_context( mcp.ClientSession( - *streams, + read_stream=read_stream, + write_stream=write_stream, read_timeout_seconds=read_timeout, - logging_callback=logging_callback, # type: ignore[arg-type] - ), + logging_callback=logging_callback, + ) ) + self.session = session else: timeout = timedelta(seconds=cfg.get("timeout", 30)) sse_read_timeout = timedelta( @@ -286,14 +287,15 @@ class MCPClient: # Create a new client session read_timeout = timedelta(seconds=cfg.get("session_read_timeout", 60)) - self.session = await self.exit_stack.enter_async_context( + session = await self.exit_stack.enter_async_context( mcp.ClientSession( read_stream=read_s, write_stream=write_s, read_timeout_seconds=read_timeout, - logging_callback=logging_callback, # type: ignore[arg-type] - ), + logging_callback=logging_callback, + ) ) + self.session = session else: cfg = _prepare_stdio_env(cfg) @@ -314,26 +316,32 @@ class MCPClient: log_msg = f"[{msg.level.upper()}] {msg.data!s}" self.server_errlogs.append(log_msg) + log_pipe = self.exit_stack.enter_context( + LogPipe( + level=logging.INFO, + logger=logger, + identifier=f"MCPServer-{name}", + callback=callback, + ) + ) + errlog_stream: TextIO = self.exit_stack.enter_context( + os.fdopen(os.dup(log_pipe.fileno()), "w") + ) stdio_transport = await self.exit_stack.enter_async_context( mcp.stdio_client( server_params, - errlog=cast( - Any, - LogPipe( - level=logging.INFO, - logger=logger, - identifier=f"MCPServer-{name}", - callback=callback, - ), - ), + errlog=errlog_stream, ), ) - self.process_pid = self._extract_stdio_process_pid(self._streams_context) + self.process_pid = self._extract_stdio_process_pid(stdio_transport) # Create a new client session - self.session = await self.exit_stack.enter_async_context( + session = await self.exit_stack.enter_async_context( mcp.ClientSession(*stdio_transport), ) + self.session = session + + assert self.session is not None await self.session.initialize() async def list_tools_and_save(self) -> mcp.ListToolsResult: diff --git a/astrbot/core/agent/message.py b/astrbot/core/agent/message.py index bde6353ff..87c60f42c 100644 --- a/astrbot/core/agent/message.py +++ b/astrbot/core/agent/message.py @@ -1,7 +1,8 @@ +from __future__ import annotations + # Inspired by MoonshotAI/kosong, credits to MoonshotAI/kosong authors for the original implementation. # License: Apache License 2.0 - -from typing import Any, ClassVar, Literal, cast +from typing import Any, ClassVar, Literal, TypeGuard from pydantic import ( BaseModel, @@ -13,10 +14,14 @@ from pydantic import ( from pydantic_core import core_schema +def _is_str_keyed_dict(value: object) -> TypeGuard[dict[str, object]]: + return isinstance(value, dict) and all(isinstance(key, str) for key in value) + + class ContentPart(BaseModel): """A part of the content in a message.""" - __content_part_registry: ClassVar[dict[str, type["ContentPart"]]] = {} + __content_part_registry: ClassVar[dict[str, type[ContentPart]]] = {} type: Literal["text", "think", "image_url", "audio_url"] @@ -33,23 +38,23 @@ class ContentPart(BaseModel): @classmethod def __get_pydantic_core_schema__( - cls, source_type: Any, handler: GetCoreSchemaHandler + cls, source_type: object, handler: GetCoreSchemaHandler ) -> core_schema.CoreSchema: # If we're dealing with the base ContentPart class, use custom validation if cls.__name__ == "ContentPart": - def validate_content_part(value: Any) -> Any: + def validate_content_part(value: object) -> ContentPart: # if it's already an instance of a ContentPart subclass, return it - if hasattr(value, "__class__") and issubclass(value.__class__, cls): + if isinstance(value, cls): return value # if it's a dict with a type field, dispatch to the appropriate subclass - if isinstance(value, dict) and "type" in value: - type_value: Any | None = cast(dict[str, Any], value).get("type") - if not isinstance(type_value, str): - raise ValueError(f"Cannot validate {value} as ContentPart") - target_class = cls.__content_part_registry[type_value] - return target_class.model_validate(value) + if _is_str_keyed_dict(value): + type_value = value.get("type") + if isinstance(type_value, str): + target_class = cls.__content_part_registry.get(type_value) + if target_class is not None: + return target_class.model_validate(value) raise ValueError(f"Cannot validate {value} as ContentPart") @@ -65,7 +70,7 @@ class TextPart(ContentPart): {'type': 'text', 'text': 'Hello, world!'} """ - type: str = "text" + type: Literal["text"] = "text" text: str @@ -75,12 +80,12 @@ class ThinkPart(ContentPart): {'type': 'think', 'think': 'I think I need to think about this.', 'encrypted': None} """ - type: str = "think" + type: Literal["think"] = "think" think: str encrypted: str | None = None """Encrypted thinking content, or signature.""" - def merge_in_place(self, other: Any) -> bool: + def merge_in_place(self, other: object) -> bool: if not isinstance(other, ThinkPart): return False if self.encrypted: @@ -103,7 +108,7 @@ class ImageURLPart(ContentPart): id: str | None = None """The ID of the image, to allow LLMs to distinguish different images.""" - type: str = "image_url" + type: Literal["image_url"] = "image_url" image_url: ImageURL @@ -119,7 +124,7 @@ class AudioURLPart(ContentPart): id: str | None = None """The ID of the audio, to allow LLMs to distinguish different audios.""" - type: str = "audio_url" + type: Literal["audio_url"] = "audio_url" audio_url: AudioURL @@ -147,7 +152,7 @@ class ToolCall(BaseModel): """The ID of the tool call.""" function: FunctionBody """The function body of the tool call.""" - extra_content: dict[str, Any] | None = None + extra_content: dict[str, object] | None = None """Extra metadata for the tool call.""" @model_serializer(mode="wrap") diff --git a/astrbot/core/agent/runners/base.py b/astrbot/core/agent/runners/base.py index e5cdd42d7..56f0457c6 100644 --- a/astrbot/core/agent/runners/base.py +++ b/astrbot/core/agent/runners/base.py @@ -1,4 +1,5 @@ import abc +import asyncio from collections.abc import AsyncGenerator from enum import Enum, auto from typing import Any, Generic @@ -25,7 +26,8 @@ class BaseAgentRunner(Generic[TContext]): def __init__( self, ): - self.tasks: set = set() + self.tasks: set[asyncio.Task[object]] = set() + self._state = AgentState.IDLE @abc.abstractmethod async def reset( @@ -54,14 +56,12 @@ class BaseAgentRunner(Generic[TContext]): ... @abc.abstractmethod - async def step(self) -> AsyncGenerator[AgentResponse, None]: + def step(self) -> AsyncGenerator[AgentResponse, None]: """Process a single step of the agent.""" ... @abc.abstractmethod - async def step_until_done( - self, max_step: int - ) -> AsyncGenerator[AgentResponse, None]: + def step_until_done(self, max_step: int) -> AsyncGenerator[AgentResponse, None]: """Process steps until the agent is done.""" ... diff --git a/astrbot/core/agent/runners/dashscope/dashscope_agent_runner.py b/astrbot/core/agent/runners/dashscope/dashscope_agent_runner.py index fef568d9b..5a81088e9 100644 --- a/astrbot/core/agent/runners/dashscope/dashscope_agent_runner.py +++ b/astrbot/core/agent/runners/dashscope/dashscope_agent_runner.py @@ -56,7 +56,7 @@ class DashscopeAgentRunner(BaseAgentRunner[TContext]): ) -> None: self.req = request self.streaming = streaming - self.final_llm_resp = None + self.final_llm_resp: LLMResponse | None = None self._state = AgentState.IDLE self.agent_hooks = agent_hooks self.run_context = run_context @@ -193,7 +193,8 @@ class DashscopeAgentRunner(BaseAgentRunner[TContext]): ), ) - chunk_text = chunk.output.get("text", "") or "" + chunk_text_value = chunk.output.get("text", "") + chunk_text = chunk_text_value if isinstance(chunk_text_value, str) else "" # RAG 引用脚标格式化 chunk_text = re.sub(r"\[(\d+)\]", r"[\1]", chunk_text) @@ -206,7 +207,10 @@ class DashscopeAgentRunner(BaseAgentRunner[TContext]): ) # 获取文档引用 - doc_references = chunk.output.get("doc_references", None) + raw_doc_references = chunk.output.get("doc_references") + doc_references = ( + raw_doc_references if isinstance(raw_doc_references, list) else None + ) return output_text, doc_references, response @@ -251,15 +255,17 @@ class DashscopeAgentRunner(BaseAgentRunner[TContext]): default="", ) # 获得会话变量 - payload_vars = self.variables.copy() - session_var = await sp.get_async( - scope="umo", - scope_id=session_id, - key="session_variables", - default={}, + payload_vars: dict = self.variables.copy() + session_var: dict = ( + await sp.get_async( + scope="umo", + scope_id=session_id, + key="session_variables", + default={}, + ) + or {} ) - payload_vars.update(session_var) # type: ignore[arg-type] - + payload_vars.update(session_var) if ( self.dashscope_app_type in ["agent", "dialog-workflow"] and not self.has_rag_options() @@ -302,7 +308,7 @@ class DashscopeAgentRunner(BaseAgentRunner[TContext]): AgentResponse 对象 """ - response_queue = queue.Queue() + response_queue: queue.Queue[tuple[str, Any]] = queue.Queue() consumer_thread = threading.Thread( target=self._consume_sync_generator, args=(response, response_queue), @@ -324,6 +330,10 @@ class DashscopeAgentRunner(BaseAgentRunner[TContext]): if item_type == "done": break elif item_type == "error": + if not isinstance(item_data, BaseException): + raise RuntimeError( + f"Unexpected Dashscope error payload: {item_data!r}" + ) raise item_data elif item_type == "data": chunk = item_data @@ -332,14 +342,14 @@ class DashscopeAgentRunner(BaseAgentRunner[TContext]): ( output_text, chunk_doc_refs, - response, + agent_response, ) = await self._process_stream_chunk(chunk, output_text) - if response: - if response.type == "err": - yield response + if agent_response: + if agent_response.type == "err": + yield agent_response return - yield response + yield agent_response if chunk_doc_refs: doc_references = chunk_doc_refs @@ -365,11 +375,12 @@ class DashscopeAgentRunner(BaseAgentRunner[TContext]): # 创建最终响应 chain = MessageChain(chain=[Comp.Plain(output_text)]) - self.final_llm_resp = LLMResponse(role="assistant", result_chain=chain) + final_llm_resp = LLMResponse(role="assistant", result_chain=chain) + self.final_llm_resp = final_llm_resp self._transition_state(AgentState.DONE) try: - await self.agent_hooks.on_agent_done(self.run_context, self.final_llm_resp) + await self.agent_hooks.on_agent_done(self.run_context, final_llm_resp) except Exception as e: logger.error(f"Error in on_agent_done hook: {e}", exc_info=True) diff --git a/astrbot/core/agent/runners/deerflow/deerflow_agent_runner.py b/astrbot/core/agent/runners/deerflow/deerflow_agent_runner.py index 01bb06680..1eabb6c23 100644 --- a/astrbot/core/agent/runners/deerflow/deerflow_agent_runner.py +++ b/astrbot/core/agent/runners/deerflow/deerflow_agent_runner.py @@ -53,6 +53,7 @@ class DeerFlowAgentRunner(BaseAgentRunner[TContext]): """DeerFlow Agent Runner via LangGraph HTTP API.""" _MAX_VALUES_HISTORY = 200 + final_llm_resp: LLMResponse | None @dataclass(frozen=True) class _RunnerConfig: diff --git a/astrbot/core/agent/runners/deerflow/deerflow_api_client.py b/astrbot/core/agent/runners/deerflow/deerflow_api_client.py index 160875af7..17999657f 100644 --- a/astrbot/core/agent/runners/deerflow/deerflow_api_client.py +++ b/astrbot/core/agent/runners/deerflow/deerflow_api_client.py @@ -144,12 +144,12 @@ class DeerFlowAPIClient: async def create_thread(self, timeout: float = 20) -> dict[str, Any]: session = self._get_session() url = f"{self.api_base}/api/langgraph/threads" - payload = {"metadata": {}} + payload: dict[str, dict[str, object]] = {"metadata": {}} async with session.post( url, json=payload, headers=self.headers, - timeout=timeout, + timeout=ClientTimeout(total=timeout), proxy=self.proxy, ) as resp: if resp.status not in (200, 201): diff --git a/astrbot/core/agent/runners/dify/dify_agent_runner.py b/astrbot/core/agent/runners/dify/dify_agent_runner.py index d8547e6a7..cd19c900e 100644 --- a/astrbot/core/agent/runners/dify/dify_agent_runner.py +++ b/astrbot/core/agent/runners/dify/dify_agent_runner.py @@ -165,13 +165,16 @@ class DifyAgentRunner(BaseAgentRunner[TContext]): # 获得会话变量 payload_vars = self.variables.copy() # 动态变量 - session_var = await sp.get_async( - scope="umo", - scope_id=session_id, - key="session_variables", - default={}, + session_var: dict = ( + await sp.get_async( + scope="umo", + scope_id=session_id, + key="session_variables", + default={}, + ) + or {} ) - payload_vars.update(session_var) # type: ignore[arg-type] + payload_vars.update(session_var) payload_vars["system_prompt"] = system_prompt # 处理不同的 API 类型 @@ -297,7 +300,7 @@ class DifyAgentRunner(BaseAgentRunner[TContext]): # Chat return MessageChain(chain=[Comp.Plain(chunk)]) - async def parse_file(item: dict): + async def parse_file(item: dict) -> Comp.BaseMessageComponent: match item["type"]: case "image": return Comp.Image(file=item["url"], url=item["url"]) @@ -313,7 +316,7 @@ class DifyAgentRunner(BaseAgentRunner[TContext]): return Comp.File(name=item["filename"], file=item["url"]) output = chunk["data"]["outputs"][self.workflow_output_key] - chains = [] + chains: list[Comp.BaseMessageComponent] = [] if isinstance(output, str): # 纯文本输出 chains.append(Comp.Plain(output)) diff --git a/astrbot/core/agent/runners/dify/dify_api_client.py b/astrbot/core/agent/runners/dify/dify_api_client.py index bd3949063..ace2ea384 100644 --- a/astrbot/core/agent/runners/dify/dify_api_client.py +++ b/astrbot/core/agent/runners/dify/dify_api_client.py @@ -4,7 +4,7 @@ from collections.abc import AsyncGenerator from typing import Any import anyio -from aiohttp import ClientResponse, ClientSession, FormData +from aiohttp import ClientResponse, ClientSession, ClientTimeout, FormData from astrbot.core import logger @@ -36,32 +36,37 @@ class DifyAPIClient: self.api_key = api_key self.api_base = api_base self.session = ClientSession(trust_env=True) - self.headers = { + self.headers: dict[str, str] = { "Authorization": f"Bearer {self.api_key}", } async def chat_messages( self, - inputs: dict, + inputs: dict[str, object], query: str, user: str, response_mode: str = "streaming", conversation_id: str = "", - files: list[dict[str, Any]] | None = None, + files: list[dict[str, object]] | None = None, request_timeout: float = 60, ) -> AsyncGenerator[dict[str, Any], None]: if files is None: files = [] url = f"{self.api_base}/chat-messages" - payload = locals() - payload.pop("self") - payload.pop("request_timeout") + payload: dict[str, object] = { + "inputs": inputs, + "query": query, + "user": user, + "response_mode": response_mode, + "conversation_id": conversation_id, + "files": files, + } logger.info(f"chat_messages payload: {payload}") async with self.session.post( url, json=payload, headers=self.headers, - timeout=request_timeout, + timeout=ClientTimeout(total=request_timeout), ) as resp: if resp.status != 200: text = await resp.text() @@ -73,24 +78,27 @@ class DifyAPIClient: async def workflow_run( self, - inputs: dict, + inputs: dict[str, object], user: str, response_mode: str = "streaming", - files: list[dict[str, Any]] | None = None, + files: list[dict[str, object]] | None = None, request_timeout: float = 60, ): if files is None: files = [] url = f"{self.api_base}/workflows/run" - payload = locals() - payload.pop("self") - payload.pop("request_timeout") + payload: dict[str, object] = { + "inputs": inputs, + "user": user, + "response_mode": response_mode, + "files": files, + } logger.info(f"workflow_run payload: {payload}") async with self.session.post( url, json=payload, headers=self.headers, - timeout=request_timeout, + timeout=ClientTimeout(total=request_timeout), ) as resp: if resp.status != 200: text = await resp.text() @@ -162,11 +170,11 @@ class DifyAPIClient: async def get_chat_convs(self, user: str, limit: int = 20): # conversations. GET url = f"{self.api_base}/conversations" - payload = { + params: dict[str, str | int] = { "user": user, "limit": limit, } - async with self.session.get(url, params=payload, headers=self.headers) as resp: + async with self.session.get(url, params=params, headers=self.headers) as resp: return await resp.json() async def delete_chat_conv(self, user: str, conversation_id: str): diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index 7a5dc186b..9796d4f1c 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -25,6 +25,7 @@ from astrbot.core.agent.context.token_counter import TokenCounter from astrbot.core.agent.hooks import BaseAgentRunHooks from astrbot.core.agent.message import ( AssistantMessageSegment, + ContentPart, ImageURLPart, Message, TextPart, @@ -175,7 +176,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): self.fallback_providers.append(fallback_provider) if fallback_id: seen_provider_ids.add(fallback_id) - self.final_llm_resp = None + self.final_llm_resp: LLMResponse | None = None self._state = AgentState.IDLE self.tool_executor = tool_executor self.agent_hooks = agent_hooks @@ -213,8 +214,8 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): m._no_save = True messages.append(m) if request.prompt is not None: - m = await request.assemble_context() - messages.append(Message.model_validate(m)) + assembled_context = await request.assemble_context() + messages.append(Message.model_validate(assembled_context)) if request.system_prompt: messages.insert( 0, @@ -766,7 +767,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): except Exception as e: logger.error(f"Error in on_tool_start hook: {e}", exc_info=True) - executor = self.tool_executor.execute( + executor = await self.tool_executor.execute( tool=func_tool, run_context=self.run_context, session_manager=self.run_context.session_manager, @@ -774,9 +775,9 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): ) _final_resp: CallToolResult | None = None - async for resp in self._iter_tool_executor_results(executor): # type: ignore[arg-type] + async for resp in self._iter_tool_executor_results(executor): if isinstance(resp, CallToolResult): - res = resp + res: CallToolResult = resp _final_resp = resp if not res.content: _append_tool_call_result( @@ -1011,7 +1012,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): self._transition_state(AgentState.DONE) self.stats.end_time = time.time() - parts = [] + parts: list[ContentPart] = [] if llm_resp.reasoning_content or llm_resp.reasoning_signature: parts.append( ThinkPart( diff --git a/astrbot/core/agent/tool.py b/astrbot/core/agent/tool.py index 7aa3e143a..0b09818ec 100644 --- a/astrbot/core/agent/tool.py +++ b/astrbot/core/agent/tool.py @@ -1,7 +1,7 @@ import copy from collections.abc import AsyncGenerator, Awaitable, Callable from dataclasses import field -from typing import Any, Generic +from typing import Any, Generic, TypedDict import jsonschema import mcp @@ -17,6 +17,12 @@ ParametersType = dict[str, Any] ToolExecResult = str | mcp.types.CallToolResult +class ToolArgumentSpec(TypedDict): + name: str + type: str + description: str + + @dataclass class ToolSchema: """A class representing the schema of a tool for function calling.""" @@ -27,7 +33,11 @@ class ToolSchema: description: str """The description of the tool.""" - parameters: ParametersType + parameters: ParametersType = field(default_factory=dict) + """The parameters of the tool, in JSON Schema format.""" + + active: bool = True + """Whether the tool is active.""" """The parameters of the tool, in JSON Schema format.""" @model_validator(mode="after") @@ -111,13 +121,13 @@ class ToolSet: convert the tools to different API formats (OpenAI, Anthropic, Google GenAI). """ - tools: list[FunctionTool] = Field(default_factory=list) + tools: list[ToolSchema] = Field(default_factory=list) def empty(self) -> bool: """Check if the tool set is empty.""" return len(self.tools) == 0 - def add_tool(self, tool: FunctionTool) -> None: + def add_tool(self, tool: ToolSchema) -> None: """Add a tool to the set. If a tool with the same name already exists: @@ -153,12 +163,13 @@ class ToolSet: """Get a tool by its name.""" for tool in self.tools: if tool.name == name: - return tool + if isinstance(tool, FunctionTool): + return tool return None def get_light_tool_set(self) -> "ToolSet": """Return a light tool set with only name/description.""" - light_tools = [] + light_tools: list[ToolSchema] = [] for tool in self.tools: if hasattr(tool, "active") and not tool.active: continue @@ -178,7 +189,7 @@ class ToolSet: def get_param_only_tool_set(self) -> "ToolSet": """Return a tool set with name/parameters only (no description).""" - param_tools = [] + param_tools: list[ToolSchema] = [] for tool in self.tools: if hasattr(tool, "active") and not tool.active: continue @@ -201,17 +212,18 @@ class ToolSet: def add_func( self, name: str, - func_args: list, + func_args: list[ToolArgumentSpec], desc: str, handler: Callable[..., Awaitable[Any]], ) -> None: """Add a function tool to the set.""" + properties: dict[str, dict[str, str]] = {} params = { "type": "object", # hard-coded here - "properties": {}, + "properties": properties, } for param in func_args: - params["properties"][param["name"]] = { + properties[param["name"]] = { "type": param["type"], "description": param["description"], } @@ -236,11 +248,11 @@ class ToolSet: @property def func_list(self) -> list[FunctionTool]: """Get the list of function tools.""" - return self.tools + return [t for t in self.tools if isinstance(t, FunctionTool)] def list_tools(self) -> list[FunctionTool]: """Get the list of function tools (alias for func_list).""" - return self.tools + return [t for t in self.tools if isinstance(t, FunctionTool)] def openai_schema(self, omit_empty_parameter_field: bool = False) -> list[dict]: """Convert tools to OpenAI API function calling schema format.""" diff --git a/astrbot/core/agent/tool_image_cache.py b/astrbot/core/agent/tool_image_cache.py index de4f9e243..9c99b3509 100644 --- a/astrbot/core/agent/tool_image_cache.py +++ b/astrbot/core/agent/tool_image_cache.py @@ -7,9 +7,7 @@ import base64 import os import time from dataclasses import dataclass, field -from typing import ClassVar, cast - -from typing_extensions import Self +from typing import ClassVar, Self from astrbot import logger from astrbot.core.utils.astrbot_path import get_astrbot_temp_path @@ -37,16 +35,20 @@ class ToolImageCache: Images are stored in data/temp/tool_images/ and can be retrieved by file path. """ - _instance: ClassVar["ToolImageCache | None"] = None + _instance: ClassVar[Self | None] = None CACHE_DIR_NAME: ClassVar[str] = "tool_images" # Cache expiry time in seconds (1 hour) CACHE_EXPIRY: ClassVar[int] = 3600 + _initialized: bool + _cache_dir: str def __new__(cls) -> Self: - if cls._instance is None: - cls._instance = super().__new__(cls) - cls._instance._initialized = False - return cast(Self, cls._instance) + instance = cls._instance + if instance is None: + instance = super().__new__(cls) + instance._initialized = False + cls._instance = instance + return instance def __init__(self) -> None: if self._initialized: