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

4b97c3b verified about 1 month ago

6.4 kB

"""
TokioAI Engine -- Output Verifier & Retry Logic
Validates model responses and orchestrates retries on failure.
The verifier catches:
1. Invalid JSON tool calls
2. Wrong tool names
3. Missing required parameters
4. Malformed output (e.g., natural language when tool expected)
5. Verbosity (strips filler from responses)
"""

import json
import re
import time
from typing import Any, Dict, List, Optional, Tuple

from .schemas import validate_tool_call, extract_tool_calls_from_text, VALID_TOOL_NAMES


class VerificationResult:
    """Result of verifying model output."""

    def __init__(
        self,
        valid: bool,
        output_type: str,  # "tool_call", "text", "error"
        tool_calls: List[Dict] = None,
        text: str = None,
        error: str = None,
        raw_output: str = None,
    ):
        self.valid = valid
        self.output_type = output_type
        self.tool_calls = tool_calls or []
        self.text = text or ""
        self.error = error or ""
        self.raw_output = raw_output or ""

    def __repr__(self):
        if self.valid:
            if self.output_type == "tool_call":
                names = [tc["name"] for tc in self.tool_calls]
                return f"VerificationResult(valid=True, tools={names})"
            return f"VerificationResult(valid=True, type={self.output_type})"
        return f"VerificationResult(valid=False, error={self.error!r})"


class Verifier:
    """
    Validates model output and provides structured results.
    """

    def __init__(self, strict: bool = False):
        """
        Args:
            strict: If True, rejects ambiguous outputs. If False, tries harder to parse.
        """
        self.strict = strict

    def verify(self, model_output: str, expected_intent: str = None) -> VerificationResult:
        """
        Verify model output.
        
        Args:
            model_output: Raw text from the model
            expected_intent: What the router classified (tool_call, text, etc.)
            
        Returns:
            VerificationResult with parsed tool calls or cleaned text
        """
        if not model_output or not model_output.strip():
            return VerificationResult(
                valid=False,
                output_type="error",
                error="Empty model output",
                raw_output=model_output,
            )

        raw = model_output.strip()

        # Step 1: Try to extract tool calls
        tool_calls = extract_tool_calls_from_text(raw)

        if tool_calls:
            # Validate each tool call
            valid_calls = []
            errors = []

            for tc in tool_calls:
                is_valid, err = validate_tool_call(tc["name"], tc["arguments"])
                if is_valid:
                    valid_calls.append(tc)
                else:
                    errors.append(err)

            if valid_calls:
                return VerificationResult(
                    valid=True,
                    output_type="tool_call",
                    tool_calls=valid_calls,
                    raw_output=raw,
                )
            elif errors:
                return VerificationResult(
                    valid=False,
                    output_type="error",
                    error="; ".join(errors),
                    raw_output=raw,
                )

        # Step 2: If expected tool_call but got text, that might be an error
        if expected_intent == "tool_call":
            # Check if the model is TRYING to make a tool call but failing
            if any(name in raw.lower() for name in VALID_TOOL_NAMES):
                return VerificationResult(
                    valid=False,
                    output_type="error",
                    error="Model mentioned a tool name but didn't format it as a valid JSON tool call",
                    raw_output=raw,
                )

        # Step 3: It's a text response -- clean it up
        cleaned = self._clean_text(raw)
        return VerificationResult(
            valid=True,
            output_type="text",
            text=cleaned,
            raw_output=raw,
        )

    def _clean_text(self, text: str) -> str:
        """Clean up verbose model output."""
        # Remove thinking tags (Qwen3 /think blocks)
        text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
        # Remove empty lines at start/end
        text = text.strip()
        # Remove excessive newlines
        text = re.sub(r'\n{3,}', '\n\n', text)
        return text


class RetryStrategy:
    """
    Manages retry logic when the verifier rejects model output.
    Each retry gets a more explicit prompt to guide the model.
    """

    def __init__(self, max_retries: int = 2):
        self.max_retries = max_retries

    def should_retry(self, attempt: int, result: VerificationResult) -> bool:
        """Should we retry after a failed verification?"""
        if result.valid:
            return False
        if attempt >= self.max_retries:
            return False
        # Don't retry empty outputs (model is broken)
        if not result.raw_output.strip():
            return False
        return True

    def get_retry_message(self, attempt: int, result: VerificationResult) -> str:
        """Generate a retry prompt based on what went wrong."""
        if attempt == 0:
            # First retry: gentle correction
            return (
                f"Your previous response was not valid: {result.error}\n"
                f"Please respond with ONLY a JSON object like this:\n"
                f'{{"name": "tool_name", "arguments": {{"param": "value"}}}}\n'
                f"No explanation, no markdown, just the JSON."
            )
        else:
            # Second retry: very explicit
            return (
                f"ERROR: {result.error}\n"
                f"You MUST respond with valid JSON only. "
                f"Valid tool names: {', '.join(sorted(VALID_TOOL_NAMES))}.\n"
                f"Format: {{\"name\": \"execute_local\", \"arguments\": {{\"command\": \"your command\"}}}}\n"
                f"Respond with NOTHING except the JSON object."
            )


