aient 1.0.29__py3-none-any.whl

aient/plugins/__init__.py

@@ -0,0 +1,32 @@
+import os
+import pkgutil
+import importlib
+
+# 首先导入registry，因为其他模块中的装饰器依赖它
+from .registry import registry, register_tool, register_agent
+
+# 自动导入当前目录下所有的插件模块
+excluded_modules = ['config', 'registry', '__init__']
+current_dir = os.path.dirname(__file__)
+
+# 先导入所有模块，确保装饰器被执行
+for _, module_name, _ in pkgutil.iter_modules([current_dir]):
+    if module_name not in excluded_modules:
+        importlib.import_module(f'.{module_name}', package=__name__)
+
+# 然后从config导入必要的定义
+from .config import *
+
+# 确保将所有工具函数添加到全局名称空间
+for tool_name, tool_func in registry.tools.items():
+    globals()[tool_name] = tool_func
+
+__all__ = [
+    'PLUGINS',
+    'function_call_list',
+    'get_tools_result_async',
+    'registry',
+    'register_tool',
+    'register_agent',
+    'update_tools_config',
+] + list(registry.tools.keys())

aient/plugins/arXiv.py

@@ -0,0 +1,48 @@
+import requests
+
+from ..utils.scripts import Document_extract
+from .registry import register_tool
+
+@register_tool()
+async def download_read_arxiv_pdf(arxiv_id: str) -> str:
+    """
+    下载指定arXiv ID的论文PDF并提取其内容。
+
+    此函数会下载arXiv上的论文PDF文件，保存到指定路径，
+    然后使用文档提取工具读取其内容。
+
+    Args:
+        arxiv_id: arXiv论文的ID，例如'2305.12345'
+
+    Returns:
+        提取的论文内容文本或失败消息
+    """
+    # 构造下载PDF的URL
+    url = f'https://arxiv.org/pdf/{arxiv_id}.pdf'
+
+    # 发送HTTP GET请求
+    response = requests.get(url)
+
+    # 检查是否成功获取内容
+    if response.status_code == 200:
+        # 将PDF内容写入文件
+        save_path = "paper.pdf"
+        with open(save_path, 'wb') as file:
+            file.write(response.content)
+        print(f'PDF下载成功，保存路径: {save_path}')
+        return await Document_extract(None, save_path)
+    else:
+        print(f'下载失败，状态码: {response.status_code}')
+        return "文件下载失败"
+
+if __name__ == '__main__':
+    # 示例使用
+    arxiv_id = '2305.12345'  # 替换为实际的arXiv ID
+
+    # 测试下载功能
+    # print(download_read_arxiv_pdf(arxiv_id))
+
+    # 测试函数转换为JSON
+    # json_result = function_to_json(download_read_arxiv_pdf)
+    # import json
+    # print(json.dumps(json_result, indent=2, ensure_ascii=False))

aient/plugins/config.py

