mirror of
https://github.com/LiteyukiStudio/nonebot-plugin-marshoai.git
synced 2025-02-07 09:06:10 +08:00
92 lines
2.2 KiB
Python
92 lines
2.2 KiB
Python
from typing import Generic, TypeVar
|
||
|
||
from nonebot import logger
|
||
|
||
from ..typing import FUNCTION_CALL_FUNC
|
||
from .params import P
|
||
|
||
F = TypeVar("F", bound=FUNCTION_CALL_FUNC)
|
||
|
||
|
||
class Caller(Generic[P]):
|
||
def __init__(self, name: str | None = None, description: str | None = None):
|
||
self._name = name
|
||
self._description = description
|
||
self._parameters: dict[str, P] = {}
|
||
self.func: FUNCTION_CALL_FUNC | None = None
|
||
|
||
def params(self, **kwargs: P) -> "Caller":
|
||
"""设置多个函数参数
|
||
Args:
|
||
**kwargs: 参数字典
|
||
Returns:
|
||
Caller: Caller对象
|
||
"""
|
||
self._parameters.update(kwargs)
|
||
return self
|
||
|
||
def param(self, name: str, param: P) -> "Caller":
|
||
"""设置一个函数参数
|
||
|
||
Args:
|
||
name (str): 参数名
|
||
param (P): 参数对象
|
||
|
||
Returns:
|
||
Caller: Caller对象
|
||
"""
|
||
self._parameters[name] = param
|
||
return self
|
||
|
||
def name(self, name: str) -> "Caller":
|
||
"""设置函数名称
|
||
|
||
Args:
|
||
name (str): 函数名称
|
||
|
||
Returns:
|
||
Caller: Caller对象
|
||
"""
|
||
self._name = name
|
||
return self
|
||
|
||
def description(self, description: str) -> "Caller":
|
||
"""设置函数描述
|
||
|
||
Args:
|
||
description (str): 函数描述
|
||
|
||
Returns:
|
||
Caller: Caller对象
|
||
"""
|
||
self._description = description
|
||
return self
|
||
|
||
def __call__(self, func: F) -> F:
|
||
"""装饰函数，注册为一个可被AI调用的function call函数
|
||
|
||
Args:
|
||
func (F): 函数对象
|
||
|
||
Returns:
|
||
F: 函数对象
|
||
"""
|
||
logger.opt(colors=True).info(
|
||
f"<y>加载函数 {func.__name__} {self._description}</y>"
|
||
)
|
||
self.func = func
|
||
return func
|
||
|
||
|
||
def on_function_call(name: str | None = None, description: str | None = None) -> Caller:
|
||
"""返回一个Caller类，可用于装饰一个函数，使其注册为一个可被AI调用的function call函数
|
||
|
||
Args:
|
||
description: 函数描述，若为None则从函数的docstring中获取
|
||
|
||
Returns:
|
||
Caller: Caller对象
|
||
"""
|
||
|
||
return Caller(name=name, description=description)
|