AstrBot 3.5.6__py3-none-any.whl → 4.7.0__py3-none-any.whl

astrbot/core/star/register/star_handler.py

@@ -1,73 +1,101 @@
 from __future__ import annotations
+
+import re
+from collections.abc import Awaitable, Callable
+from typing import Any
+
 import docstring_parser
 
-from ..star_handler import star_handlers_registry, StarHandlerMetadata, EventType
+from astrbot.core import logger
+from astrbot.core.agent.agent import Agent
+from astrbot.core.agent.handoff import HandoffTool
+from astrbot.core.agent.hooks import BaseAgentRunHooks
+from astrbot.core.agent.tool import FunctionTool
+from astrbot.core.astr_agent_context import AstrAgentContext
+from astrbot.core.provider.func_tool_manager import PY_TO_JSON_TYPE, SUPPORTED_TYPES
+from astrbot.core.provider.register import llm_tools
+
 from ..filter.command import CommandFilter
 from ..filter.command_group import CommandGroupFilter
-from ..filter.event_message_type import EventMessageTypeFilter, EventMessageType
+from ..filter.custom_filter import CustomFilterAnd, CustomFilterOr
+from ..filter.event_message_type import EventMessageType, EventMessageTypeFilter
+from ..filter.permission import PermissionType, PermissionTypeFilter
 from ..filter.platform_adapter_type import (
-    PlatformAdapterTypeFilter,
     PlatformAdapterType,
+    PlatformAdapterTypeFilter,
 )
-from ..filter.permission import PermissionTypeFilter, PermissionType
-from ..filter.custom_filter import CustomFilterAnd, CustomFilterOr
 from ..filter.regex import RegexFilter
-from typing import Awaitable
-from astrbot.core.provider.func_tool_manager import SUPPORTED_TYPES
-from astrbot.core.provider.register import llm_tools
+from ..star_handler import EventType, StarHandlerMetadata, star_handlers_registry
 
 
-def get_handler_full_name(awaitable: Awaitable) -> str:
+def get_handler_full_name(awaitable: Callable[..., Awaitable[Any]]) -> str:
     """获取 Handler 的全名"""
     return f"{awaitable.__module__}_{awaitable.__name__}"
 
 
 def get_handler_or_create(
-    handler: Awaitable, event_type: EventType, dont_add=False, **kwargs
+    handler: Callable[..., Awaitable[Any]],
+    event_type: EventType,
+    dont_add=False,
+    **kwargs,
 ) -> StarHandlerMetadata:
     """获取 Handler 或者创建一个新的 Handler"""
     handler_full_name = get_handler_full_name(handler)
     md = star_handlers_registry.get_handler_by_full_name(handler_full_name)
     if md:
         return md
-    else:
-        md = StarHandlerMetadata(
-            event_type=event_type,
-            handler_full_name=handler_full_name,
-            handler_name=handler.__name__,
-            handler_module_path=handler.__module__,
-            handler=handler,
-            event_filters=[],
-        )
-
-        # 插件handler的附加额外信息
-        if handler.__doc__:
-            md.desc = handler.__doc__.strip()
-        if "desc" in kwargs:
-            md.desc = kwargs["desc"]
-            del kwargs["desc"]
-        md.extras_configs = kwargs
-
-        if not dont_add:
-            star_handlers_registry.append(md)
-        return md
+    md = StarHandlerMetadata(
+        event_type=event_type,
+        handler_full_name=handler_full_name,
+        handler_name=handler.__name__,
+        handler_module_path=handler.__module__,
+        handler=handler,
+        event_filters=[],
+    )
+
+    # 插件handler的附加额外信息
+    if handler.__doc__:
+        md.desc = handler.__doc__.strip()
+    if "desc" in kwargs:
+        md.desc = kwargs["desc"]
+        del kwargs["desc"]
+    md.extras_configs = kwargs
+
+    if not dont_add:
+        star_handlers_registry.append(md)
+    return md
 
 
 def register_command(
-    command_name: str = None, sub_command: str = None, alias: set = None, **kwargs
+    command_name: str | None = None,
+    sub_command: str | None = None,
+    alias: set | None = None,
+    **kwargs,
 ):
     """注册一个 Command."""
     new_command = None
     add_to_event_filters = False
     if isinstance(command_name, RegisteringCommandable):
         # 子指令
