TokioAI-Coder IoT Package v1.0 -- Engine + IoT System Prompt + Tools

4b97c3b verified about 1 month ago

10.7 kB

"""
TokioAI Engine v3.0 -- Server-Side Tool Executor
Executes tools LOCALLY on the Pi without CLI round-trips.
This is what makes the Engine think like Opus: it can investigate autonomously.
Safety: only execute_local, read_file, search_files, diagnose are allowed server-side.
Write operations (write_file, edit_file) are ALWAYS delegated to the CLI for user approval.
Remote operations (execute_raspi, execute_gcp, execute_router) are also delegated.
"""

import os
import re
import subprocess
import time
import glob as glob_module
from typing import Dict, Tuple, Optional


# Tools that the Engine can execute directly (read-only / diagnostic)
SAFE_TOOLS = {
    "execute_local",
    "read_file",
    "search_files",
    "diagnose",
    "robot_vision",
}

# Tools that must be delegated to the CLI (writes, remote)
DELEGATE_TOOLS = {
    "write_file",
    "edit_file",
    "execute_raspi",
    "execute_gcp",
    "execute_router",
    "memory",
    "task",
}

# Commands that are NEVER allowed server-side (destructive)
BLOCKED_COMMANDS = [
    r'\brm\s+-rf\s+/',
    r'\bmkfs\b',
    r'\bdd\s+if=',
    r'\b(reboot|shutdown|poweroff|halt)\b',
    r'\bsudo\s+(rm|mkfs|dd|reboot|shutdown|poweroff|halt)\b',
    r'>\s*/dev/',
    r'\bchmod\s+-R\s+777\s+/',
    r'\bchown\s+-R.*\s+/',
]


def is_safe_tool(tool_name: str) -> bool:
    """Check if a tool can be executed server-side."""
    return tool_name in SAFE_TOOLS


def is_command_blocked(command: str) -> bool:
    """Check if a shell command is blocked for safety."""
    for pattern in BLOCKED_COMMANDS:
        if re.search(pattern, command, re.IGNORECASE):
            return True
    return False


def execute_tool(tool_name: str, arguments: Dict, timeout: int = 30) -> Tuple[bool, str]:
    """
    Execute a tool server-side.
    
    Returns:
        (success: bool, output: str)
        
    If the tool is not in SAFE_TOOLS, returns (False, "DELEGATE") to signal
    that this tool should be sent back to the CLI.
    """
    if tool_name not in SAFE_TOOLS:
        return False, "DELEGATE"
    
    try:
        if tool_name == "execute_local":
            return _execute_local(arguments, timeout)
        elif tool_name == "read_file":
            return _read_file(arguments)
        elif tool_name == "search_files":
            return _search_files(arguments)
        elif tool_name == "diagnose":
            return _diagnose(arguments)
        elif tool_name == "robot_vision":
            return _robot_vision(arguments)
        else:
            return False, f"Unknown tool: {tool_name}"
    except Exception as e:
        return False, f"Error executing {tool_name}: {str(e)}"


def _execute_local(arguments: Dict, default_timeout: int = 30) -> Tuple[bool, str]:
    """Execute a shell command locally."""
    command = arguments.get("command", "")
    if not command:
        return False, "No command specified"
    
    if is_command_blocked(command):
        return False, f"Command blocked for safety: {command}"
    
    timeout = min(arguments.get("timeout", default_timeout), 60)  # Cap at 60s for server-side
    
    try:
        result = subprocess.run(
            command,
            shell=True,
            capture_output=True,
            text=True,
            timeout=timeout,
            env={**os.environ, "TERM": "dumb"},  # prevent color codes
        )
        
        output = result.stdout
        if result.stderr:
            output += ("\n" if output else "") + result.stderr
        
        output = output.strip()
        
        if not output:
            if result.returncode == 0:
                output = "(no output)"
            else:
                output = f"(exit code {result.returncode}, no output)"
        
        # Truncate very long output to prevent context explosion
        if len(output) > 3000:
            output = output[:2500] + f"\n... [truncated, {len(output)} total chars]"
        
        success = (result.returncode == 0)
        return success, output
        
    except subprocess.TimeoutExpired:
        return False, f"Command timed out after {timeout}s"
    except Exception as e:
        return False, f"Execution error: {str(e)}"


