From b8fb0bee245a9ac4778eb3f540337ae6b6087ce6 Mon Sep 17 00:00:00 2001 From: Soulter <905617992@qq.com> Date: Wed, 8 Jan 2025 02:10:27 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=88=9D=E6=AD=A5=E5=AE=9E=E7=8E=B0?= =?UTF-8?q?=E4=BB=A3=E7=A0=81=E6=89=A7=E8=A1=8C=E5=99=A8=20#210?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- astrbot/core/provider/sources/dify_source.py | 1 - astrbot/core/star/register/star_handler.py | 2 +- packages/python_interpreter/main.py | 219 +++++++++++++++++++ packages/python_interpreter/requirements.txt | 0 packages/python_interpreter/shared/api.py | 13 ++ requirements.txt | 3 +- 6 files changed, 235 insertions(+), 3 deletions(-) create mode 100644 packages/python_interpreter/main.py create mode 100644 packages/python_interpreter/requirements.txt create mode 100644 packages/python_interpreter/shared/api.py diff --git a/astrbot/core/provider/sources/dify_source.py b/astrbot/core/provider/sources/dify_source.py index 64755c968..a8b2f73bd 100644 --- a/astrbot/core/provider/sources/dify_source.py +++ b/astrbot/core/provider/sources/dify_source.py @@ -1,4 +1,3 @@ -import base64 from typing import List from .. import Provider from ..entites import LLMResponse diff --git a/astrbot/core/star/register/star_handler.py b/astrbot/core/star/register/star_handler.py index a46c05ae3..f0a86b0c0 100644 --- a/astrbot/core/star/register/star_handler.py +++ b/astrbot/core/star/register/star_handler.py @@ -185,7 +185,7 @@ def register_llm_tool(name: str = None): "description": arg.description }) md = get_handler_or_create(awaitable, EventType.OnCallingFuncToolEvent) - llm_tools.add_func(llm_tool_name, args, docstring.short_description, md.handler) + llm_tools.add_func(llm_tool_name, args, docstring.description, md.handler) logger.debug(f"LLM 函数工具 {llm_tool_name} 已注册") return awaitable diff --git a/packages/python_interpreter/main.py b/packages/python_interpreter/main.py new file mode 100644 index 000000000..394ae6276 --- /dev/null +++ b/packages/python_interpreter/main.py @@ -0,0 +1,219 @@ +import os +import aiohttp +import uuid +import asyncio +import re +import astrbot.api.star as star +import aiodocker +from astrbot.api.event import AstrMessageEvent +from astrbot.api import llm_tool, logger +from astrbot.api.message_components import Image + +PROMPT = """ +## Task +You need to generate python codes to solve user's problem: {prompt} + +{extra_input} + +## Limit +1. The python libraries you can use include python standard libraries and `Pillow`, `requests`, `numpy`, `matplotlib`. +2. You must not generate malicious code. +3. You can only output text, image. For Image, you need save it to `output` folder. +4. Use given `shared.api` package to output the result. +5. Your must only output the code, do not output the result of the code and other any information. +6. The output language is same as the user's input language. + +## Example +1. The user's problem is: `please solve the fabonacci sequence problem.` +Output: +```python +from shared.api import send_text, send_image + +def fabonacci(n): + if n <= 1: + return n + else: + return fabonacci(n-1) + fabonacci(n-2) + +result = fabonacci(10) +# introduce the fabonacci sequence briefly +send_text("The fabonacci sequence is a series of numbers in which each number is the sum of the two preceding ones, starting from 0 and 1.") +send_text("Let's calculate the fabonacci sequence of 10: " + result) # send_text is a function to send pure text to user +``` + +2. The user's problem is: `please draw a sin(x) function.` +Output: +```python +from shared.api import send_text, send_image +import numpy as np +import matplotlib.pyplot as plt + +x = np.linspace(0, 2*np.pi, 100) +y = np.sin(x) +plt.plot(x, y) +plt.savefig("output/sin_x.png") +send_text("The sin(x) is a periodic function with a period of 2π, and the value range is [-1, 1]. The following is the image of sin(x).") # introduce the sin(x) function briefly +send_image("output/sin_x.png") # send_image is a function to send image to user +send_text("If you need more information, please let me know :)") +``` + +{extra_prompt} +""" + +@star.register(name="astrbot-python-interpreter", desc="Python 代码执行器", author="Soulter", version="0.0.1") +class Main(star.Star): + '''基于 Docker 沙箱的 Python 代码执行器''' + def __init__(self, context: star.Context) -> None: + self.context = context + self.curr_dir = os.path.dirname(os.path.abspath(__file__)) + self.workplace_path = os.path.join(self.curr_dir, "workplace") + self.shared_path = os.path.join(self.curr_dir, "shared") + os.makedirs(self.workplace_path, exist_ok=True) + + async def gen_magic_code(self) -> str: + return uuid.uuid4().hex[:8] + + async def download_image(self, image_url: str, workplace_path: str, filename: str) -> str: + '''Download image from url to workplace_path''' + async with aiohttp.ClientSession() as session: + async with session.get(image_url) as resp: + if resp.status != 200: + return "" + image_path = os.path.join(workplace_path, f"{filename}.jpg") + with open(image_path, 'wb') as f: + f.write(await resp.read()) + return f"{filename}.jpg" + + async def tidy_code(self, code: str) -> str: + '''Tidy the code''' + pattern = r"```(?:py|python)?\n(.*?)\n```" + match = re.search(pattern, code, re.DOTALL) + if match is None: + raise ValueError("The code is not in the code block.") + return match.group(1) + + @llm_tool("python_interpreter") + async def python_interpreter(self, event: AstrMessageEvent): + '''Use this tool only if user really want to solve a complex problem and the problem can be solved very well by Python code. For example, user can use this tool to solve a math problem, edit Image, etc. + ''' + plain_text = event.message_str + + # 创建必要的工作目录和幻术码 + magic_code = await self.gen_magic_code() + workplace_path = os.path.join(self.workplace_path, magic_code) + output_path = os.path.join(workplace_path, "output") + os.makedirs(workplace_path, exist_ok=True) + os.makedirs(output_path, exist_ok=True) + + # 图片 + images = [] + idx = 1 + for comp in event.message_obj.message: + if isinstance(comp, Image): + image_url = comp.url if comp.url else comp.file + if image_url.startswith("http"): + image_path = await self.download_image(image_url, workplace_path, f"img_{idx}") + if image_path: + images.append(image_path) + idx += 1 + + obs = "" + n = 5 + + for i in range(n): + if i > 0: + logger.info(f"Try {i+1}/{n}") + + PROMPT_ = PROMPT.format( + prompt=plain_text, + extra_input=f"User provided images: {images}", + extra_prompt=obs, + ) + provider = self.context.get_using_provider() + llm_response = await provider.text_chat(prompt=PROMPT_, session_id=event.session_id) + + logger.debug("code interpreter llm gened code:" + llm_response.completion_text) + + # 整理代码并保存 + code_clean = await self.tidy_code(llm_response.completion_text) + with open(os.path.join(workplace_path, "exec.py"), "w") as f: + f.write(code_clean) + + # 启动容器 + docker = aiodocker.Docker() + + # 检查有没有image + try: + await docker.images.get("cjie.eu.org/soulter/astrbot-code-interpreter-sandbox") + except aiodocker.exceptions.DockerError: + # 拉取镜像 + logger.debug("Pulling image soulter/astrbot-code-interpreter-sandbo...") + await docker.images.pull("cjie.eu.org/soulter/astrbot-code-interpreter-sandbox") + + yield event.plain_result(f"使用沙箱执行代码中，请稍等...(尝试次数: {i+1}/{n})") + + container = await docker.containers.run({ + "Image": "cjie.eu.org/soulter/astrbot-code-interpreter-sandbox", + "Cmd": ["python", "exec.py"], + "Memory": 512 * 1024 * 1024, + "NanoCPUs": 1000000000, + "HostConfig": { + "Binds": [ + f"{self.shared_path}:/astrbot_sandbox/shared:ro", + f"{output_path}:/astrbot_sandbox/output:rw", + f"{workplace_path}:/astrbot_sandbox:rw", + ] + }, + "Env": [ + f"MAGIC_CODE={magic_code}" + ], + "AutoRemove": True + }) + + logger.debug(f"Container {container.id} created.") + logs = await self.run_container(container) + + logger.debug(f"Container {container.id} finished.") + logger.debug(f"Container {container.id} logs: {logs}") + + # 发送结果 + pattern = r"\[ASTRBOT_(TEXT|IMAGE)_OUTPUT#\w+\]: (.*)" + ok = False + traceback = "" + for idx, log in enumerate(logs): + match = re.match(pattern, log) + if match: + ok = True + if match.group(1) == "TEXT": + yield event.plain_result(match.group(2)) + elif match.group(1) == "IMAGE": + image_path = os.path.join(workplace_path, match.group(2)) + logger.debug(f"Sending image: {image_path}") + yield event.image_result(image_path) + elif "Traceback (most recent call last)" in log: + traceback = "\n".join(logs[idx:]) + + if not ok: + if traceback: + obs = f"## Observation\When execute the code: ```python\n{code_clean}\n```\n\n Error occured:\n\n{traceback}\n Need to improve/fix the code." + else: + logger.warning(f"未从沙箱输出中捕获到合法的输出。沙箱输出日志: {logs}") + break + else: + return + + yield event.plain_result("经过多次尝试后，未从沙箱输出中捕获到合法的输出，请更换问法或者查看日志。") + + + async def run_container(self, container: aiodocker.docker.DockerContainer, timeout: int = 20) -> list[str]: + '''Run the container and get the output''' + try: + await container.wait(timeout=timeout) + logs = await container.log(stdout=True, stderr=True) + return logs + except asyncio.TimeoutError: + logger.warning(f"Container {container.id} timeout.") + await container.kill() + return f"Container has been killed due to timeout ({timeout}s)." + finally: + await container.delete() \ No newline at end of file diff --git a/packages/python_interpreter/requirements.txt b/packages/python_interpreter/requirements.txt new file mode 100644 index 000000000..e69de29bb diff --git a/packages/python_interpreter/shared/api.py b/packages/python_interpreter/shared/api.py new file mode 100644 index 000000000..320936e01 --- /dev/null +++ b/packages/python_interpreter/shared/api.py @@ -0,0 +1,13 @@ +import os + +def _get_magic_code(): + '''防止注入攻击''' + return os.getenv("MAGIC_CODE") + +def send_text(text: str): + print(f"[ASTRBOT_TEXT_OUTPUT#{_get_magic_code()}]: {text}") + +def send_image(image_path: str): + if not os.path.exists(image_path): + raise Exception(f"Image file not found: {image_path}") + print(f"[ASTRBOT_IMAGE_OUTPUT#{_get_magic_code()}]: {image_path}") \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index c02fc9b0d..965651870 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,4 +15,5 @@ colorlog aiocqhttp pyjwt apscheduler -docstring_parser \ No newline at end of file +docstring_parser +aiodocker \ No newline at end of file