@@ -0,0 +1,178 @@
+import os
+import json
+import inspect
+
+from .registry import registry
+from ..utils.scripts import cut_message, safe_get
+from ..utils.prompt import search_key_word_prompt, arxiv_doc_user_prompt
+
+async def get_tools_result_async(function_call_name, function_full_response, function_call_max_tokens, engine, robot, api_key, api_url, use_plugins, model, add_message, convo_id, language):
+    function_response = ""
+    if function_call_name in registry.tools:
+        function_to_call = registry.tools[function_call_name]
+        function_args = registry.tools_info[function_call_name].args
+        # required_args = registry.tools_info[function_call_name].required
+        if function_args:
+            arg = function_args[0]
+        else:
+            arg = None
+    if function_call_name == "get_search_results":
+        prompt = json.loads(function_full_response)["query"]
+        yield "message_search_stage_1"
+        llm = robot(api_key=api_key, api_url=api_url.source_api_url, engine=engine, use_plugins=use_plugins)
+        keywords = (await llm.ask_async(search_key_word_prompt.format(source=prompt), model=model)).split("\n")
+        print("keywords", keywords)
+        keywords = [item.replace("三行关键词是：", "") for item in keywords if "\\x" not in item if item != ""]
+        keywords = [prompt] + keywords
+        keywords = keywords[:3]
+        print("select keywords", keywords)
+        async for chunk in function_to_call(keywords):
+            if type(chunk) == str:
+                yield chunk
+            else:
+                function_response = "\n\n".join(chunk)
+            # function_response = yield chunk
+        # function_response = yield from eval(function_call_name)(prompt, keywords)
+        function_call_max_tokens = 32000
+        function_response, text_len = cut_message(function_response, function_call_max_tokens, engine)
+        if function_response:
+            function_response = (
+                f"You need to response the following question: {prompt}. Search results is provided inside <Search_results></Search_results> XML tags. Your task is to think about the question step by step and then answer the above question in {language} based on the Search results provided. Please response in {language} and adopt a style that is logical, in-depth, and detailed. Note: In order to make the answer appear highly professional, you should be an expert in textual analysis, aiming to make the answer precise and comprehensive. Directly response markdown format, without using markdown code blocks. For each sentence quoting search results, a markdown ordered superscript number url link must be used to indicate the source, e.g., [¹](https://www.example.com)"
+                "Here is the Search results, inside <Search_results></Search_results> XML tags:"
+                "<Search_results>"
+                "{}"
+                "</Search_results>"
+            ).format(function_response)
+        else:
+            function_response = "无法找到相关信息，停止使用 tools"
+        # user_prompt = f"You need to response the following question: {prompt}. Search results is provided inside <Search_results></Search_results> XML tags. Your task is to think about the question step by step and then answer the above question in {config.language} based on the Search results provided. Please response in {config.language} and adopt a style that is logical, in-depth, and detailed. Note: In order to make the answer appear highly professional, you should be an expert in textual analysis, aiming to make the answer precise and comprehensive. Directly response markdown format, without using markdown code blocks"
+        # self.add_to_conversation(user_prompt, "user", convo_id=convo_id)
+    elif arg: # generate_image get_url_content run_python_script
+        prompt = safe_get(json.loads(function_full_response), arg, default=".")
+        if inspect.iscoroutinefunction(function_to_call):
+            function_response = await function_to_call(prompt)
+        else:
+            function_response = function_to_call(prompt)
+        function_response, text_len = cut_message(function_response, function_call_max_tokens, engine)
+    else: # get_date_time_weekday
+        if inspect.iscoroutinefunction(function_to_call):
+            function_response = await function_to_call()
+        else:
+            function_response = function_to_call()
+        function_response, text_len = cut_message(function_response, function_call_max_tokens, engine)
+
+    if function_call_name == "download_read_arxiv_pdf":
+        add_message(arxiv_doc_user_prompt, "user", convo_id=convo_id)
+
+    function_response = (
+        f"function_response:{function_response}"
+    )
+    yield function_response
+    # return function_response
+
+def function_to_json(func) -> dict:
+    """
+    将Python函数转换为JSON可序列化的字典，描述函数的签名，包括名称、描述和参数。
+
+    Args:
+        func: 要转换的函数
+
+    Returns:
+        表示函数签名的JSON格式字典
+    """
+    type_map = {
+        str: "string",
+        int: "integer",
+        float: "number",
+        bool: "boolean",
+        type(None): "null",
+    }
+
+    try:
+        signature = inspect.signature(func)
+    except ValueError as e:
+        raise ValueError(f"获取函数{func.__name__}签名失败: {str(e)}")
+
+    parameters = {}
+    for param in signature.parameters.values():
+        try:
+            if param.annotation == inspect._empty:
+                parameters[param.name] = {"type": "string"}
+            else:
+                parameters[param.name] = {"type": type_map.get(param.annotation, "string")}
+        except KeyError as e:
+            raise KeyError(f"未知类型注解 {param.annotation} 用于参数 {param.name}: {str(e)}")
+
+    required = [
+        param.name
+        for param in signature.parameters.values()
+        if param.default == inspect._empty
+    ]
+
+    return {
+        "name": func.__name__,
+        "description": func.__doc__ or "",
+        "parameters": {
+            "type": "object",
+            "properties": parameters,
+            "required": required,
+        },
+    }
+
+def gpt2claude_tools_json(json_dict):
+    import copy
+    json_dict = copy.deepcopy(json_dict)
+    keys_to_change = {
+        "parameters": "input_schema",
+    }
+    for old_key, new_key in keys_to_change.items():
+        if old_key in json_dict:
+            if new_key:
+                json_dict[new_key] = json_dict.pop(old_key)
+            else:
+                json_dict.pop(old_key)
+        else:
+            if new_key and "description" in json_dict.keys():
+                json_dict[new_key] = {
+                    "type": "object",
+                    "properties": {}
+                }
+    if "tools" in json_dict.keys():
+        json_dict["tool_choice"] = {
+            "type": "auto"
+        }
+    return json_dict
+
+# print("registry.tools", json.dumps(registry.tools_info.get('get_time', {}), indent=4, ensure_ascii=False))
+# print("registry.tools", json.dumps(registry.tools_info['run_python_script'].to_dict(), indent=4, ensure_ascii=False))
+
+# 修改PLUGINS定义，使用registry中的工具
+def get_plugins():
+    return {
+        tool_name: (os.environ.get(tool_name, "False") == "False") == False
+        for tool_name in registry.tools.keys()
+    }
+
+# 修改function_call_list定义，使用registry中的工具
+def get_function_call_list():
+    function_list = {}
+    for tool_name, tool_func in registry.tools.items():
+        function_list[tool_name] = function_to_json(tool_func)
+    return function_list
+
+def get_claude_tools_list():
+    function_list = get_function_call_list()
+    return {f"{key}": gpt2claude_tools_json(function_list[key]) for key in function_list.keys()}
+
+# 初始化默认配置
+PLUGINS = get_plugins()
+function_call_list = get_function_call_list()
+claude_tools_list = get_claude_tools_list()
+
+# 动态更新工具函数配置
+def update_tools_config():
+    global PLUGINS, function_call_list, claude_tools_list
+    PLUGINS = get_plugins()
+    function_call_list = get_function_call_list()
+    claude_tools_list = get_claude_tools_list()
+    return PLUGINS, function_call_list, claude_tools_list