-        parent_command_names = command_name.parent_group.get_complete_command_names()
-        new_command = CommandFilter(
-            sub_command, alias, None, parent_command_names=parent_command_names
-        )
-        command_name.parent_group.add_sub_command_filter(new_command)
+        if sub_command is not None:
+            parent_command_names = (
+                command_name.parent_group.get_complete_command_names()
+            )
+            new_command = CommandFilter(
+                sub_command,
+                alias,
+                None,
+                parent_command_names=parent_command_names,
+            )
+            command_name.parent_group.add_sub_command_filter(new_command)
+        else:
+            logger.warning(
+                f"注册指令{command_name} 的子指令时未提供 sub_command 参数。",
+            )
+    # 裸指令
+    elif command_name is None:
+        logger.warning("注册裸指令时未提供 command_name 参数。")
     else:
-        # 裸指令
         new_command = CommandFilter(command_name, alias, None)
         add_to_event_filters = True
 
@@ -77,10 +105,13 @@ def register_command(
                 True  # 打一个标记，表示这是一个子指令，再 wakingstage 阶段这个 handler 将会直接被跳过（其父指令会接管）
             )
         handler_md = get_handler_or_create(
-            awaitable, EventType.AdapterMessageEvent, **kwargs
+            awaitable,
+            EventType.AdapterMessageEvent,
+            **kwargs,
         )
-        new_command.init_handler_md(handler_md)
-        handler_md.event_filters.append(new_command)
+        if new_command:
+            new_command.init_handler_md(handler_md)
+            handler_md.event_filters.append(new_command)
         return awaitable
 
     return decorator
@@ -93,6 +124,7 @@ def register_custom_filter(custom_type_filter, *args, **kwargs):
         custom_type_filter: 在裸指令时为CustomFilter对象
                                         在指令组时为父指令的RegisteringCommandable对象，即self或者command_group的返回
         raise_error: 如果没有权限，是否抛出错误到消息平台，并且停止事件传播。默认为 True
