Team Ai
Modelpublic

MrMoz33/tokioai-coder-iot

sourceHugging Facemitupdated 5d agoView on Hugging Face
0likes
executor.py332 linesDownload Raw Back to engine
1"""2TokioAI Engine v3.0 -- Server-Side Tool Executor3Executes tools LOCALLY on the Pi without CLI round-trips.4This is what makes the Engine think like Opus: it can investigate autonomously.5 6Safety: only execute_local, read_file, search_files, diagnose are allowed server-side.7Write operations (write_file, edit_file) are ALWAYS delegated to the CLI for user approval.8Remote operations (execute_raspi, execute_gcp, execute_router) are also delegated.9"""10 11import os12import re13import subprocess14import time15import glob as glob_module16from typing import Dict, Tuple, Optional17 18 19# Tools that the Engine can execute directly (read-only / diagnostic)20SAFE_TOOLS = {21    "execute_local",22    "read_file",23    "search_files",24    "diagnose",25    "robot_vision",26}27 28# Tools that must be delegated to the CLI (writes, remote)29DELEGATE_TOOLS = {30    "write_file",31    "edit_file",32    "execute_raspi",33    "execute_gcp",34    "execute_router",35    "memory",36    "task",37}38 39# Commands that are NEVER allowed server-side (destructive)40BLOCKED_COMMANDS = [41    r'\brm\s+-rf\s+/',42    r'\bmkfs\b',43    r'\bdd\s+if=',44    r'\b(reboot|shutdown|poweroff|halt)\b',45    r'\bsudo\s+(rm|mkfs|dd|reboot|shutdown|poweroff|halt)\b',46    r'>\s*/dev/',47    r'\bchmod\s+-R\s+777\s+/',48    r'\bchown\s+-R.*\s+/',49]50 51 52def is_safe_tool(tool_name: str) -> bool:53    """Check if a tool can be executed server-side."""54    return tool_name in SAFE_TOOLS55 56 57def is_command_blocked(command: str) -> bool:58    """Check if a shell command is blocked for safety."""59    for pattern in BLOCKED_COMMANDS:60        if re.search(pattern, command, re.IGNORECASE):61            return True62    return False63 64 65def execute_tool(tool_name: str, arguments: Dict, timeout: int = 30) -> Tuple[bool, str]:66    """67    Execute a tool server-side.68    69    Returns:70        (success: bool, output: str)71        72    If the tool is not in SAFE_TOOLS, returns (False, "DELEGATE") to signal73    that this tool should be sent back to the CLI.74    """75    if tool_name not in SAFE_TOOLS:76        return False, "DELEGATE"77    78    try:79        if tool_name == "execute_local":80            return _execute_local(arguments, timeout)81        elif tool_name == "read_file":82            return _read_file(arguments)83        elif tool_name == "search_files":84            return _search_files(arguments)85        elif tool_name == "diagnose":86            return _diagnose(arguments)87        elif tool_name == "robot_vision":88            return _robot_vision(arguments)89        else:90            return False, f"Unknown tool: {tool_name}"91    except Exception as e:92        return False, f"Error executing {tool_name}: {str(e)}"93 94 95def _execute_local(arguments: Dict, default_timeout: int = 30) -> Tuple[bool, str]:96    """Execute a shell command locally."""97    command = arguments.get("command", "")98    if not command:99        return False, "No command specified"100    101    if is_command_blocked(command):102        return False, f"Command blocked for safety: {command}"103    104    timeout = min(arguments.get("timeout", default_timeout), 60)  # Cap at 60s for server-side105    106    try:107        result = subprocess.run(108            command,109            shell=True,110            capture_output=True,111            text=True,112            timeout=timeout,113            env={**os.environ, "TERM": "dumb"},  # prevent color codes114        )115        116        output = result.stdout117        if result.stderr:118            output += ("\n" if output else "") + result.stderr119        120        output = output.strip()121        122        if not output:123            if result.returncode == 0:124                output = "(no output)"125            else:126                output = f"(exit code {result.returncode}, no output)"127        128        # Truncate very long output to prevent context explosion129        if len(output) > 3000:130            output = output[:2500] + f"\n... [truncated, {len(output)} total chars]"131        132        success = (result.returncode == 0)133        return success, output134        135    except subprocess.TimeoutExpired:136        return False, f"Command timed out after {timeout}s"137    except Exception as e:138        return False, f"Execution error: {str(e)}"139 140 141def _read_file(arguments: Dict) -> Tuple[bool, str]:142    """Read a file from the local filesystem."""143    path = arguments.get("path", "")144    if not path:145        return False, "No path specified"146    147    # Only allow local paths (no raspi: or gcp: prefix)148    if ":" in path and not path.startswith("/"):149        return False, "DELEGATE"  # Remote file, delegate to CLI150    151    # Expand user home152    path = os.path.expanduser(path)153    154    if not os.path.exists(path):155        return False, f"No such file or directory: {path}"156    157    if os.path.isdir(path):158        return False, f"Is a directory: {path}"159    160    max_lines = arguments.get("lines", None)161    162    try:163        with open(path, "r", errors="replace") as f:164            if max_lines:165                content = "".join(f.readline() for _ in range(max_lines))166            else:167                content = f.read()168        169        # Truncate170        if len(content) > 5000:171            content = content[:4500] + f"\n... [truncated, {len(content)} total chars]"172        173        return True, content.strip() if content.strip() else "(empty file)"174        175    except PermissionError:176        return False, f"Permission denied: {path}"177    except Exception as e:178        return False, f"Error reading {path}: {str(e)}"179 180 181def _search_files(arguments: Dict) -> Tuple[bool, str]:182    """Search for patterns in files (like grep -rn)."""183    pattern = arguments.get("pattern", "")184    if not pattern:185        return False, "No pattern specified"186    187    search_path = arguments.get("path", ".")188    189    # Only local paths190    if ":" in search_path and not search_path.startswith("/"):191        return False, "DELEGATE"192    193    search_path = os.path.expanduser(search_path)194    file_glob = arguments.get("glob", "*")195    196    try:197        # Use grep for efficiency198        cmd = f"grep -rn --include='{file_glob}' '{pattern}' {search_path} 2>/dev/null | head -50"199        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=15)200        output = result.stdout.strip()201        202        if not output:203            return True, f"No matches found for '{pattern}' in {search_path}"204        205        if len(output) > 3000:206            output = output[:2500] + "\n... [truncated]"207        208        return True, output209        210    except Exception as e:211        return False, f"Search error: {str(e)}"212 213 214def _diagnose(arguments: Dict) -> Tuple[bool, str]:215    """Run basic diagnostics on the local system."""216    target = arguments.get("target", "local")217    218    if target != "local":219        return False, "DELEGATE"  # Remote diagnostics delegate to CLI220    221    checks = []222    223    # CPU load224    try:225        with open("/proc/loadavg", "r") as f:226            load = f.read().strip()227        checks.append(f"Load: {load}")228    except:229        pass230    231    # Memory232    try:233        result = subprocess.run("free -h | head -3", shell=True, capture_output=True, text=True, timeout=5)234        checks.append(f"Memory:\n{result.stdout.strip()}")235    except:236        pass237    238    # Disk239    try:240        result = subprocess.run("df -h / | tail -1", shell=True, capture_output=True, text=True, timeout=5)241        checks.append(f"Disk: {result.stdout.strip()}")242    except:243        pass244    245    # Docker246    try:247        result = subprocess.run("docker ps --format 'table {{.Names}}\t{{.Status}}' 2>/dev/null | head -10", shell=True, capture_output=True, text=True, timeout=5)248        if result.returncode == 0 and result.stdout.strip():249            checks.append(f"Docker:\n{result.stdout.strip()}")250    except:251        pass252    253    # Key services254    try:255        result = subprocess.run("systemctl list-units --type=service --state=running --no-pager | head -20", shell=True, capture_output=True, text=True, timeout=5)256        checks.append(f"Services:\n{result.stdout.strip()}")257    except:258        pass259    260    return True, "\n\n".join(checks) if checks else "Could not gather diagnostics"261 262 263def _robot_vision(arguments: Dict) -> Tuple[bool, str]:264    """Capture a snapshot from the robot's camera and analyze with moondream vision model."""265    import requests266    import base64267    import time268 269    ROBOT_PROXY = os.environ.get("ROBOT_PROXY_URL", "http://localhost:5002")270    OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:8200")271    VISION_MODEL = os.environ.get("VISION_MODEL", "moondream")272 273    question = arguments.get("question", "")274    if not question:275        question = "Describe what you see in this image in detail. What objects, people, furniture, or environment is visible? Be specific."276 277    # 1. Capture snapshot from robot camera278    t0 = time.time()279    try:280        resp = requests.get(f"{ROBOT_PROXY}/snapshot", timeout=10)281        resp.raise_for_status()282        img_data = resp.content283    except requests.ConnectionError:284        return False, "Robot camera unavailable -- cannot connect to robot proxy (localhost:5002). Is picar-forward.service running?"285    except Exception as e:286        return False, f"Error capturing snapshot from robot camera: {e}"287 288    t_capture = time.time() - t0289 290    if len(img_data) < 500:291        return False, f"Robot camera returned suspiciously small image ({len(img_data)} bytes). Camera may be offline."292 293    # 2. Encode to base64294    img_b64 = base64.b64encode(img_data).decode('utf-8')295 296    # 3. Analyze with vision model297    t1 = time.time()298    try:299        resp = requests.post(300            f"{OLLAMA_URL}/api/generate",301            json={302                "model": VISION_MODEL,303                "prompt": question,304                "images": [img_b64],305                "stream": False,306            },307            timeout=60,308        )309        resp.raise_for_status()310        result = resp.json()311    except requests.ConnectionError:312        return False, "Vision model unavailable -- cannot connect to Ollama (localhost:8200). Is tokioai-coder-tunnel.service running?"313    except Exception as e:314        return False, f"Error analyzing image with vision model: {e}"315 316    t_analysis = time.time() - t1317 318    description = result.get("response", "").strip()319    if not description:320        return False, "Vision model returned empty response."321 322    tokens = result.get("eval_count", 0)323 324    output = (325        f"[Robot Vision -- moondream]\n"326        f"Image: {len(img_data):,} bytes | Capture: {t_capture:.1f}s | Analysis: {t_analysis:.1f}s | Tokens: {tokens}\n"327        f"Question: {question}\n"328        f"Answer: {description}"329    )330 331    return True, output332