laoxong/AstrBot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

style: format codes

Browse Source
This commit is contained in:
Soulter
2025-03-03 13:30:55 +08:00
parent d305211a36
commit 49bba9bf98
148 changed files with 6135 additions and 4401 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/python_interpreter/main.py
+140 -106
View File
@@ -86,19 +86,26 @@ send_text("If you need more information, please let me know :)")
DEFAULT_CONFIG = {
"sandbox": {
"image": "soulter/astrbot-code-interpreter-sandbox",
"docker_mirror": "", # cjie.eu.org
"docker_mirror": "", # cjie.eu.org
},
"docker_host_astrbot_abs_path": ""
"docker_host_astrbot_abs_path": "",
}
PATH = "data/config/python_interpreter.json"
@star.register(name="astrbot-python-interpreter", desc="Python 代码执行器", author="Soulter", version="0.0.1")
@star.register(
name="astrbot-python-interpreter",
desc="Python 代码执行器",
author="Soulter",
version="0.0.1",
)
class Main(star.Star):
'''基于 Docker 沙箱的 Python 代码执行器'''
"""基于 Docker 沙箱的 Python 代码执行器"""
def __init__(self, context: star.Context) -> None:
self.context = context
self.curr_dir = os.path.dirname(os.path.abspath(__file__))
self.shared_path = os.path.join("data", "py_interpreter_shared")
if not os.path.exists(self.shared_path):
# 复制 api.py 到 shared 目录
@@ -107,12 +114,12 @@ class Main(star.Star):
shutil.copy(shared_api_file, self.shared_path)
self.workplace_path = os.path.join("data", "py_interpreter_workplace")
os.makedirs(self.workplace_path, exist_ok=True)
self.user_file_msg_buffer = defaultdict(list)
'''存放用户上传的文件和图片'''
"""存放用户上传的文件和图片"""
self.user_waiting = {}
'''正在等待用户的文件或图片'''
"""正在等待用户的文件或图片"""
# 加载配置
if not os.path.exists(PATH):
self.config = DEFAULT_CONFIG
@@ -120,33 +127,38 @@ class Main(star.Star):
else:
with open(PATH, "r") as f:
self.config = json.load(f)
async def initialize(self):
ok = await self.is_docker_available()
if not ok:
logger.info("Docker 不可用，代码解释器将无法使用，astrbot-python-interpreter 将自动禁用。")
await self.context._star_manager.turn_off_plugin("astrbot-python-interpreter")
logger.info(
"Docker 不可用，代码解释器将无法使用，astrbot-python-interpreter 将自动禁用。"
)
await self.context._star_manager.turn_off_plugin(
"astrbot-python-interpreter"
)
async def file_upload(self, file_path: str):
'''
"""
上传图像文件到 S3
'''
"""
ext = os.path.splitext(file_path)[1]
S3_URL = "https://s3.neko.soulter.top/astrbot-s3"
with open(file_path, "rb") as f:
file = f.read()
s3_file_url = f"{S3_URL}/{uuid.uuid4().hex}{ext}"
async with aiohttp.ClientSession(headers = {"Accept": "application/json"}, trust_env=True) as session:
async with aiohttp.ClientSession(
headers={"Accept": "application/json"}, trust_env=True
) as session:
async with session.put(s3_file_url, data=file) as resp:
if resp.status != 200:
raise Exception(f"Failed to upload image: {resp.status}")
return s3_file_url
async def is_docker_available(self) -> bool:
'''Check if docker is available'''
"""Check if docker is available"""
try:
docker = aiodocker.Docker()
await docker.version()
@@ -155,42 +167,44 @@ class Main(star.Star):
except BaseException as e:
logger.info(f"检查 Docker 可用性: {e}")
return False
async def get_image_name(self) -> str:
'''Get the image name'''
"""Get the image name"""
if self.config["sandbox"]["docker_mirror"]:
return f"{self.config['sandbox']['docker_mirror']}/{self.config['sandbox']['image']}"
return self.config["sandbox"]["image"]
def _save_config(self):
with open(PATH, "w") as f:
json.dump(self.config, f)
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 def download_image(
self, image_url: str, workplace_path: str, filename: str
) -> str:
"""Download image from url to workplace_path"""
async with aiohttp.ClientSession(trust_env=True) 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:
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'''
"""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)
@filter.event_message_type(filter.EventMessageType.ALL)
async def on_message(self, event: AstrMessageEvent):
'''处理消息'''
"""处理消息"""
uid = event.get_sender_id()
if uid not in self.user_waiting:
return
@@ -220,23 +234,24 @@ class Main(star.Star):
yield event.plain_result(f"代码执行器: 图片已经上传: {image_path}")
if uid in self.user_waiting:
del self.user_waiting[uid]
@filter.on_llm_request()
async def on_llm_req(self, event: AstrMessageEvent, request: ProviderRequest):
if event.get_session_id() in self.user_file_msg_buffer:
files = self.user_file_msg_buffer[event.get_session_id()]
request.prompt += f"\nUser provided files: {files}"
@filter.command_group("pi")
def pi(self):
pass
@pi.command("absdir")
async def pi_absdir(self, event: AstrMessageEvent, path: str = ""):
'''设置 Docker 宿主机绝对路径'''
"""设置 Docker 宿主机绝对路径"""
if not path:
yield event.plain_result(f"当前 Docker 宿主机绝对路径: {self.config.get('docker_host_astrbot_abs_path', '')}")
yield event.plain_result(
f"当前 Docker 宿主机绝对路径: {self.config.get('docker_host_astrbot_abs_path', '')}"
)
else:
self.config["docker_host_astrbot_abs_path"] = path
self._save_config()
@@ -244,9 +259,9 @@ class Main(star.Star):
@pi.command("mirror")
async def pi_mirror(self, event: AstrMessageEvent, url: str = ""):
'''Docker 镜像地址'''
"""Docker 镜像地址"""
if not url:
yield event.plain_result(f"""当前 Docker 镜像地址: {self.config['sandbox']['docker_mirror']}。
yield event.plain_result(f"""当前 Docker 镜像地址: {self.config["sandbox"]["docker_mirror"]}。
使用 `pi mirror <url>` 来设置 Docker 镜像地址。
您所设置的 Docker 镜像地址将会自动加在 Docker 镜像名前。如: `soulter/astrbot-code-interpreter-sandbox` -> `cjie.eu.org/soulter/astrbot-code-interpreter-sandbox`。
""")
@@ -257,7 +272,7 @@ class Main(star.Star):
@pi.command("repull")
async def pi_repull(self, event: AstrMessageEvent):
'''重新拉取沙箱镜像'''
"""重新拉取沙箱镜像"""
docker = aiodocker.Docker()
image_name = await self.get_image_name()
try:
@@ -267,28 +282,29 @@ class Main(star.Star):
pass
await docker.images.pull(image_name)
yield event.plain_result("重新拉取沙箱镜像成功。")
@pi.command("file")
async def pi_file(self, event: AstrMessageEvent):
'''在规定秒数(60s)内上传一个文件'''
"""在规定秒数(60s)内上传一个文件"""
uid = event.get_sender_id()
self.user_waiting[uid] = time.time()
tip = "文件"
yield event.plain_result(f"代码执行器: 请在 60s 内上传一个{tip}。")
await asyncio.sleep(60)
if uid in self.user_waiting:
yield event.plain_result(f"代码执行器: {event.get_sender_name()}/{event.get_sender_id()} 未在规定时间内上传{tip}。")
yield event.plain_result(
f"代码执行器: {event.get_sender_name()}/{event.get_sender_id()} 未在规定时间内上传{tip}。"
)
self.user_waiting.pop(uid)
@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.
"""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 math problems, edit image, docx, pptx, pdf, etc.
'''
"""
if not await self.is_docker_available():
yield event.plain_result("Docker 在当前机器不可用，无法沙箱化执行代码。")
plain_text = event.message_str
# 创建必要的工作目录和幻术码
@@ -297,7 +313,7 @@ class Main(star.Star):
output_path = os.path.join(workplace_path, "output")
os.makedirs(workplace_path, exist_ok=True)
os.makedirs(output_path, exist_ok=True)
files = []
# 文件
for file_path in self.user_file_msg_buffer[event.get_session_id()]:
@@ -310,9 +326,9 @@ class Main(star.Star):
file_name = os.path.basename(file_path)
shutil.copy(file_path, os.path.join(workplace_path, file_name))
files.append(file_name)
logger.debug(f"user query: {plain_text}, files: {files}")
# 整理额外输入
extra_inputs = ""
if files:
@@ -320,30 +336,33 @@ class Main(star.Star):
obs = ""
n = 5
for i in range(n):
if i > 0:
logger.info(f"Try {i+1}/{n}")
logger.info(f"Try {i + 1}/{n}")
PROMPT_ = PROMPT.format(
prompt=plain_text,
prompt=plain_text,
extra_input=extra_inputs,
extra_prompt=obs,
)
provider = self.context.get_using_provider()
llm_response = await provider.text_chat(prompt=PROMPT_, session_id=f"{event.session_id}_{magic_code}_{str(i)}")
logger.debug("code interpreter llm gened code:" + llm_response.completion_text)
llm_response = await provider.text_chat(
prompt=PROMPT_, session_id=f"{event.session_id}_{magic_code}_{str(i)}"
)
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
image_name = await self.get_image_name()
try:
@@ -352,47 +371,58 @@ class Main(star.Star):
# 拉取镜像
logger.info(f"未找到沙箱镜像，正在尝试拉取 {image_name}...")
await docker.images.pull(image_name)
yield event.plain_result(f"使用沙箱执行代码中，请稍等...(尝试次数: {i+1}/{n})")
self.docker_host_astrbot_abs_path = self.config.get("docker_host_astrbot_abs_path", "")
yield event.plain_result(
f"使用沙箱执行代码中，请稍等...(尝试次数: {i + 1}/{n})"
)
self.docker_host_astrbot_abs_path = self.config.get(
"docker_host_astrbot_abs_path", ""
)
if self.docker_host_astrbot_abs_path:
host_shared = os.path.join(self.docker_host_astrbot_abs_path, self.shared_path)
host_output = os.path.join(self.docker_host_astrbot_abs_path, output_path)
host_workplace = os.path.join(self.docker_host_astrbot_abs_path, workplace_path)
host_shared = os.path.join(
self.docker_host_astrbot_abs_path, self.shared_path
)
host_output = os.path.join(
self.docker_host_astrbot_abs_path, output_path
)
host_workplace = os.path.join(
self.docker_host_astrbot_abs_path, workplace_path
)
else:
host_shared = os.path.abspath(self.shared_path)
host_output = os.path.abspath(output_path)
host_workplace = os.path.abspath(workplace_path)
logger.debug(f"host_shared: {host_shared}, host_output: {host_output}, host_workplace: {host_workplace}")
container = await docker.containers.run({
"Image": image_name,
"Cmd": ["python", "exec.py"],
"Memory": 512 * 1024 * 1024,
"NanoCPUs": 1000000000,
"HostConfig": {
"Binds": [
f"{host_shared}:/astrbot_sandbox/shared:ro",
f"{host_output}:/astrbot_sandbox/output:rw",
f"{host_workplace}:/astrbot_sandbox:rw",
]
},
"Env": [
f"MAGIC_CODE={magic_code}"
],
"AutoRemove": True
})
logger.debug(
f"host_shared: {host_shared}, host_output: {host_output}, host_workplace: {host_workplace}"
)
container = await docker.containers.run(
{
"Image": image_name,
"Cmd": ["python", "exec.py"],
"Memory": 512 * 1024 * 1024,
"NanoCPUs": 1000000000,
"HostConfig": {
"Binds": [
f"{host_shared}:/astrbot_sandbox/shared:ro",
f"{host_output}:/astrbot_sandbox/output:rw",
f"{host_workplace}:/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|FILE)_OUTPUT#\w+\]: (.*)"
ok = False
@@ -415,39 +445,43 @@ class Main(star.Star):
file_name = os.path.basename(file_path)
chain = [File(name=file_name, file=file_s3_url)]
yield event.set_result(MessageEventResult(chain=chain))
elif "Traceback (most recent call last)" in log \
or "[Error]: " in log:
elif "Traceback (most recent call last)" in log or "[Error]: " in log:
traceback = "\n".join(logs[idx:])
if not ok:
if traceback:
obs = f"## Observation \n When execute the code: ```python\n{code_clean}\n```\n\n Error occurred:\n\n{traceback}\n Need to improve/fix the code."
else:
logger.warning(f"未从沙箱输出中捕获到合法的输出。沙箱输出日志: {logs}")
logger.warning(
f"未从沙箱输出中捕获到合法的输出。沙箱输出日志: {logs}"
)
break
else:
# 成功了
self.user_file_msg_buffer.pop(event.get_session_id())
return
yield event.plain_result("经过多次尝试后，未从沙箱输出中捕获到合法的输出，请更换问法或者查看日志。")
yield event.plain_result(
"经过多次尝试后，未从沙箱输出中捕获到合法的输出，请更换问法或者查看日志。"
)
@pi.command("cleanfile")
async def pi_cleanfile(self, event: AstrMessageEvent):
'''清理用户上传的文件'''
"""清理用户上传的文件"""
for file in self.user_file_msg_buffer[event.get_session_id()]:
try:
os.remove(file)
except BaseException as e:
logger.error(f"删除文件 {file} 失败: {e}")
self.user_file_msg_buffer.pop(event.get_session_id())
yield event.plain_result(f"用户 {event.get_session_id()} 上传的文件已清理。")
async def run_container(self, container: aiodocker.docker.DockerContainer, timeout: int = 20) -> list[str]:
'''Run the container and get the output'''
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)
packages/python_interpreter/shared/api.py
+8 -4
View File
@@ -1,18 +1,22 @@
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}")
def send_file(file_path: str):
if not os.path.exists(file_path):
raise Exception(f"File not found: {file_path}")
print(f"[ASTRBOT_FILE_OUTPUT#{_get_magic_code()}]: {file_path}")
print(f"[ASTRBOT_FILE_OUTPUT#{_get_magic_code()}]: {file_path}")