+
     """
     add_to_event_filters = False
     raise_error = True
@@ -117,19 +149,20 @@ def register_custom_filter(custom_type_filter, *args, **kwargs):
     def decorator(awaitable):
         # 裸指令，子指令与指令组的区分，指令组会因为标记跳过wake。
         if (
-            not add_to_event_filters
-            and isinstance(awaitable, RegisteringCommandable)
-            or (add_to_event_filters and isinstance(awaitable, RegisteringCommandable))
-        ):
+            not add_to_event_filters and isinstance(awaitable, RegisteringCommandable)
+        ) or (add_to_event_filters and isinstance(awaitable, RegisteringCommandable)):
             # 指令组 与 根指令组，添加到本层的grouphandle中一起判断
             awaitable.parent_group.add_custom_filter(custom_filter)
         else:
             handler_md = get_handler_or_create(
-                awaitable, EventType.AdapterMessageEvent, **kwargs
+                awaitable,
+                EventType.AdapterMessageEvent,
+                **kwargs,
             )
 
             if not add_to_event_filters and not isinstance(
-                awaitable, RegisteringCommandable
+                awaitable,
+                RegisteringCommandable,
             ):
                 # 底层子指令
                 handle_full_name = get_handler_full_name(awaitable)
@@ -148,7 +181,9 @@ def register_custom_filter(custom_type_filter, *args, **kwargs):
             else:
                 # 裸指令
                 handler_md = get_handler_or_create(
-                    awaitable, EventType.AdapterMessageEvent, **kwargs
+                    awaitable,
+                    EventType.AdapterMessageEvent,
+                    **kwargs,
                 )
                 handler_md.event_filters.append(custom_filter)
 
@@ -158,26 +193,41 @@ def register_custom_filter(custom_type_filter, *args, **kwargs):
 
 
 def register_command_group(
-    command_group_name: str = None, sub_command: str = None, alias: set = None, **kwargs
+    command_group_name: str | None = None,
+    sub_command: str | None = None,
+    alias: set | None = None,
+    **kwargs,
 ):
     """注册一个 CommandGroup"""
     new_group = None
     if isinstance(command_group_name, RegisteringCommandable):
         # 子指令组
-        new_group = CommandGroupFilter(
-            sub_command, alias, parent_group=command_group_name.parent_group
-        )
-        command_group_name.parent_group.add_sub_command_filter(new_group)
+        if sub_command is None:
+            logger.warning(f"{command_group_name} 指令组的子指令组 sub_command 未指定")
+        else:
+            new_group = CommandGroupFilter(
+                sub_command,
+                alias,
+                parent_group=command_group_name.parent_group,
+            )
+            command_group_name.parent_group.add_sub_command_filter(new_group)
+    # 根指令组
+    elif command_group_name is None:
+        logger.warning("根指令组的名称未指定")
     else:
-        # 根指令组
         new_group = CommandGroupFilter(command_group_name, alias)
 
     def decorator(obj):
-        # 根指令组
-        handler_md = get_handler_or_create(obj, EventType.AdapterMessageEvent, **kwargs)
-        handler_md.event_filters.append(new_group)
+        if new_group:
+            handler_md = get_handler_or_create(
+                obj,
+                EventType.AdapterMessageEvent,
+                **kwargs,
+            )
+            handler_md.event_filters.append(new_group)
 
-        return RegisteringCommandable(new_group)
+            return RegisteringCommandable(new_group)
+        raise ValueError("注册指令组失败。")
 
     return decorator
 
@@ -185,9 +235,9 @@ def register_command_group(
 class RegisteringCommandable:
     """用于指令组级联注册"""
 
-    group: CommandGroupFilter = register_command_group
-    command: CommandFilter = register_command
-    custom_filter = register_custom_filter
+    group: Callable[..., Callable[..., RegisteringCommandable]] = register_command_group
+    command: Callable[..., Callable[..., None]] = register_command
+    custom_filter: Callable[..., Callable[..., None]] = register_custom_filter
 
     def __init__(self, parent_group: CommandGroupFilter):
         self.parent_group = parent_group
@@ -198,7 +248,9 @@ def register_event_message_type(event_message_type: EventMessageType, **kwargs):
 
     def decorator(awaitable):
         handler_md = get_handler_or_create(
-            awaitable, EventType.AdapterMessageEvent, **kwargs
+            awaitable,
+            EventType.AdapterMessageEvent,
+            **kwargs,
         )
         handler_md.event_filters.append(EventMessageTypeFilter(event_message_type))
         return awaitable
@@ -207,14 +259,15 @@ def register_event_message_type(event_message_type: EventMessageType, **kwargs):
 
 
 def register_platform_adapter_type(
-    platform_adapter_type: PlatformAdapterType, **kwargs
+    platform_adapter_type: PlatformAdapterType,
+    **kwargs,
 ):
     """注册一个 PlatformAdapterType"""
 
     def decorator(awaitable):
         handler_md = get_handler_or_create(awaitable, EventType.AdapterMessageEvent)
         handler_md.event_filters.append(
-            PlatformAdapterTypeFilter(platform_adapter_type)
+            PlatformAdapterTypeFilter(platform_adapter_type),
         )
         return awaitable
 
@@ -226,7 +279,9 @@ def register_regex(regex: str, **kwargs):
 
     def decorator(awaitable):
         handler_md = get_handler_or_create(
-            awaitable, EventType.AdapterMessageEvent, **kwargs
+            awaitable,
+            EventType.AdapterMessageEvent,
+            **kwargs,
         )
         handler_md.event_filters.append(RegexFilter(regex))
         return awaitable
@@ -240,12 +295,13 @@ def register_permission_type(permission_type: PermissionType, raise_error: bool
     Args:
         permission_type: PermissionType
         raise_error: 如果没有权限，是否抛出错误到消息平台，并且停止事件传播。默认为 True
+
     """
 
     def decorator(awaitable):
         handler_md = get_handler_or_create(awaitable, EventType.AdapterMessageEvent)
         handler_md.event_filters.append(
-            PermissionTypeFilter(permission_type, raise_error)
+            PermissionTypeFilter(permission_type, raise_error),
         )
         return awaitable
 