def _read_file(arguments: Dict) -> Tuple[bool, str]:
    """Read a file from the local filesystem."""
    path = arguments.get("path", "")
    if not path:
        return False, "No path specified"
    
    # Only allow local paths (no raspi: or gcp: prefix)
    if ":" in path and not path.startswith("/"):
        return False, "DELEGATE"  # Remote file, delegate to CLI
    
    # Expand user home
    path = os.path.expanduser(path)
    
    if not os.path.exists(path):
        return False, f"No such file or directory: {path}"
    
    if os.path.isdir(path):
        return False, f"Is a directory: {path}"
    
    max_lines = arguments.get("lines", None)
    
    try:
        with open(path, "r", errors="replace") as f:
            if max_lines:
                content = "".join(f.readline() for _ in range(max_lines))
            else:
                content = f.read()
        
        # Truncate
        if len(content) > 5000:
            content = content[:4500] + f"\n... [truncated, {len(content)} total chars]"
        
        return True, content.strip() if content.strip() else "(empty file)"
        
    except PermissionError:
        return False, f"Permission denied: {path}"
    except Exception as e:
        return False, f"Error reading {path}: {str(e)}"


def _search_files(arguments: Dict) -> Tuple[bool, str]:
    """Search for patterns in files (like grep -rn)."""
    pattern = arguments.get("pattern", "")
    if not pattern:
        return False, "No pattern specified"
    
    search_path = arguments.get("path", ".")
    
    # Only local paths
    if ":" in search_path and not search_path.startswith("/"):
        return False, "DELEGATE"
    
    search_path = os.path.expanduser(search_path)
    file_glob = arguments.get("glob", "*")
    
    try:
        # Use grep for efficiency
        cmd = f"grep -rn --include='{file_glob}' '{pattern}' {search_path} 2>/dev/null | head -50"
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=15)
        output = result.stdout.strip()
        
        if not output:
            return True, f"No matches found for '{pattern}' in {search_path}"
        
        if len(output) > 3000:
            output = output[:2500] + "\n... [truncated]"
        
        return True, output
        
    except Exception as e:
        return False, f"Search error: {str(e)}"


def _diagnose(arguments: Dict) -> Tuple[bool, str]:
    """Run basic diagnostics on the local system."""
    target = arguments.get("target", "local")
    
    if target != "local":
        return False, "DELEGATE"  # Remote diagnostics delegate to CLI
    
    checks = []
    
    # CPU load
    try:
        with open("/proc/loadavg", "r") as f:
            load = f.read().strip()
        checks.append(f"Load: {load}")
    except:
        pass
    
    # Memory
    try:
        result = subprocess.run("free -h | head -3", shell=True, capture_output=True, text=True, timeout=5)
        checks.append(f"Memory:\n{result.stdout.strip()}")
    except:
        pass
    
    # Disk
    try:
        result = subprocess.run("df -h / | tail -1", shell=True, capture_output=True, text=True, timeout=5)
        checks.append(f"Disk: {result.stdout.strip()}")
    except:
        pass
    
    # Docker
    try:
        result = subprocess.run("docker ps --format 'table {{.Names}}\t{{.Status}}' 2>/dev/null | head -10", shell=True, capture_output=True, text=True, timeout=5)
        if result.returncode == 0 and result.stdout.strip():
            checks.append(f"Docker:\n{result.stdout.strip()}")
    except:
        pass
    
    # Key services
    try:
        result = subprocess.run("systemctl list-units --type=service --state=running --no-pager | head -20", shell=True, capture_output=True, text=True, timeout=5)
        checks.append(f"Services:\n{result.stdout.strip()}")
    except:
        pass
    
    return True, "\n\n".join(checks) if checks else "Could not gather diagnostics"


