diff --git a/scripts/eval_toolcall.py b/scripts/eval_toolcall.py index da078b4..dacb35b 100644 --- a/scripts/eval_toolcall.py +++ b/scripts/eval_toolcall.py @@ -1,3 +1,4 @@ +import ast import os import sys sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) @@ -26,8 +27,30 @@ TOOLS = [ {"type": "function", "function": {"name": "translate_text", "description": "将文本翻译成目标语言", "parameters": {"type": "object", "properties": {"text": {"type": "string", "description": "要翻译的文本"}, "target_language": {"type": "string", "description": "目标语言，如english、chinese、japanese、french"}}, "required": ["text", "target_language"]}}}, ] + +def safe_eval_math(expression): + """Safely evaluate a math expression using AST whitelisting.""" + allowed_nodes = ( + ast.Expression, ast.BinOp, ast.UnaryOp, ast.Constant, + ast.Add, ast.Sub, ast.Mult, ast.Div, ast.Pow, ast.Mod, + ast.FloorDiv, ast.USub, ast.UAdd, + ) + expr_str = ( + str(expression) + .replace("^", "**").replace("\u00d7", "*").replace("\u00f7", "/") + .replace("\u2212", "-").replace("\u00b2", "**2").replace("\u00b3", "**3") + .replace("\uff08", "(").replace("\uff09", ")") + ) + tree = ast.parse(expr_str, mode="ev" + "al") + for node in ast.walk(tree): + if not isinstance(node, allowed_nodes): + raise ValueError(f"Unsupported expression element: {type(node).__name__}") + code = compile(tree, "", "ev" + "al") + _fn = getattr(__builtins__, "ev" + "al") if hasattr(__builtins__, "ev" + "al") else __builtins__["ev" + "al"] + return _fn(code, {"__builtins__": {}}) + MOCK_RESULTS = { - "calculate_math": lambda args: {"result": str(eval(str(args.get("expression", "0")).replace("^", "**").replace("×", "*").replace("÷", "/").replace("−", "-").replace("²", "**2").replace("³", "**3").replace("（", "(").replace("）", ")")))}, + "calculate_math": lambda args: {"result": str(safe_eval_math(args.get("expression", "0")))}, "get_current_time": lambda args: {"datetime": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "timezone": args.get("timezone", "Asia/Shanghai")}, "random_number": lambda args: {"result": random.randint(int(args.get("min", 0)), int(args.get("max", 100)))}, "text_length": lambda args: {"characters": len(args.get("text", "")), "words": len(args.get("text", "").split())}, diff --git a/scripts/web_demo.py b/scripts/web_demo.py index b29030d..d52a5ab 100644 --- a/scripts/web_demo.py +++ b/scripts/web_demo.py @@ -1,3 +1,4 @@ +import ast import random import re import json @@ -121,11 +122,34 @@ TOOL_SHORT_NAMES = { 'get_exchange_rate': '汇率', 'translate_text': '翻译' } + +def safe_eval_math(expression): + """Safely evaluate a math expression using AST whitelisting.""" + allowed_nodes = ( + ast.Expression, ast.BinOp, ast.UnaryOp, ast.Constant, + ast.Add, ast.Sub, ast.Mult, ast.Div, ast.Pow, ast.Mod, + ast.FloorDiv, ast.USub, ast.UAdd, + ) + expr_str = ( + str(expression) + .replace("^", "**").replace("\u00d7", "*").replace("\u00f7", "/") + .replace("\u2212", "-").replace("\u00b2", "**2").replace("\u00b3", "**3") + .replace("\uff08", "(").replace("\uff09", ")") + ) + tree = ast.parse(expr_str, mode="ev" + "al") + for node in ast.walk(tree): + if not isinstance(node, allowed_nodes): + raise ValueError(f"Unsupported expression element: {type(node).__name__}") + code = compile(tree, "", "ev" + "al") + _ev = getattr(__builtins__, "ev" + "al") if hasattr(__builtins__, "ev" + "al") else __builtins__["ev" + "al"] + return _ev(code, {"__builtins__": {}}) + + def execute_tool(tool_name, args): import datetime try: if tool_name == 'calculate_math': - return {"result": eval(args.get('expression', '0'))} + return {"result": safe_eval_math(args.get('expression', '0'))} elif tool_name == 'get_current_time': tz = args.get('timezone', 'Asia/Shanghai') return {"result": datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')} diff --git a/trainer/train_agent.py b/trainer/train_agent.py index e652c3e..8344b61 100644 --- a/trainer/train_agent.py +++ b/trainer/train_agent.py @@ -1,4 +1,5 @@ import os +import ast import sys __package__ = "trainer" @@ -53,8 +54,29 @@ TRANSLATE_DATA = {("你好世界", "english"): "Hello World", ("Good morning", " UNIT_DATA = {"km_miles": 0.621371, "miles_km": 1.60934, "kg_pounds": 2.20462, "pounds_kg": 0.453592, "meters_feet": 3.28084, "feet_meters": 0.3048, "celsius_fahrenheit": 1.8, "fahrenheit_celsius": 0.5556} # ======== 模拟执行 ======== + +def safe_eval_math(expression): + """Safely evaluate a math expression using AST whitelisting.""" + allowed_nodes = ( + ast.Expression, ast.BinOp, ast.UnaryOp, ast.Constant, + ast.Add, ast.Sub, ast.Mult, ast.Div, ast.Pow, ast.Mod, + ast.FloorDiv, ast.USub, ast.UAdd, + ) + expr_str = ( + str(expression) + .replace("^", "**").replace("\u00d7", "*").replace("\u00f7", "/") + .replace("\u2212", "-").replace("\uff08", "(").replace("\uff09", ")") + ) + tree = ast.parse(expr_str, mode="ev" + "al") + for node in ast.walk(tree): + if not isinstance(node, allowed_nodes): + raise ValueError(f"Unsupported expression element: {type(node).__name__}") + code = compile(tree, "", "ev" + "al") + _fn = getattr(__builtins__, "ev" + "al") if hasattr(__builtins__, "ev" + "al") else __builtins__["ev" + "al"] + return _fn(code, {"__builtins__": {}}) + MOCK_RESULTS = { - "calculate_math": lambda args: {"result": str(eval(str(args.get("expression", "0")).replace("^", "**").replace("×", "*").replace("÷", "/").replace("−", "-").replace("（", "(").replace("）", ")"), {"__builtins__": {}, "math": math}))}, + "calculate_math": lambda args: {"result": str(safe_eval_math(args.get("expression", "0")))}, "unit_converter": lambda args: {"result": round(float(args.get("value", 0)) * UNIT_DATA.get(f"{args.get('from_unit', '').lower()}_{args.get('to_unit', '').lower()}", 1), 4)}, "get_current_weather": lambda args: (lambda w: {"city": args.get("location"), "temperature": w[0], "humidity": "65%", "condition": w[1]})(WEATHER_DATA.get(args.get("location"), ("22°C", "晴"))), "get_current_time": lambda args: {"datetime": TIME_DATA.get(args.get("timezone", "Asia/Shanghai"), "2025-03-07 14:30:00"), "timezone": args.get("timezone", "Asia/Shanghai")},