@@ -262,6 +318,16 @@ def register_on_astrbot_loaded(**kwargs):
     return decorator
 
 
+def register_on_platform_loaded(**kwargs):
+    """当平台加载完成时"""
+
+    def decorator(awaitable):
+        _ = get_handler_or_create(awaitable, EventType.OnPlatformLoadedEvent, **kwargs)
+        return awaitable
+
+    return decorator
+
+
 def register_on_llm_request(**kwargs):
     """当有 LLM 请求时的事件
 
@@ -275,6 +341,7 @@ def register_on_llm_request(**kwargs):
     ```
 
     请务必接收两个参数：event, request
+
     """
 
     def decorator(awaitable):
@@ -297,6 +364,7 @@ def register_on_llm_response(**kwargs):
     ```
 
     请务必接收两个参数：event, request
+
     """
 
     def decorator(awaitable):
@@ -306,7 +374,7 @@ def register_on_llm_response(**kwargs):
     return decorator
 
 
-def register_llm_tool(name: str = None):
+def register_llm_tool(name: str | None = None, **kwargs):
     """为函数调用（function-calling / tools-use）添加工具。
 
     请务必按照以下格式编写一个工具（包括函数注释，AstrBot 会尝试解析该函数注释）
@@ -316,7 +384,7 @@ def register_llm_tool(name: str = None):
     async def get_weather(event: AstrMessageEvent, location: str):
         \'\'\'获取天气信息。
 
-        Args:
+    Args:
             location(string): 地点
         \'\'\'
         # 处理逻辑
@@ -337,41 +405,121 @@ def register_llm_tool(name: str = None):
     event.stop_event()
     yield
     ```
-    """
 
+    """
     name_ = name
+    registering_agent = None
+    if kwargs.get("registering_agent"):
+        registering_agent = kwargs["registering_agent"]
 
-    def decorator(awaitable: Awaitable):
+    def decorator(awaitable: Callable[..., Awaitable[Any]]):
         llm_tool_name = name_ if name_ else awaitable.__name__
-        docstring = docstring_parser.parse(awaitable.__doc__)
+        func_doc = awaitable.__doc__ or ""
+        docstring = docstring_parser.parse(func_doc)
         args = []
         for arg in docstring.params:
-            if arg.type_name not in SUPPORTED_TYPES:
+            sub_type_name = None
+            type_name = arg.type_name
+            if not type_name:
                 raise ValueError(
-                    f"LLM 函数工具 {awaitable.__module__}_{llm_tool_name} 不支持的参数类型：{arg.type_name}"
+                    f"LLM 函数工具 {awaitable.__module__}_{llm_tool_name} 的参数 {arg.arg_name} 缺少类型注释。",
                 )
