MrMoz33/tokioai-coder-iot

1"""2TokioAI Engine -- Output Verifier & Retry Logic3Validates model responses and orchestrates retries on failure.4 5The verifier catches:61. Invalid JSON tool calls72. Wrong tool names83. Missing required parameters94. Malformed output (e.g., natural language when tool expected)105. Verbosity (strips filler from responses)11"""12 13import json14import re15import time16from typing import Any, Dict, List, Optional, Tuple17 18from .schemas import validate_tool_call, extract_tool_calls_from_text, VALID_TOOL_NAMES19 20 21class VerificationResult:22    """Result of verifying model output."""23 24    def __init__(25        self,26        valid: bool,27        output_type: str,  # "tool_call", "text", "error"28        tool_calls: List[Dict] = None,29        text: str = None,30        error: str = None,31        raw_output: str = None,32    ):33        self.valid = valid34        self.output_type = output_type35        self.tool_calls = tool_calls or []36        self.text = text or ""37        self.error = error or ""38        self.raw_output = raw_output or ""39 40    def __repr__(self):41        if self.valid:42            if self.output_type == "tool_call":43                names = [tc["name"] for tc in self.tool_calls]44                return f"VerificationResult(valid=True, tools={names})"45            return f"VerificationResult(valid=True, type={self.output_type})"46        return f"VerificationResult(valid=False, error={self.error!r})"47 48 49class Verifier:50    """51    Validates model output and provides structured results.52    """53 54    def __init__(self, strict: bool = False):55        """56        Args:57            strict: If True, rejects ambiguous outputs. If False, tries harder to parse.58        """59        self.strict = strict60 61    def verify(self, model_output: str, expected_intent: str = None) -> VerificationResult:62        """63        Verify model output.64        65        Args:66            model_output: Raw text from the model67            expected_intent: What the router classified (tool_call, text, etc.)68            69        Returns:70            VerificationResult with parsed tool calls or cleaned text71        """72        if not model_output or not model_output.strip():73            return VerificationResult(74                valid=False,75                output_type="error",76                error="Empty model output",77                raw_output=model_output,78            )79 80        raw = model_output.strip()81 82        # Step 1: Try to extract tool calls83        tool_calls = extract_tool_calls_from_text(raw)84 85        if tool_calls:86            # Validate each tool call87            valid_calls = []88            errors = []89 90            for tc in tool_calls:91                is_valid, err = validate_tool_call(tc["name"], tc["arguments"])92                if is_valid:93                    valid_calls.append(tc)94                else:95                    errors.append(err)96 97            if valid_calls:98                return VerificationResult(99                    valid=True,100                    output_type="tool_call",101                    tool_calls=valid_calls,102                    raw_output=raw,103                )104            elif errors:105                return VerificationResult(106                    valid=False,107                    output_type="error",108                    error="; ".join(errors),109                    raw_output=raw,110                )111 112        # Step 2: If expected tool_call but got text, that might be an error113        if expected_intent == "tool_call":114            # Check if the model is TRYING to make a tool call but failing115            if any(name in raw.lower() for name in VALID_TOOL_NAMES):116                return VerificationResult(117                    valid=False,118                    output_type="error",119                    error="Model mentioned a tool name but didn't format it as a valid JSON tool call",120                    raw_output=raw,121                )122 123        # Step 3: It's a text response -- clean it up124        cleaned = self._clean_text(raw)125        return VerificationResult(126            valid=True,127            output_type="text",128            text=cleaned,129            raw_output=raw,130        )131 132    def _clean_text(self, text: str) -> str:133        """Clean up verbose model output."""134        # Remove thinking tags (Qwen3 /think blocks)135        text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)136        # Remove empty lines at start/end137        text = text.strip()138        # Remove excessive newlines139        text = re.sub(r'\n{3,}', '\n\n', text)140        return text141 142 143class RetryStrategy:144    """145    Manages retry logic when the verifier rejects model output.146    Each retry gets a more explicit prompt to guide the model.147    """148 149    def __init__(self, max_retries: int = 2):150        self.max_retries = max_retries151 152    def should_retry(self, attempt: int, result: VerificationResult) -> bool:153        """Should we retry after a failed verification?"""154        if result.valid:155            return False156        if attempt >= self.max_retries:157            return False158        # Don't retry empty outputs (model is broken)159        if not result.raw_output.strip():160            return False161        return True162 163    def get_retry_message(self, attempt: int, result: VerificationResult) -> str:164        """Generate a retry prompt based on what went wrong."""165        if attempt == 0:166            # First retry: gentle correction167            return (168                f"Your previous response was not valid: {result.error}\n"169                f"Please respond with ONLY a JSON object like this:\n"170                f'{{"name": "tool_name", "arguments": {{"param": "value"}}}}\n'171                f"No explanation, no markdown, just the JSON."172            )173        else:174            # Second retry: very explicit175            return (176                f"ERROR: {result.error}\n"177                f"You MUST respond with valid JSON only. "178                f"Valid tool names: {', '.join(sorted(VALID_TOOL_NAMES))}.\n"179                f"Format: {{\"name\": \"execute_local\", \"arguments\": {{\"command\": \"your command\"}}}}\n"180                f"Respond with NOTHING except the JSON object."181            )182 183 184def verify_and_extract(model_output: str, expected_intent: str = None) -> VerificationResult:185    """Convenience function: verify model output with default settings."""186    return Verifier().verify(model_output, expected_intent)187