aient/plugins/image.py

@@ -0,0 +1,72 @@
+import os
+import requests
+import json
+from ..models.base import BaseLLM
+from .registry import register_tool
+
+API = os.environ.get('API', None)
+API_URL = os.environ.get('API_URL', None)
+
+class dalle3(BaseLLM):
+    def __init__(
+        self,
+        api_key: str,
+        api_url: str = (os.environ.get("API_URL") or "https://api.openai.com/v1/images/generations"),
+        timeout: float = 20,
+    ):
+        super().__init__(api_key, api_url=api_url, timeout=timeout)
+        self.engine: str = "dall-e-3"
+
+    def generate(
+        self,
+        prompt: str,
+        model: str = "",
+        **kwargs,
+    ):
+        url = self.api_url.image_url
+        headers = {"Authorization": f"Bearer {kwargs.get('api_key', self.api_key)}"}
+
+        json_post = {
+                "model": os.environ.get("IMAGE_MODEL_NAME") or model or self.engine,
+                "prompt": prompt,
+                "n": 1,
+                "size": "1024x1024",
+        }
+        try:
+            response = self.session.post(
+                url,
+                headers=headers,
+                json=json_post,
+                timeout=kwargs.get("timeout", self.timeout),
+                stream=True,
+            )
+        except ConnectionError:
+            print("连接错误，请检查服务器状态或网络连接。")
+            return
+        except requests.exceptions.ReadTimeout:
+            print("请求超时，请检查网络连接或增加超时时间。{e}")
+            return
+        except Exception as e:
+            print(f"发生了未预料的错误: {e}")
+            return
+
+        if response.status_code != 200:
+            raise Exception(f"{response.status_code} {response.reason} {response.text}")
+        json_data = json.loads(response.text)
+        url = json_data["data"][0]["url"]
+        yield url
+
+@register_tool()
+def generate_image(text):
+    """
+    生成图像
+
+    参数:
+        text: 描述图像的文本
+
+    返回:
+        图像的URL
+    """
+    dallbot = dalle3(api_key=f"{API}")
+    for data in dallbot.generate(text):
+        return data

aient/plugins/registry.py

