File size: 6,403 Bytes
4b97c3b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | """
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)
|