Luxx/alcor/tools/builtin/data.py

"""数据处理工具"""
import re
import json
import hashlib
import base64
import urllib.parse
from typing import Dict, Any, List

from alcor.tools.factory import tool


@tool(
    name="calculate",
    description="Perform mathematical calculations",
    parameters={
        "type": "object",
        "properties": {
            "expression": {
                "type": "string",
                "description": "Mathematical expression to evaluate (e.g., '2 + 2', 'sqrt(16)', 'sin(pi/2)')"
            }
        },
        "required": ["expression"]
    },
    category="data"
)
def calculate(arguments: Dict[str, Any]) -> Dict[str, Any]:
    """执行数学计算"""
    expression = arguments.get("expression", "")
    
    if not expression:
        return {"success": False, "error": "Expression is required"}
    
    try:
        # 安全替换数学函数
        safe_dict = {
            "abs": abs,
            "round": round,
            "min": min,
            "max": max,
            "sum": sum,
            "pow": pow,
            "sqrt": lambda x: x ** 0.5,
            "sin": lambda x: __import__("math").sin(x),
            "cos": lambda x: __import__("math").cos(x),
            "tan": lambda x: __import__("math").tan(x),
            "log": lambda x: __import__("math").log(x),
            "pi": __import__("math").pi,
            "e": __import__("math").e,
        }
        
        # 移除危险字符，只保留数字和运算符
        safe_expr = re.sub(r"[^0-9+\-*/().%sqrtinsclogmaxminpowabsroundte, ]", "", expression)
        result = eval(safe_expr, {"__builtins__": {}, **safe_dict})
        
        return {
            "success": True,
            "data": {
                "expression": expression,
                "result": float(result) if isinstance(result, (int, float)) else result
            }
        }
    except ZeroDivisionError:
        return {"success": False, "error": "Division by zero"}
    except Exception as e:
        return {"success": False, "error": f"Calculation error: {str(e)}"}


@tool(
    name="text_process",
    description="Process and transform text",
    parameters={
        "type": "object",
        "properties": {
            "text": {
                "type": "string",
                "description": "Input text"
            },
            "operation": {
                "type": "string",
                "description": "Operation to perform: upper, lower, title, reverse, word_count, char_count, reverse_words",
                "enum": ["upper", "lower", "title", "reverse", "word_count", "char_count", "reverse_words"]
            }
        },
        "required": ["text", "operation"]
    },
    category="data"
)
def text_process(arguments: Dict[str, Any]) -> Dict[str, Any]:
    """文本处理"""
    text = arguments.get("text", "")
    operation = arguments.get("operation", "")
    
    if not text:
        return {"success": False, "error": "Text is required"}
    
    operations = {
        "upper": lambda t: t.upper(),
        "lower": lambda t: t.lower(),
        "title": lambda t: t.title(),
        "reverse": lambda t: t[::-1],
        "word_count": lambda t: len(t.split()),
        "char_count": lambda t: len(t),
        "reverse_words": lambda t: " ".join(t.split()[::-1])
    }
    
    if operation not in operations:
        return {
            "success": False, 
            "error": f"Unknown operation: {operation}"
        }
    
    try:
        result = operations[operation](text)
        return {
            "success": True,
            "data": {
                "operation": operation,
                "input": text,
                "result": result
            }
        }
    except Exception as e:
        return {"success": False, "error": str(e)}


@tool(
    name="hash_text",
    description="Generate hash of text using various algorithms",
    parameters={
        "type": "object",
        "properties": {
            "text": {
                "type": "string",
                "description": "Text to hash"
            },
            "algorithm": {
                "type": "string",
                "description": "Hash algorithm: md5, sha1, sha256, sha512",
                "default": "sha256"
            }
        },
        "required": ["text"]
    },
    category="data"
)
def hash_text(arguments: Dict[str, Any]) -> Dict[str, Any]:
    """生成文本哈希"""
    text = arguments.get("text", "")
    algorithm = arguments.get("algorithm", "sha256")
    
    if not text:
        return {"success": False, "error": "Text is required"}
    
    hash_funcs = {
        "md5": hashlib.md5,
        "sha1": hashlib.sha1,
        "sha256": hashlib.sha256,
        "sha512": hashlib.sha512
    }
    
    if algorithm not in hash_funcs:
        return {
            "success": False,
            "error": f"Unsupported algorithm: {algorithm}"
        }
    
    try:
        hash_obj = hash_funcs[algorithm](text.encode("utf-8"))
        return {
            "success": True,
            "data": {
                "algorithm": algorithm,
                "hash": hash_obj.hexdigest()
            }
        }
    except Exception as e:
        return {"success": False, "error": str(e)}


@tool(
    name="url_encode_decode",
    description="URL encode or decode text",
    parameters={
        "type": "object",
        "properties": {
            "text": {
                "type": "string",
                "description": "Text to encode or decode"
            },
            "operation": {
                "type": "string",
                "description": "Operation: encode or decode",
                "enum": ["encode", "decode"]
            }
        },
        "required": ["text", "operation"]
    },
    category="data"
)
def url_encode_decode(arguments: Dict[str, Any]) -> Dict[str, Any]:
    """URL编码/解码"""
    text = arguments.get("text", "")
    operation = arguments.get("operation", "encode")
    
    if not text:
        return {"success": False, "error": "Text is required"}
    
    try:
        if operation == "encode":
            result = urllib.parse.quote(text)
        else:
            result = urllib.parse.unquote(text)
        
        return {
            "success": True,
            "data": {
                "operation": operation,
                "input": text,
                "result": result
            }
        }
    except Exception as e:
        return {"success": False, "error": str(e)}


@tool(
    name="base64_encode_decode",
    description="Base64 encode or decode text",
    parameters={
        "type": "object",
        "properties": {
            "text": {
                "type": "string",
                "description": "Text to encode or decode"
            },
            "operation": {
                "type": "string",
                "description": "Operation: encode or decode",
                "enum": ["encode", "decode"]
            }
        },
        "required": ["text", "operation"]
    },
    category="data"
)
def base64_encode_decode(arguments: Dict[str, Any]) -> Dict[str, Any]:
    """Base64编码/解码"""
    text = arguments.get("text", "")
    operation = arguments.get("operation", "encode")
    
    if not text:
        return {"success": False, "error": "Text is required"}
    
    try:
        if operation == "encode":
            result = base64.b64encode(text.encode("utf-8")).decode("utf-8")
        else:
            result = base64.b64decode(text.encode("utf-8")).decode("utf-8")
        
        return {
            "success": True,
            "data": {
                "operation": operation,
                "input": text,
                "result": result
            }
        }
    except Exception as e:
        return {"success": False, "error": str(e)}