def verify_and_extract(model_output: str, expected_intent: str = None) -> VerificationResult:
    """Convenience function: verify model output with default settings."""
    return Verifier().verify(model_output, expected_intent)

	"""
	TokioAI Engine -- Output Verifier & Retry Logic
	Validates model responses and orchestrates retries on failure.

	The verifier catches:
	1. Invalid JSON tool calls
	2. Wrong tool names
	3. Missing required parameters
	4. Malformed output (e.g., natural language when tool expected)
	5. Verbosity (strips filler from responses)
	"""

	import json
	import re
	import time
	from typing import Any, Dict, List, Optional, Tuple

	from .schemas import validate_tool_call, extract_tool_calls_from_text, VALID_TOOL_NAMES


	class VerificationResult:
	"""Result of verifying model output."""

	def __init__(
	self,
	valid: bool,
	output_type: str, # "tool_call", "text", "error"
	tool_calls: List[Dict] = None,
	text: str = None,
	error: str = None,
	raw_output: str = None,
	):
	self.valid = valid
	self.output_type = output_type
	self.tool_calls = tool_calls or []
	self.text = text or ""
	self.error = error or ""
	self.raw_output = raw_output or ""

	def __repr__(self):
	if self.valid:
	if self.output_type == "tool_call":
	names = [tc["name"] for tc in self.tool_calls]
	return f"VerificationResult(valid=True, tools={names})"
	return f"VerificationResult(valid=True, type={self.output_type})"
	return f"VerificationResult(valid=False, error={self.error!r})"


	class Verifier:
	"""
	Validates model output and provides structured results.
	"""

	def __init__(self, strict: bool = False):
	"""
	Args:
	strict: If True, rejects ambiguous outputs. If False, tries harder to parse.
	"""
	self.strict = strict

	def verify(self, model_output: str, expected_intent: str = None) -> VerificationResult:
	"""
	Verify model output.

	Args:
	model_output: Raw text from the model
	expected_intent: What the router classified (tool_call, text, etc.)

	Returns:
	VerificationResult with parsed tool calls or cleaned text
	"""
	if not model_output or not model_output.strip():
	return VerificationResult(
	valid=False,
	output_type="error",
	error="Empty model output",
	raw_output=model_output,
	)

	raw = model_output.strip()

	# Step 1: Try to extract tool calls
	tool_calls = extract_tool_calls_from_text(raw)

	if tool_calls:
	# Validate each tool call
	valid_calls = []
	errors = []

	for tc in tool_calls:
	is_valid, err = validate_tool_call(tc["name"], tc["arguments"])
	if is_valid:
	valid_calls.append(tc)
	else:
	errors.append(err)

	if valid_calls:
	return VerificationResult(
	valid=True,
	output_type="tool_call",
	tool_calls=valid_calls,
	raw_output=raw,
	)
	elif errors:
	return VerificationResult(
	valid=False,
	output_type="error",
	error="; ".join(errors),
	raw_output=raw,
	)

	# Step 2: If expected tool_call but got text, that might be an error
	if expected_intent == "tool_call":
	# Check if the model is TRYING to make a tool call but failing
	if any(name in raw.lower() for name in VALID_TOOL_NAMES):
	return VerificationResult(
	valid=False,
	output_type="error",
	error="Model mentioned a tool name but didn't format it as a valid JSON tool call",
	raw_output=raw,
	)

	# Step 3: It's a text response -- clean it up
	cleaned = self._clean_text(raw)
	return VerificationResult(
	valid=True,
	output_type="text",
	text=cleaned,
	raw_output=raw,
	)

	def _clean_text(self, text: str) -> str:
	"""Clean up verbose model output."""
	# Remove thinking tags (Qwen3 /think blocks)
	text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
	# Remove empty lines at start/end
	text = text.strip()
	# Remove excessive newlines
	text = re.sub(r'\n{3,}', '\n\n', text)
	return text


	class RetryStrategy:
	"""
	Manages retry logic when the verifier rejects model output.
	Each retry gets a more explicit prompt to guide the model.
	"""

	def __init__(self, max_retries: int = 2):
	self.max_retries = max_retries

	def should_retry(self, attempt: int, result: VerificationResult) -> bool:
	"""Should we retry after a failed verification?"""
	if result.valid:
	return False
	if attempt >= self.max_retries:
	return False
	# Don't retry empty outputs (model is broken)
	if not result.raw_output.strip():
	return False
	return True

	def get_retry_message(self, attempt: int, result: VerificationResult) -> str:
	"""Generate a retry prompt based on what went wrong."""
	if attempt == 0:
	# First retry: gentle correction
	return (
	f"Your previous response was not valid: {result.error}\n"
	f"Please respond with ONLY a JSON object like this:\n"
	f'{{"name": "tool_name", "arguments": {{"param": "value"}}}}\n'
	f"No explanation, no markdown, just the JSON."
	)
	else:
	# Second retry: very explicit
	return (
	f"ERROR: {result.error}\n"
	f"You MUST respond with valid JSON only. "
	f"Valid tool names: {', '.join(sorted(VALID_TOOL_NAMES))}.\n"
	f"Format: {{\"name\": \"execute_local\", \"arguments\": {{\"command\": \"your command\"}}}}\n"
	f"Respond with NOTHING except the JSON object."
	)


	def verify_and_extract(model_output: str, expected_intent: str = None) -> VerificationResult:
	"""Convenience function: verify model output with default settings."""
	return Verifier().verify(model_output, expected_intent)