-            args.append(
-                {
-                    "type": arg.type_name,
-                    "name": arg.arg_name,
-                    "description": arg.description,
-                }
-            )
-        md = get_handler_or_create(awaitable, EventType.OnCallingFuncToolEvent)
-        llm_tools.add_func(
-            llm_tool_name, args, docstring.description.strip(), md.handler
-        )
+            # parse type_name to handle cases like "list[string]"
+            match = re.match(r"(\w+)\[(\w+)\]", type_name)
+            if match:
+                type_name = match.group(1)
+                sub_type_name = match.group(2)
+            type_name = PY_TO_JSON_TYPE.get(type_name, type_name)
+            if sub_type_name:
+                sub_type_name = PY_TO_JSON_TYPE.get(sub_type_name, sub_type_name)
+            if type_name not in SUPPORTED_TYPES or (
+                sub_type_name and sub_type_name not in SUPPORTED_TYPES
+            ):
+                raise ValueError(
+                    f"LLM 函数工具 {awaitable.__module__}_{llm_tool_name} 不支持的参数类型：{arg.type_name}",
+                )
+
+            arg_json_schema = {
+                "type": type_name,
+                "name": arg.arg_name,
+                "description": arg.description,
+            }
+            if sub_type_name:
+                if type_name == "array":
+                    arg_json_schema["items"] = {"type": sub_type_name}
+            args.append(arg_json_schema)
+
+        if not registering_agent:
+            doc_desc = docstring.description.strip() if docstring.description else ""
+            md = get_handler_or_create(awaitable, EventType.OnCallingFuncToolEvent)
+            llm_tools.add_func(llm_tool_name, args, doc_desc, md.handler)
+        else:
+            assert isinstance(registering_agent, RegisteringAgent)
+            # print(f"Registering tool {llm_tool_name} for agent", registering_agent._agent.name)
+            if registering_agent._agent.tools is None:
+                registering_agent._agent.tools = []
+
+            desc = docstring.description.strip() if docstring.description else ""
+            tool = llm_tools.spec_to_func(llm_tool_name, args, desc, awaitable)
+            registering_agent._agent.tools.append(tool)
+
         return awaitable
 
     return decorator
 
 
+class RegisteringAgent:
+    """用于 Agent 注册"""
+
+    def llm_tool(self, *args, **kwargs):
+        kwargs["registering_agent"] = self
+        return register_llm_tool(*args, **kwargs)
+
+    def __init__(self, agent: Agent[AstrAgentContext]):
+        self._agent = agent
+
+
+def register_agent(
+    name: str,
+    instruction: str,
+    tools: list[str | FunctionTool] | None = None,
+    run_hooks: BaseAgentRunHooks[AstrAgentContext] | None = None,
+):
+    """注册一个 Agent
+
+    Args:
+        name: Agent 的名称
+        instruction: Agent 的指令
+        tools: Agent 使用的工具列表
+        run_hooks: Agent 运行时的钩子函数
+
+    """
+    tools_ = tools or []
+
+    def decorator(awaitable: Callable[..., Awaitable[Any]]):
+        AstrAgent = Agent[AstrAgentContext]
+        agent = AstrAgent(
+            name=name,
+            instructions=instruction,
+            tools=tools_,
+            run_hooks=run_hooks or BaseAgentRunHooks[AstrAgentContext](),
+        )
+        handoff_tool = HandoffTool(agent=agent)
+        handoff_tool.handler = awaitable
+        llm_tools.func_list.append(handoff_tool)
+        return RegisteringAgent(agent)
+
+    return decorator
+
+
 def register_on_decorating_result(**kwargs):
     """在发送消息前的事件"""
 
     def decorator(awaitable):
         _ = get_handler_or_create(
-            awaitable, EventType.OnDecoratingResultEvent, **kwargs
+            awaitable,
+            EventType.OnDecoratingResultEvent,
+            **kwargs,
         )
         return awaitable
 
@@ -383,7 +531,9 @@ def register_after_message_sent(**kwargs):
 
     def decorator(awaitable):
         _ = get_handler_or_create(
-            awaitable, EventType.OnAfterMessageSentEvent, **kwargs
+            awaitable,
+            EventType.OnAfterMessageSentEvent,
+            **kwargs,
         )
         return awaitable