def _robot_vision(arguments: Dict) -> Tuple[bool, str]:
    """Capture a snapshot from the robot's camera and analyze with moondream vision model."""
    import requests
    import base64
    import time

    ROBOT_PROXY = os.environ.get("ROBOT_PROXY_URL", "http://localhost:5002")
    OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:8200")
    VISION_MODEL = os.environ.get("VISION_MODEL", "moondream")

    question = arguments.get("question", "")
    if not question:
        question = "Describe what you see in this image in detail. What objects, people, furniture, or environment is visible? Be specific."

    # 1. Capture snapshot from robot camera
    t0 = time.time()
    try:
        resp = requests.get(f"{ROBOT_PROXY}/snapshot", timeout=10)
        resp.raise_for_status()
        img_data = resp.content
    except requests.ConnectionError:
        return False, "Robot camera unavailable -- cannot connect to robot proxy (localhost:5002). Is picar-forward.service running?"
    except Exception as e:
        return False, f"Error capturing snapshot from robot camera: {e}"

    t_capture = time.time() - t0

    if len(img_data) < 500:
        return False, f"Robot camera returned suspiciously small image ({len(img_data)} bytes). Camera may be offline."

    # 2. Encode to base64
    img_b64 = base64.b64encode(img_data).decode('utf-8')

    # 3. Analyze with vision model
    t1 = time.time()
    try:
        resp = requests.post(
            f"{OLLAMA_URL}/api/generate",
            json={
                "model": VISION_MODEL,
                "prompt": question,
                "images": [img_b64],
                "stream": False,
            },
            timeout=60,
        )
        resp.raise_for_status()
        result = resp.json()
    except requests.ConnectionError:
        return False, "Vision model unavailable -- cannot connect to Ollama (localhost:8200). Is tokioai-coder-tunnel.service running?"
    except Exception as e:
        return False, f"Error analyzing image with vision model: {e}"

    t_analysis = time.time() - t1

    description = result.get("response", "").strip()
    if not description:
        return False, "Vision model returned empty response."

    tokens = result.get("eval_count", 0)

    output = (
        f"[Robot Vision -- moondream]\n"
        f"Image: {len(img_data):,} bytes | Capture: {t_capture:.1f}s | Analysis: {t_analysis:.1f}s | Tokens: {tokens}\n"
        f"Question: {question}\n"
        f"Answer: {description}"
    )

    return True, output

	"""
	TokioAI Engine v3.0 -- Server-Side Tool Executor
	Executes tools LOCALLY on the Pi without CLI round-trips.
	This is what makes the Engine think like Opus: it can investigate autonomously.

	Safety: only execute_local, read_file, search_files, diagnose are allowed server-side.
	Write operations (write_file, edit_file) are ALWAYS delegated to the CLI for user approval.
	Remote operations (execute_raspi, execute_gcp, execute_router) are also delegated.
	"""

	import os
	import re
	import subprocess
	import time
	import glob as glob_module
	from typing import Dict, Tuple, Optional


	# Tools that the Engine can execute directly (read-only / diagnostic)
	SAFE_TOOLS = {
	"execute_local",
	"read_file",
	"search_files",
	"diagnose",
	"robot_vision",
	}

	# Tools that must be delegated to the CLI (writes, remote)
	DELEGATE_TOOLS = {
	"write_file",
	"edit_file",
	"execute_raspi",
	"execute_gcp",
	"execute_router",
	"memory",
	"task",
	}

	# Commands that are NEVER allowed server-side (destructive)
	BLOCKED_COMMANDS = [
	r'\brm\s+-rf\s+/',
	r'\bmkfs\b',
	r'\bdd\s+if=',
	r'\b(reboot\|shutdown\|poweroff\|halt)\b',
	r'\bsudo\s+(rm\|mkfs\|dd\|reboot\|shutdown\|poweroff\|halt)\b',
	r'>\s*/dev/',
	r'\bchmod\s+-R\s+777\s+/',
	r'\bchown\s+-R.*\s+/',
	]


	def is_safe_tool(tool_name: str) -> bool:
	"""Check if a tool can be executed server-side."""
	return tool_name in SAFE_TOOLS


	def is_command_blocked(command: str) -> bool:
	"""Check if a shell command is blocked for safety."""
	for pattern in BLOCKED_COMMANDS:
	if re.search(pattern, command, re.IGNORECASE):
	return True
	return False


	def execute_tool(tool_name: str, arguments: Dict, timeout: int = 30) -> Tuple[bool, str]:
	"""
	Execute a tool server-side.

	Returns:
	(success: bool, output: str)

	If the tool is not in SAFE_TOOLS, returns (False, "DELEGATE") to signal
	that this tool should be sent back to the CLI.
	"""
	if tool_name not in SAFE_TOOLS:
	return False, "DELEGATE"

	try:
	if tool_name == "execute_local":
	return _execute_local(arguments, timeout)
	elif tool_name == "read_file":
	return _read_file(arguments)
	elif tool_name == "search_files":
	return _search_files(arguments)
	elif tool_name == "diagnose":
	return _diagnose(arguments)
	elif tool_name == "robot_vision":
	return _robot_vision(arguments)
	else:
	return False, f"Unknown tool: {tool_name}"
	except Exception as e:
	return False, f"Error executing {tool_name}: {str(e)}"


	def _execute_local(arguments: Dict, default_timeout: int = 30) -> Tuple[bool, str]:
	"""Execute a shell command locally."""
	command = arguments.get("command", "")
	if not command:
	return False, "No command specified"

	if is_command_blocked(command):
	return False, f"Command blocked for safety: {command}"

	timeout = min(arguments.get("timeout", default_timeout), 60) # Cap at 60s for server-side

	try:
	result = subprocess.run(
	command,
	shell=True,
	capture_output=True,
	text=True,
	timeout=timeout,
	env={**os.environ, "TERM": "dumb"}, # prevent color codes
	)

	output = result.stdout
	if result.stderr:
	output += ("\n" if output else "") + result.stderr

	output = output.strip()

	if not output:
	if result.returncode == 0:
	output = "(no output)"
	else:
	output = f"(exit code {result.returncode}, no output)"

	# Truncate very long output to prevent context explosion
	if len(output) > 3000:
	output = output[:2500] + f"\n... [truncated, {len(output)} total chars]"

	success = (result.returncode == 0)
	return success, output

	except subprocess.TimeoutExpired:
	return False, f"Command timed out after {timeout}s"
	except Exception as e:
	return False, f"Execution error: {str(e)}"


	def _read_file(arguments: Dict) -> Tuple[bool, str]:
	"""Read a file from the local filesystem."""
	path = arguments.get("path", "")
	if not path:
	return False, "No path specified"

	# Only allow local paths (no raspi: or gcp: prefix)
	if ":" in path and not path.startswith("/"):
	return False, "DELEGATE" # Remote file, delegate to CLI

	# Expand user home
	path = os.path.expanduser(path)

	if not os.path.exists(path):
	return False, f"No such file or directory: {path}"

	if os.path.isdir(path):
	return False, f"Is a directory: {path}"

	max_lines = arguments.get("lines", None)

	try:
	with open(path, "r", errors="replace") as f:
	if max_lines:
	content = "".join(f.readline() for _ in range(max_lines))
	else:
	content = f.read()

	# Truncate
	if len(content) > 5000:
	content = content[:4500] + f"\n... [truncated, {len(content)} total chars]"

	return True, content.strip() if content.strip() else "(empty file)"

	except PermissionError:
	return False, f"Permission denied: {path}"
	except Exception as e:
	return False, f"Error reading {path}: {str(e)}"


	def _search_files(arguments: Dict) -> Tuple[bool, str]:
	"""Search for patterns in files (like grep -rn)."""
	pattern = arguments.get("pattern", "")
	if not pattern:
	return False, "No pattern specified"

	search_path = arguments.get("path", ".")

	# Only local paths
	if ":" in search_path and not search_path.startswith("/"):
	return False, "DELEGATE"

	search_path = os.path.expanduser(search_path)
	file_glob = arguments.get("glob", "*")

	try:
	# Use grep for efficiency
	cmd = f"grep -rn --include='{file_glob}' '{pattern}' {search_path} 2>/dev/null \| head -50"
	result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=15)
	output = result.stdout.strip()

	if not output:
	return True, f"No matches found for '{pattern}' in {search_path}"

	if len(output) > 3000:
	output = output[:2500] + "\n... [truncated]"

	return True, output

	except Exception as e:
	return False, f"Search error: {str(e)}"


	def _diagnose(arguments: Dict) -> Tuple[bool, str]:
	"""Run basic diagnostics on the local system."""
	target = arguments.get("target", "local")

	if target != "local":
	return False, "DELEGATE" # Remote diagnostics delegate to CLI

	checks = []

	# CPU load
	try:
	with open("/proc/loadavg", "r") as f:
	load = f.read().strip()
	checks.append(f"Load: {load}")
	except:
	pass

	# Memory
	try:
	result = subprocess.run("free -h \| head -3", shell=True, capture_output=True, text=True, timeout=5)
	checks.append(f"Memory:\n{result.stdout.strip()}")
	except:
	pass

	# Disk
	try:
	result = subprocess.run("df -h / \| tail -1", shell=True, capture_output=True, text=True, timeout=5)
	checks.append(f"Disk: {result.stdout.strip()}")
	except:
	pass

	# Docker
	try:
	result = subprocess.run("docker ps --format 'table {{.Names}}\t{{.Status}}' 2>/dev/null \| head -10", shell=True, capture_output=True, text=True, timeout=5)
	if result.returncode == 0 and result.stdout.strip():
	checks.append(f"Docker:\n{result.stdout.strip()}")
	except:
	pass

	# Key services
	try:
	result = subprocess.run("systemctl list-units --type=service --state=running --no-pager \| head -20", shell=True, capture_output=True, text=True, timeout=5)
	checks.append(f"Services:\n{result.stdout.strip()}")
	except:
	pass

	return True, "\n\n".join(checks) if checks else "Could not gather diagnostics"


	def _robot_vision(arguments: Dict) -> Tuple[bool, str]:
	"""Capture a snapshot from the robot's camera and analyze with moondream vision model."""
	import requests
	import base64
	import time

	ROBOT_PROXY = os.environ.get("ROBOT_PROXY_URL", "http://localhost:5002")
	OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:8200")
	VISION_MODEL = os.environ.get("VISION_MODEL", "moondream")

	question = arguments.get("question", "")
	if not question:
	question = "Describe what you see in this image in detail. What objects, people, furniture, or environment is visible? Be specific."

	# 1. Capture snapshot from robot camera
	t0 = time.time()
	try:
	resp = requests.get(f"{ROBOT_PROXY}/snapshot", timeout=10)
	resp.raise_for_status()
	img_data = resp.content
	except requests.ConnectionError:
	return False, "Robot camera unavailable -- cannot connect to robot proxy (localhost:5002). Is picar-forward.service running?"
	except Exception as e:
	return False, f"Error capturing snapshot from robot camera: {e}"

	t_capture = time.time() - t0

	if len(img_data) < 500:
	return False, f"Robot camera returned suspiciously small image ({len(img_data)} bytes). Camera may be offline."

	# 2. Encode to base64
	img_b64 = base64.b64encode(img_data).decode('utf-8')

	# 3. Analyze with vision model
	t1 = time.time()
	try:
	resp = requests.post(
	f"{OLLAMA_URL}/api/generate",
	json={
	"model": VISION_MODEL,
	"prompt": question,
	"images": [img_b64],
	"stream": False,
	},
	timeout=60,
	)
	resp.raise_for_status()
	result = resp.json()
	except requests.ConnectionError:
	return False, "Vision model unavailable -- cannot connect to Ollama (localhost:8200). Is tokioai-coder-tunnel.service running?"
	except Exception as e:
	return False, f"Error analyzing image with vision model: {e}"

	t_analysis = time.time() - t1

	description = result.get("response", "").strip()
	if not description:
	return False, "Vision model returned empty response."

	tokens = result.get("eval_count", 0)

	output = (
	f"[Robot Vision -- moondream]\n"
	f"Image: {len(img_data):,} bytes \| Capture: {t_capture:.1f}s \| Analysis: {t_analysis:.1f}s \| Tokens: {tokens}\n"
	f"Question: {question}\n"
	f"Answer: {description}"
	)

	return True, output