feat: 增强编码能力 - 新增代码编辑器、代码搜索、Git集成
- 新增 plugins/code_editor.py: 精准代码编辑（搜索替换、批量编辑、diff对比） - 新增 plugins/code_search.py: 代码搜索与理解（跨文件搜索、符号提取、依赖分析） - 新增 plugins/git_tools.py: Git 版本控制（status/diff/log/commit/branch） - 新增 core/ 模块: 配置管理、引擎管理、对话处理、显示输出、代码沙箱 - 修复 requirements.txt: 使用正确的 pip 包名（Pillow、opencv-python 等） - 更新 unified_tool_manager.py: 兼容 Plugin 和 Liugin 两种类名 - 更新系统提示词: 增强编程工作流指导 本次更新由 mimo 自动化执行，与开发者无关。
This commit is contained in:
Xiaoli Dev committed
commit
e22652fbbb
53 files changed
+20067
No files matched your search
@@ -0,0 +1,42 @@
|
||||
"""核心配置管理模块"""
|
||||
import os
|
||||
import json
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 常量
|
||||
CONFIG_FILE = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "config.json")
|
||||
MAX_FILE_SIZE = 52428800 # 50MB
|
||||
DEFAULT_MAX_HISTORY = 999999
|
||||
|
||||
def load_config():
|
||||
"""加载配置文件"""
|
||||
try:
|
||||
with open(CONFIG_FILE, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
except FileNotFoundError:
|
||||
return {"api": {"engines": {}}, "system": {}, "plugins": {}}
|
||||
except json.JSONDecodeError:
|
||||
return {"api": {"engines": {}}, "system": {}, "plugins": {}}
|
||||
|
||||
def save_config(config):
|
||||
"""保存配置文件"""
|
||||
try:
|
||||
with open(CONFIG_FILE, 'w', encoding='utf-8') as f:
|
||||
json.dump(config, f, ensure_ascii=False, indent=2)
|
||||
except Exception as e:
|
||||
logger.error(f"保存配置失败: {e}")
|
||||
|
||||
def get_system_config(key, default=None):
|
||||
"""获取系统配置"""
|
||||
config = load_config()
|
||||
return config.get("system", {}).get(key, default)
|
||||
|
||||
def set_system_config(key, value):
|
||||
"""设置系统配置"""
|
||||
config = load_config()
|
||||
if "system" not in config:
|
||||
config["system"] = {}
|
||||
config["system"][key] = value
|
||||
save_config(config)
|
||||
@@ -0,0 +1,315 @@
|
||||
"""
|
||||
对话处理模块 - 核心对话循环和工具调用
|
||||
"""
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
import threading
|
||||
import logging
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from colorama import Fore, Style
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ConversationManager:
|
||||
"""对话管理器 - 处理用户输入、AI 响应、工具调用循环"""
|
||||
|
||||
def __init__(self, cli):
|
||||
self.cli = cli
|
||||
self.history = []
|
||||
self.max_history = 999999
|
||||
|
||||
def process(self, user_input: str):
|
||||
"""处理完整的对话循环"""
|
||||
self.history.append({"role": "user", "content": user_input})
|
||||
|
||||
tool_prompts = self._get_tool_prompts()
|
||||
loop_count = 0
|
||||
current_input = user_input
|
||||
max_loops = 20
|
||||
|
||||
while loop_count < max_loops:
|
||||
response = self._generate_with_animation(current_input, tool_prompts)
|
||||
processed = self._process_thinking(response)
|
||||
processed = self._process_code_blocks(processed)
|
||||
|
||||
try:
|
||||
text_content, json_data = self._parse_response(processed)
|
||||
|
||||
if json_data:
|
||||
if text_content:
|
||||
self._add_and_display(text_content)
|
||||
|
||||
if isinstance(json_data, dict):
|
||||
result = self._handle_dict(json_data, text_content, processed)
|
||||
if result == "break":
|
||||
break
|
||||
elif result == "continue":
|
||||
loop_count += 1
|
||||
continue
|
||||
elif isinstance(json_data, list):
|
||||
result = self._handle_list(json_data, text_content, processed)
|
||||
if result == "break":
|
||||
break
|
||||
elif result == "continue":
|
||||
loop_count += 1
|
||||
continue
|
||||
else:
|
||||
display = text_content if text_content else processed
|
||||
self._add_and_display(display)
|
||||
break
|
||||
|
||||
except json.JSONDecodeError:
|
||||
text_content, json_data = self._parse_response(processed)
|
||||
if json_data:
|
||||
# Re-process with parsed data
|
||||
pass
|
||||
else:
|
||||
self._add_and_display(processed)
|
||||
break
|
||||
|
||||
if loop_count >= max_loops:
|
||||
self.cli.display.output(f"AI 已达到最大循环次数 ({max_loops})，已停止")
|
||||
|
||||
def _handle_dict(self, data, text_content, processed):
|
||||
"""处理 dict 类型的 JSON 数据"""
|
||||
if data.get('action') == 'use_tool':
|
||||
self.cli._execute_single_tool(data)
|
||||
return "break"
|
||||
|
||||
if data.get('action') == 'continue':
|
||||
content = data.get('content', '')
|
||||
if content:
|
||||
self._add_and_display(content)
|
||||
return "continue"
|
||||
|
||||
if data.get('continue') or data.get('need_continue') or data.get('think_more'):
|
||||
msg = data.get('message', '')
|
||||
if msg:
|
||||
self.history.append({"role": "assistant", "content": msg})
|
||||
self.cli.display.response(msg, is_continue=True)
|
||||
return "continue"
|
||||
|
||||
msg = data.get('message', '')
|
||||
if msg:
|
||||
self._add_and_display(msg)
|
||||
return "break"
|
||||
|
||||
if text_content:
|
||||
self._add_and_display(text_content)
|
||||
else:
|
||||
self._add_and_display(processed)
|
||||
return "break"
|
||||
|
||||
def _handle_list(self, data_list, text_content, processed):
|
||||
"""处理 list 类型的 JSON 数据"""
|
||||
tool_calls = [item for item in data_list
|
||||
if isinstance(item, dict) and item.get('action') == 'use_tool']
|
||||
|
||||
if tool_calls:
|
||||
self.cli._execute_concurrent_tools(tool_calls)
|
||||
return "break"
|
||||
|
||||
continue_items = [item for item in data_list
|
||||
if isinstance(item, dict) and
|
||||
(item.get('continue') or item.get('need_continue') or item.get('think_more'))]
|
||||
|
||||
if continue_items:
|
||||
msg = continue_items[0].get('message', '')
|
||||
if msg:
|
||||
self.history.append({"role": "assistant", "content": msg})
|
||||
self.cli.display.response(msg, is_continue=True)
|
||||
return "continue"
|
||||
|
||||
if text_content:
|
||||
self._add_and_display(text_content)
|
||||
else:
|
||||
self._add_and_display(processed)
|
||||
return "break"
|
||||
|
||||
def _add_and_display(self, content):
|
||||
"""添加到历史并显示"""
|
||||
self.history.append({"role": "assistant", "content": content})
|
||||
self.cli.display.response(content)
|
||||
|
||||
def _generate_with_animation(self, user_input, tool_prompts=None):
|
||||
"""生成 AI 响应（带动画）"""
|
||||
system_prompt = self._build_system_prompt(tool_prompts)
|
||||
|
||||
if not self.cli.engine_manager.current_engine:
|
||||
return "错误: 当前没有可用的AI引擎"
|
||||
|
||||
# TUI 模式：直接生成
|
||||
if self.cli.display.tui_callback:
|
||||
response = self.cli.engine_manager.generate_response(user_input, system_prompt)
|
||||
if response:
|
||||
self.history.append({"role": "assistant", "content": response})
|
||||
return response
|
||||
|
||||
# CLI 模式：带动画
|
||||
stop_event = threading.Event()
|
||||
esc_event = threading.Event()
|
||||
|
||||
# 加载等待句子
|
||||
sentences = self._load_love_sentences()
|
||||
|
||||
anim_thread = threading.Thread(
|
||||
target=self.cli.display.thinking_animation,
|
||||
args=(sentences, stop_event, esc_event), daemon=True)
|
||||
anim_thread.start()
|
||||
|
||||
try:
|
||||
response = self.cli.engine_manager.generate_response(user_input, system_prompt)
|
||||
if response:
|
||||
self.cli.display.typeprint(response, Fore.CYAN)
|
||||
self.history.append({"role": "assistant", "content": response})
|
||||
return response
|
||||
finally:
|
||||
stop_event.set()
|
||||
anim_thread.join()
|
||||
|
||||
def _build_system_prompt(self, tool_prompts=None):
|
||||
"""构建系统提示词"""
|
||||
base = ("你叫小狸，是一个强大的智能编程助手。你具备精准的代码编辑能力、"
|
||||
"完整的 Git 工作流支持、以及跨文件代码搜索和理解能力。"
|
||||
"你可以帮助用户编写、修改、调试代码，管理版本控制，理解复杂代码库。")
|
||||
|
||||
tool_section = ""
|
||||
if tool_prompts:
|
||||
tool_section = f"\n\n{tool_prompts}"
|
||||
|
||||
call_format = """
|
||||
【工具调用格式】
|
||||
当需要执行操作时，在回复中包含 JSON 指令：
|
||||
|
||||
1. 纯 JSON: {"action": "use_tool", "tool": "工具名", "args": "参数"}
|
||||
2. 文本+JSON: 好的，我来修改这个文件。{"action": "use_tool", "tool": "code_editor", "args": "edit file.py old_code <<<>>> new_code"}
|
||||
|
||||
【核心工具】
|
||||
- code_editor: 精准代码编辑（edit/search/diff/insert/create/write）
|
||||
- code_search: 代码搜索与理解（find/regex/symbols/imports/structure）
|
||||
- git_tools: Git 操作（status/diff/log/add/commit/branch）
|
||||
- cmd_executor: 执行 shell 命令
|
||||
- file_manager: 文件管理
|
||||
|
||||
【编码工作流】
|
||||
1. 先用 code_search 了解代码结构
|
||||
2. 用 code_editor.read_range 查看相关代码
|
||||
3. 用 code_editor.edit 精准修改
|
||||
4. 用 git_tools 提交变更
|
||||
|
||||
【注意事项】
|
||||
- code_editor.edit 的格式: old_text <<<>>> new_text（用 <<<>>> 分隔）
|
||||
- args 必须是字符串
|
||||
- 可以一次返回多个 JSON 指令顺序执行
|
||||
- 使用 continue 补充回答: {"action": "continue", "content": "补充内容"}"""
|
||||
|
||||
return base + tool_section + call_format
|
||||
|
||||
def _process_thinking(self, response):
|
||||
"""处理深度思考标记"""
|
||||
# <think/> 格式
|
||||
if '<think/>' in response:
|
||||
parts = response.split('<think/>', 1)
|
||||
if len(parts) >= 2:
|
||||
thinking = parts[0].strip()
|
||||
reply = parts[1].strip()
|
||||
if thinking:
|
||||
print(f"{Fore.LIGHTBLACK_EX}[思考: {thinking[:200]}...]{Style.RESET_ALL}" if len(thinking) > 200 else f"{Fore.LIGHTBLACK_EX}[思考: {thinking}]{Style.RESET_ALL}")
|
||||
return reply if reply else ""
|
||||
|
||||
# <thinking>...</thinking> 格式
|
||||
engine = self.cli.engine_manager.current_engine
|
||||
if engine:
|
||||
start = getattr(engine, 'thinking_start_marker', None)
|
||||
end = getattr(engine, 'thinking_end_marker', None)
|
||||
if start and end and start in response and end in response:
|
||||
parts = response.split(start, 1)
|
||||
if len(parts) > 1 and end in parts[1]:
|
||||
thinking, reply = parts[1].split(end, 1)
|
||||
if thinking.strip():
|
||||
print(f"{Fore.LIGHTBLACK_EX}[思考: {thinking.strip()[:200]}]{Style.RESET_ALL}")
|
||||
return reply.strip()
|
||||
|
||||
return response
|
||||
|
||||
def _process_code_blocks(self, text):
|
||||
"""处理代码块标记"""
|
||||
pattern = r'(""".*?""")'
|
||||
def replace(m):
|
||||
inner = m.group(1)[3:-3]
|
||||
return f'{Fore.RED}{inner}{Style.RESET_ALL}'
|
||||
return re.sub(pattern, replace, text, flags=re.DOTALL)
|
||||
|
||||
def _parse_response(self, response):
|
||||
"""解析混合响应（文本+JSON）"""
|
||||
response = response.strip()
|
||||
|
||||
# 直接 JSON
|
||||
try:
|
||||
parsed = json.loads(response)
|
||||
if isinstance(parsed, (dict, list)):
|
||||
return "", parsed
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# 代码块中的 JSON
|
||||
pattern = r"```(?:json)?\s*({.*?})\s*```"
|
||||
matches = re.findall(pattern, response, re.DOTALL)
|
||||
if matches:
|
||||
try:
|
||||
data = json.loads(matches[-1])
|
||||
if isinstance(data, (dict, list)):
|
||||
text = re.sub(pattern, "", response, flags=re.DOTALL).strip()
|
||||
return text, data
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# 内联 JSON
|
||||
markers = ['{"action": "use_tool"', '{"action": "continue"',
|
||||
'{"action":"use_tool"', '{"action":"continue"']
|
||||
for marker in markers:
|
||||
idx = response.find(marker)
|
||||
if idx != -1:
|
||||
brace_count = 0
|
||||
for i in range(idx, len(response)):
|
||||
if response[i] == '{':
|
||||
brace_count += 1
|
||||
elif response[i] == '}':
|
||||
brace_count -= 1
|
||||
if brace_count == 0:
|
||||
try:
|
||||
data = json.loads(response[idx:i+1])
|
||||
if isinstance(data, dict) and data.get('action') in ['use_tool', 'continue']:
|
||||
return response[:idx].strip(), data
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
break
|
||||
|
||||
return response, None
|
||||
|
||||
def _get_tool_prompts(self):
|
||||
"""获取工具提示词"""
|
||||
if not self.cli.tool_manager or not self.cli.tool_manager.tools:
|
||||
return "当前没有可用的工具"
|
||||
|
||||
lines = ["可用工具列表：\n"]
|
||||
for tool in self.cli.tool_manager.tools:
|
||||
name = tool.get('name', '')
|
||||
desc = tool.get('description', '')
|
||||
lines.append(f"- {name}: {desc}")
|
||||
|
||||
return '\n'.join(lines)
|
||||
|
||||
def _load_love_sentences(self):
|
||||
"""加载等待动画句子"""
|
||||
import os
|
||||
path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "love.txt")
|
||||
if os.path.exists(path):
|
||||
try:
|
||||
with open(path, 'r', encoding='utf-8') as f:
|
||||
return [l.strip() for l in f.readlines() if l.strip()]
|
||||
except:
|
||||
pass
|
||||
return ["AI正在思考中...", "请稍等片刻...", "正在处理您的请求..."]
|
||||
+131
@@ -0,0 +1,131 @@
|
||||
"""
|
||||
显示与输出模块 - 统一 CLI/TUI 输出
|
||||
"""
|
||||
import sys
|
||||
import os
|
||||
import time
|
||||
import random
|
||||
import shutil
|
||||
import threading
|
||||
from colorama import Fore, Style
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Display:
|
||||
"""统一显示管理器"""
|
||||
|
||||
def __init__(self, user_id="", tui_callback=None):
|
||||
self.user_id = user_id
|
||||
self.tui_callback = tui_callback
|
||||
|
||||
def output(self, message):
|
||||
"""统一输出方法"""
|
||||
if self.tui_callback:
|
||||
self.tui_callback(message)
|
||||
else:
|
||||
print(message)
|
||||
|
||||
def response(self, content, is_continue=False):
|
||||
"""显示 AI 响应"""
|
||||
user_id_display = f"[用户ID: {self.user_id}]"
|
||||
lines = content.split('\n')
|
||||
|
||||
for i, line in enumerate(lines):
|
||||
if i == 0:
|
||||
prefix = f"{Fore.YELLOW}✦{Style.RESET_ALL} " if is_continue else "✦ "
|
||||
line = prefix + line
|
||||
if i == len(lines) - 1:
|
||||
print(f"{line} {user_id_display}")
|
||||
else:
|
||||
print(line)
|
||||
print()
|
||||
|
||||
def tool_result(self, tool_name, tool_args, full_result, is_error=False):
|
||||
"""显示工具执行结果"""
|
||||
if full_result is None:
|
||||
full_result = "无结果"
|
||||
|
||||
self.output("")
|
||||
user_id_display = f"[用户ID: {self.user_id}]"
|
||||
|
||||
if is_error:
|
||||
self.output(f"{Fore.RED} X 工具: {tool_name}{Style.RESET_ALL}")
|
||||
self.output(f"{Fore.RED} 参数: {tool_args} {user_id_display}{Style.RESET_ALL}")
|
||||
else:
|
||||
self.output(f"{Fore.GREEN} OK 工具: {tool_name}{Style.RESET_ALL}")
|
||||
self.output(f"{Fore.GREEN} 参数: {tool_args} {user_id_display}{Style.RESET_ALL}")
|
||||
|
||||
display_result = self._limit_lines(full_result, max_lines=2)
|
||||
color = Fore.RED if is_error else Fore.GREEN
|
||||
for line in display_result.split('\n'):
|
||||
self.output(f"{color} {line}{Style.RESET_ALL}")
|
||||
|
||||
def error(self, message):
|
||||
"""显示错误"""
|
||||
self.output(f"{Fore.RED}{message}{Style.RESET_ALL}")
|
||||
|
||||
def success(self, message):
|
||||
"""显示成功"""
|
||||
self.output(f"{Fore.GREEN}{message}{Style.RESET_ALL}")
|
||||
|
||||
def info(self, message):
|
||||
"""显示信息"""
|
||||
self.output(f"{Fore.CYAN}{message}{Style.RESET_ALL}")
|
||||
|
||||
def warning(self, message):
|
||||
"""显示警告"""
|
||||
self.output(f"{Fore.YELLOW}{message}{Style.RESET_ALL}")
|
||||
|
||||
def typeprint(self, text, color=None, delay=0.01):
|
||||
"""打字机效果输出"""
|
||||
for char in text:
|
||||
if color:
|
||||
print(f"{color}{char}{Style.RESET_ALL}", end="", flush=True)
|
||||
else:
|
||||
print(char, end="", flush=True)
|
||||
time.sleep(delay)
|
||||
print()
|
||||
|
||||
def thinking_animation(self, sentences, stop_event, esc_event):
|
||||
"""等待动画"""
|
||||
if not sentences:
|
||||
sentences = ["AI正在思考中...", "请稍等片刻...", "正在处理您的请求..."]
|
||||
|
||||
last_change = time.time()
|
||||
current_sentence = random.choice(sentences)
|
||||
chars = "|/-\\"
|
||||
idx = 0
|
||||
|
||||
while not stop_event.is_set():
|
||||
now = time.time()
|
||||
if now - last_change >= 5:
|
||||
current_sentence = random.choice(sentences)
|
||||
last_change = now
|
||||
|
||||
char = chars[idx % len(chars)]
|
||||
display = current_sentence[:47] + "..." if len(current_sentence) > 50 else current_sentence
|
||||
print(f"\r{Fore.MAGENTA}{display} {char}{Style.RESET_ALL}", end="", flush=True)
|
||||
idx += 1
|
||||
time.sleep(0.1)
|
||||
|
||||
print("\r" + " " * 60 + "\r", end="", flush=True)
|
||||
|
||||
def _limit_lines(self, text, max_lines=2):
|
||||
"""限制输出行数"""
|
||||
lines = text.split('\n')
|
||||
if len(lines) <= max_lines:
|
||||
return text
|
||||
return '\n'.join(lines[:max_lines]) + f"\n... (已截断，共{len(lines)}行)"
|
||||
|
||||
def show_banner(self, engine_name, engines, user_id):
|
||||
"""显示启动横幅"""
|
||||
self.success(f"小狸 Pro-CLI 已启动!")
|
||||
self.success(f"用户ID: {user_id}")
|
||||
print(f"{Fore.GREEN}输入 '/help' 查看帮助信息{Style.RESET_ALL}")
|
||||
print(f"{Fore.GREEN}输入 '/quit' 退出程序{Style.RESET_ALL}")
|
||||
print(f"{Fore.GREEN}当前引擎: {engine_name}{Style.RESET_ALL}")
|
||||
if len(engines) > 1:
|
||||
print(f"{Fore.GREEN}已加载引擎: {', '.join(engines)}{Style.RESET_ALL}")
|
||||
print(f"{Fore.GREEN}{'-' * 50}{Style.RESET_ALL}")
|
||||
+115
@@ -0,0 +1,115 @@
|
||||
"""
|
||||
AI 引擎管理模块 - 动态加载和切换 AI 引擎
|
||||
"""
|
||||
import os
|
||||
import importlib.util
|
||||
import logging
|
||||
from colorama import Fore, Style
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class EngineManager:
|
||||
"""AI 引擎管理器"""
|
||||
|
||||
def __init__(self):
|
||||
self.engines = {}
|
||||
self.current_engine = None
|
||||
|
||||
def load_engines(self, engines_dir: str) -> dict:
|
||||
"""动态加载所有 AI 引擎"""
|
||||
print(f"{Fore.GREEN}正在加载AI引擎插件...{Style.RESET_ALL}")
|
||||
|
||||
if not os.path.exists(engines_dir):
|
||||
print(f"{Fore.YELLOW}引擎目录不存在: {engines_dir}{Style.RESET_ALL}")
|
||||
return self.engines
|
||||
|
||||
for filename in os.listdir(engines_dir):
|
||||
if not filename.endswith('.py') or filename == '__init__.py':
|
||||
continue
|
||||
|
||||
engine_name = filename[:-3]
|
||||
engine_path = os.path.join(engines_dir, filename)
|
||||
|
||||
try:
|
||||
spec = importlib.util.spec_from_file_location(engine_name, engine_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
class_name = self._find_engine_class(module, engine_name)
|
||||
if class_name and hasattr(module, class_name):
|
||||
engine_instance = getattr(module, class_name)()
|
||||
engine_key = getattr(engine_instance, 'name', engine_name)
|
||||
|
||||
if engine_key not in self.engines:
|
||||
self.engines[engine_key] = engine_instance
|
||||
print(f"{Fore.GREEN}已加载AI引擎: {engine_key} ({class_name}){Style.RESET_ALL}")
|
||||
else:
|
||||
print(f"{Fore.YELLOW}AI引擎 {engine_key} 已存在，跳过{Style.RESET_ALL}")
|
||||
else:
|
||||
print(f"{Fore.YELLOW}引擎插件 {filename} 中未找到合适的引擎类{Style.RESET_ALL}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"加载引擎失败 {filename}: {e}")
|
||||
print(f"{Fore.RED}加载引擎失败 {filename}: {e}{Style.RESET_ALL}")
|
||||
|
||||
print(f"{Fore.GREEN}总共加载了 {len(self.engines)} 个AI引擎{Style.RESET_ALL}")
|
||||
return self.engines
|
||||
|
||||
def _find_engine_class(self, module, engine_name):
|
||||
"""查找引擎类"""
|
||||
possible_names = [
|
||||
engine_name.capitalize() + "AI",
|
||||
engine_name.replace('_', '').replace('-', '').capitalize() + "AI",
|
||||
''.join(w.capitalize() for w in engine_name.replace('-', ' ').replace('_', ' ').split()) + "AI",
|
||||
''.join(w.capitalize() for w in engine_name.replace('_engine', '').replace('-', ' ').replace('_', ' ').split()) + "AI",
|
||||
]
|
||||
|
||||
for name in possible_names:
|
||||
if hasattr(module, name):
|
||||
return name
|
||||
|
||||
# 通用查找
|
||||
for attr_name in dir(module):
|
||||
attr = getattr(module, attr_name)
|
||||
if (isinstance(attr, type) and
|
||||
attr_name.endswith('AI') and
|
||||
hasattr(attr, 'generate_response')):
|
||||
return attr_name
|
||||
return None
|
||||
|
||||
def set_default(self, name: str) -> bool:
|
||||
"""设置默认引擎"""
|
||||
if name in self.engines:
|
||||
self.current_engine = self.engines[name]
|
||||
return True
|
||||
return False
|
||||
|
||||
def switch(self, name: str) -> bool:
|
||||
"""切换引擎"""
|
||||
if name not in self.engines:
|
||||
return False
|
||||
|
||||
old_engine = self.current_engine
|
||||
self.current_engine = self.engines[name]
|
||||
|
||||
if hasattr(self.current_engine, 'cli'):
|
||||
self.current_engine.cli = getattr(old_engine, 'cli', None)
|
||||
|
||||
return True
|
||||
|
||||
def get_current_name(self) -> str:
|
||||
"""获取当前引擎名称"""
|
||||
if self.current_engine:
|
||||
return getattr(self.current_engine, 'name', '未设置')
|
||||
return '未设置'
|
||||
|
||||
def list_engines(self) -> list:
|
||||
"""列出所有引擎"""
|
||||
return list(self.engines.keys())
|
||||
|
||||
def generate_response(self, user_input: str, system_prompt: str = None) -> str:
|
||||
"""使用当前引擎生成响应"""
|
||||
if not self.current_engine:
|
||||
return "错误: 没有可用的AI引擎"
|
||||
return self.current_engine.generate_response(user_input, system_prompt=system_prompt)
|
||||
+210
@@ -0,0 +1,210 @@
|
||||
"""
|
||||
安全代码沙箱 - 在隔离环境中执行用户/插件代码
|
||||
"""
|
||||
import io
|
||||
import ast
|
||||
import time
|
||||
import re
|
||||
import threading
|
||||
import traceback
|
||||
from contextlib import redirect_stdout, redirect_stderr
|
||||
from datetime import datetime
|
||||
from colorama import Fore, Style
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CodeSandbox:
|
||||
"""安全代码执行沙箱"""
|
||||
|
||||
# 危险模块黑名单
|
||||
DANGEROUS_MODULES = {
|
||||
'os', 'sys', 'subprocess', 'shutil', 'socket', 'pickle',
|
||||
'marshal', 'ctypes', 'multiprocessing', 'threading',
|
||||
'signal', 'resource', 'posix', 'nt', 'builtins',
|
||||
'importlib', 'code', 'codeop', 'commands', 'popen2',
|
||||
}
|
||||
|
||||
# 危险函数/属性黑名单
|
||||
DANGEROUS_FUNCTIONS = {
|
||||
'exec', 'eval', 'compile', 'execfile', 'open',
|
||||
'input', 'raw_input', 'globals', 'locals', 'vars',
|
||||
'dir', 'getattr', 'setattr', 'delattr', 'hasattr',
|
||||
'__import__', '__builtins__', '__class__', '__bases__',
|
||||
'__subclasses__', '__mro__', '__dict__', '__getattribute__',
|
||||
}
|
||||
|
||||
def __init__(self, timeout=30, whitelist=None):
|
||||
self.enabled = True
|
||||
self.timeout = timeout
|
||||
self.whitelist = whitelist or [
|
||||
'json', 'math', 'datetime', 'random', 'collections',
|
||||
'itertools', 'functools', 'operator', 'typing', 're'
|
||||
]
|
||||
self.history = []
|
||||
self.max_history = 100
|
||||
|
||||
def is_safe(self, code: str) -> bool:
|
||||
"""检查代码是否安全（AST 分析 + 字符串快速检查）"""
|
||||
# 快速字符串检查
|
||||
dangerous_patterns = [
|
||||
r'__', r'\beval\b', r'\bexec\b', r'\bcompile\b',
|
||||
r'\bglobals\b', r'\blocals\b', r'\bgetattr\b',
|
||||
r'\bsetattr\b', r'\b__import__\b',
|
||||
]
|
||||
for pattern in dangerous_patterns:
|
||||
if re.search(pattern, code, re.IGNORECASE):
|
||||
return False
|
||||
|
||||
# AST 分析
|
||||
try:
|
||||
tree = ast.parse(code)
|
||||
except SyntaxError:
|
||||
return False
|
||||
|
||||
class SafetyVisitor(ast.NodeVisitor):
|
||||
def __init__(self):
|
||||
self.is_safe = True
|
||||
|
||||
def _check_name(self, name):
|
||||
if name in CodeSandbox.DANGEROUS_FUNCTIONS:
|
||||
self.is_safe = False
|
||||
|
||||
def visit_Import(self, node):
|
||||
for alias in node.names:
|
||||
if alias.name.split('.')[0] in CodeSandbox.DANGEROUS_MODULES:
|
||||
self.is_safe = False
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_ImportFrom(self, node):
|
||||
if node.module:
|
||||
if node.module.split('.')[0] in CodeSandbox.DANGEROUS_MODULES:
|
||||
self.is_safe = False
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_Call(self, node):
|
||||
if isinstance(node.func, ast.Name):
|
||||
self._check_name(node.func.id)
|
||||
elif isinstance(node.func, ast.Attribute):
|
||||
self._check_name(node.func.attr)
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_Attribute(self, node):
|
||||
self._check_name(node.attr)
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_Name(self, node):
|
||||
self._check_name(node.id)
|
||||
self.generic_visit(node)
|
||||
|
||||
visitor = SafetyVisitor()
|
||||
visitor.visit(tree)
|
||||
return visitor.is_safe
|
||||
|
||||
def execute(self, code: str, variables: dict = None) -> dict:
|
||||
"""
|
||||
在安全沙箱中执行代码
|
||||
|
||||
Returns:
|
||||
{
|
||||
'success': bool,
|
||||
'result': str,
|
||||
'stdout': str,
|
||||
'stderr': str,
|
||||
'execution_time': float,
|
||||
}
|
||||
"""
|
||||
if not self.enabled:
|
||||
return {'success': False, 'result': '代码执行功能已禁用',
|
||||
'stdout': '', 'stderr': '', 'execution_time': 0}
|
||||
|
||||
if not self.is_safe(code):
|
||||
return {'success': False, 'result': '代码包含潜在危险操作，已阻止执行',
|
||||
'stdout': '', 'stderr': '', 'execution_time': 0}
|
||||
|
||||
# 准备变量
|
||||
exec_vars = {}
|
||||
if variables:
|
||||
exec_vars.update(variables)
|
||||
|
||||
# 添加安全的内置函数
|
||||
safe_builtins = {
|
||||
'print': print, 'len': len, 'str': str, 'int': int,
|
||||
'float': float, 'list': list, 'dict': dict, 'tuple': tuple,
|
||||
'set': set, 'range': range, 'enumerate': enumerate,
|
||||
'zip': zip, 'sorted': sorted, 'sum': sum, 'max': max,
|
||||
'min': min, 'abs': abs, 'round': round, 'bool': bool,
|
||||
'type': type, 'isinstance': isinstance,
|
||||
}
|
||||
exec_vars.update(safe_builtins)
|
||||
|
||||
# 添加白名单模块
|
||||
for module_name in self.whitelist:
|
||||
try:
|
||||
exec_vars[module_name] = __import__(module_name)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# 捕获输出
|
||||
stdout_capture = io.StringIO()
|
||||
stderr_capture = io.StringIO()
|
||||
|
||||
start_time = time.time()
|
||||
result = {'success': True, 'error': None, 'traceback': None}
|
||||
|
||||
def _exec():
|
||||
try:
|
||||
exec(code, exec_vars)
|
||||
except Exception as e:
|
||||
result['success'] = False
|
||||
result['error'] = str(e)
|
||||
result['traceback'] = traceback.format_exc()
|
||||
|
||||
exec_thread = threading.Thread(target=_exec, daemon=True)
|
||||
exec_thread.start()
|
||||
exec_thread.join(timeout=self.timeout)
|
||||
|
||||
exec_time = time.time() - start_time
|
||||
|
||||
if exec_thread.is_alive():
|
||||
return {
|
||||
'success': False,
|
||||
'result': f'代码执行超时（超过 {self.timeout} 秒）',
|
||||
'stdout': stdout_capture.getvalue(),
|
||||
'stderr': stderr_capture.getvalue(),
|
||||
'execution_time': self.timeout,
|
||||
}
|
||||
|
||||
if result['success']:
|
||||
exec_result = exec_vars.get('result', None)
|
||||
result_str = str(exec_result) if exec_result is not None else '代码执行成功（无返回值）'
|
||||
return {
|
||||
'success': True,
|
||||
'result': result_str,
|
||||
'stdout': stdout_capture.getvalue(),
|
||||
'stderr': stderr_capture.getvalue(),
|
||||
'execution_time': exec_time,
|
||||
}
|
||||
else:
|
||||
return {
|
||||
'success': False,
|
||||
'result': f"代码执行错误: {result['error']}",
|
||||
'stdout': stdout_capture.getvalue(),
|
||||
'stderr': stderr_capture.getvalue() + '\n' + (result['traceback'] or ''),
|
||||
'execution_time': exec_time,
|
||||
}
|
||||
|
||||
def add_to_history(self, record: dict):
|
||||
"""添加执行记录"""
|
||||
self.history.append(record)
|
||||
if len(self.history) > self.max_history:
|
||||
self.history.pop(0)
|
||||
|
||||
def get_history(self, limit=10):
|
||||
"""获取执行历史"""
|
||||
return self.history[-limit:]
|
||||
|
||||
def clear_history(self):
|
||||
"""清空执行历史"""
|
||||
self.history = []
|
||||
Reference in new issue
Block a user