@@ -0,0 +1,116 @@
+from typing import Callable, Dict, Literal, List, Optional
+from dataclasses import dataclass, asdict
+import inspect
+
+@dataclass
+class FunctionInfo:
+    name: str
+    func: Callable
+    args: List[str]
+    docstring: Optional[str]
+    body: str
+    return_type: Optional[str]
+    def to_dict(self) -> dict:
+        # using asdict, but exclude func field because it cannot be serialized
+        d = asdict(self)
+        d.pop('func')  # remove func field
+        return d
+
+    @classmethod
+    def from_dict(cls, data: dict) -> 'FunctionInfo':
+        # if you need to create an object from a dictionary
+        if 'func' not in data:
+            data['func'] = None  # or other default value
+        return cls(**data)
+
+class Registry:
+    _instance = None
+    _registry: Dict[str, Dict[str, Callable]] = {
+        "tools": {},
+        "agents": {}
+    }
+    _registry_info: Dict[str, Dict[str, FunctionInfo]] = {
+        "tools": {},
+        "agents": {}
+    }
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def register(self,
+                type: Literal["tool", "agent"],
+                name: str = None):
+        """
+        统一的注册装饰器
+        Args:
+            type: 注册类型，"tool" 或 "agent"
+            name: 可选的注册名称
+        """
+        def decorator(func: Callable):
+            nonlocal name
+            if name is None:
+                name = func.__name__
+                # if type == "agent" and name.startswith('get_'):
+                #     name = name[4:]  # 对 agent 移除 'get_' 前缀
+
+            # 获取函数信息
+            signature = inspect.signature(func)
+            args = list(signature.parameters.keys())
+            docstring = inspect.getdoc(func)
+
+            # 获取函数体
+            source_lines = inspect.getsource(func)
+            # 移除装饰器和函数定义行
+            body_lines = source_lines.split('\n')[1:]  # 跳过装饰器行
+            while body_lines and (body_lines[0].strip().startswith('@') or 'def ' in body_lines[0]):
+                body_lines = body_lines[1:]
+            body = '\n'.join(body_lines)
+
+            # 获取返回类型提示
+            return_type = None
+            if signature.return_annotation != inspect.Signature.empty:
+                return_type = str(signature.return_annotation)
+
+            # 创建函数信息对象
+            func_info = FunctionInfo(
+                name=name,
+                func=func,
+                args=args,
+                docstring=docstring,
+                body=body,
+                return_type=return_type
+            )
+
+            registry_type = f"{type}s"
+            self._registry[registry_type][name] = func
+            self._registry_info[registry_type][name] = func_info
+            return func
+        return decorator
+
+    @property
+    def tools(self) -> Dict[str, Callable]:
+        return self._registry["tools"]
+
+    @property
+    def agents(self) -> Dict[str, Callable]:
+        return self._registry["agents"]
+
+    @property
+    def tools_info(self) -> Dict[str, FunctionInfo]:
+        return self._registry_info["tools"]
+
+    @property
+    def agents_info(self) -> Dict[str, FunctionInfo]:
+        return self._registry_info["agents"]
+
+# 创建全局实例
+registry = Registry()
+
+# 便捷的注册函数
+def register_tool(name: str = None):
+    return registry.register(type="tool", name=name)
+
+def register_agent(name: str = None):
+    return registry.register(type="agent", name=name)

aient/plugins/run_python.py

@@ -0,0 +1,156 @@
+import os
+import ast
+import asyncio
+import logging
+import tempfile
+from .registry import register_tool
+
+def get_dangerous_attributes(node):
+    # 简单的代码审查，检查是否包含某些危险关键词
+    dangerous_keywords = ['os', 'subprocess', 'sys', 'import', 'eval', 'exec', 'open']
+    if isinstance(node, ast.Name):
+        return node.id in dangerous_keywords
+    elif isinstance(node, ast.Attribute):
+        return node.attr in dangerous_keywords
+    return False
+
+def check_code_safety(code):
+    try:
+        # 解析代码为 AST
+        tree = ast.parse(code)
+
+        # 检查所有节点
+        for node in ast.walk(tree):
+            # 检查危险属性访问
+            if get_dangerous_attributes(node):
+                return False
+
+            # 检查危险的调用
+            if isinstance(node, ast.Call):
+                if isinstance(node.func, (ast.Name, ast.Attribute)):
+                    if get_dangerous_attributes(node.func):
+                        return False
+
+            # 检查字符串编码/解码操作
+            if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
+                if node.func.attr in ('encode', 'decode'):
+                    return False
+
+        return True
+    except SyntaxError:
+        return False
+
+@register_tool()
+async def run_python_script(code):
+    """
+    执行 Python 代码
+
+    参数:
+        code: 要执行的 Python 代码字符串
+
+    返回:
+        执行结果字符串
+    """
+
+    timeout = 10
+    # 检查代码安全性
+    if not check_code_safety(code):
+        return "Code contains potentially dangerous operations.\n\n"
+
+    # 添加一段捕获代码，确保最后表达式的值会被输出
+    # 这种方式比 ast 解析更可靠
+    wrapper_code = """
+import sys
+_result = None
+
+def _capture_last_result(code_to_run):
+    global _result
+    namespace = {{}}
+    exec(code_to_run, namespace)
+    if "_last_expr" in namespace:
+        _result = namespace["_last_expr"]
+
+# 用户代码
+_user_code = '''
+{}
+'''
+
+# 处理用户代码，尝试提取最后一个表达式
+lines = _user_code.strip().split('\\n')
+if lines:
+    # 检查最后一行是否是表达式
+    last_line = lines[-1].strip()
+    if last_line and not last_line.startswith(('def ', 'class ', 'if ', 'for ', 'while ', 'try:', 'with ')):
+        if not any(last_line.startswith(kw) for kw in ['return', 'print', 'raise', 'assert', 'import', 'from ']):
+            if not last_line.endswith(':') and not last_line.endswith('='):
+                # 可能是表达式，修改它
+                lines[-1] = "_last_expr = " + last_line
+                _user_code = '\\n'.join(lines)
+
+_capture_last_result(_user_code)
+
+# 输出结果
+if _result is not None:
+    print("\\nResult:", repr(_result))
+""".format(code)
+
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as temp_file:
+        temp_file.write(wrapper_code)
+        temp_file_name = temp_file.name
+
+    try:
+        process = await asyncio.create_subprocess_exec(
+            'python', temp_file_name,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE
+        )
+
+        try:
+            stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout)
+            stdout = stdout.decode()
+            stderr = stderr.decode()
+            return_code = process.returncode
+        except asyncio.TimeoutError:
+            # 使用 SIGTERM 信号终止进程
+            process.terminate()
+            await asyncio.sleep(0.1)  # 给进程一点时间来终止
+            if process.returncode is None:
+                # 如果进程还没有终止，使用 SIGKILL
+                process.kill()
+            return "Process execution timed out."
+
+        mess = (
+            f"Execution result:\n{stdout}\n",
+            f"Stderr:\n{stderr}\n" if stderr else "",
+            f"Return Code: {return_code}\n" if return_code else "",
+        )
+        mess = "".join(mess)
+        return mess
+
+    except Exception as e:
+        logging.error(f"Error executing code: {str(e)}")
+        return f"Error: {str(e)}"
+
+    finally:
+        try:
+            os.unlink(temp_file_name)
+        except Exception as e:
+            logging.error(f"Error deleting temporary file: {str(e)}")
+
+# 使用示例
+async def main():
+    code = """
+print("Hello, World!")
+"""
+    code = """
+def add(a, b):
+    return a + b
+
+result = add(5, 3)
+print(result)
+    """
+    result = await run_python_script(code)
+    print(result)
+
+if __name__ == "__main__":
+    asyncio.run(main())

aient/plugins/today.py

@@ -0,0 +1,19 @@
+import pytz
+import datetime
+
+from .registry import register_tool
+
+# Plugins 获取日期时间
+@register_tool()
+def get_date_time_weekday():
+    """
+    获取当前日期时间及星期几
+
+    返回:
+        包含当前日期时间及星期几的字符串
+    """
+    tz = pytz.timezone('Asia/Shanghai')  # 为东八区设置时区
+    now = datetime.datetime.now(tz)  # 获取东八区当前时间
+    weekday = now.weekday()
+    weekday_str = ['星期一', '星期二', '星期三', '星期四', '星期五', '星期六', '星期日'][weekday]
+    return "今天是：" + str(now.date()) + "，现在的时间是：" + str(now.time())[:-7] + "，" + weekday_str