File size: 8,357 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 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 | """
TokioAI Engine -- Test Suite
Tests router, verifier, memory, and full pipeline.
"""
import json
import time
import sys
from .router import classify_intent, build_system_prompt, IntentType
from .verifier import Verifier, verify_and_extract
from .memory import EngineMemory
from .schemas import validate_tool_call, extract_tool_calls_from_text
from .guided import parse_guided_output
def run_tests(verbose=False, ollama_host=None, model=None):
"""Run all tests."""
passed = 0
failed = 0
def check(name, condition, detail=""):
nonlocal passed, failed
if condition:
passed += 1
print(f" \033[32mPASS\033[0m {name}")
else:
failed += 1
print(f" \033[31mFAIL\033[0m {name} {detail}")
# ========================
print("\n=== ROUTER TESTS ===")
# ========================
intent, conf = classify_intent("check disk space")
check("tool: check disk space", intent == IntentType.TOOL_CALL, f"got {intent}")
intent, conf = classify_intent("list docker containers")
check("tool: docker containers", intent == IntentType.TOOL_CALL, f"got {intent}")
intent, conf = classify_intent("read /etc/hostname")
check("tool: read file", intent == IntentType.TOOL_CALL, f"got {intent}")
intent, conf = classify_intent("run nmap on 192.168.1.1")
check("tool: nmap scan", intent == IntentType.TOOL_CALL, f"got {intent}")
intent, conf = classify_intent("what is SQL injection")
check("text: SQL injection", intent in (IntentType.TEXT, IntentType.CODE), f"got {intent}")
intent, conf = classify_intent("hi")
check("text: greeting", intent == IntentType.TEXT, f"got {intent}")
intent, conf = classify_intent("thanks")
check("text: thanks", intent == IntentType.TEXT, f"got {intent}")
intent, conf = classify_intent("explain the TCP handshake")
check("text: TCP explanation", intent in (IntentType.TEXT, IntentType.CODE), f"got {intent}")
intent, conf = classify_intent("what is docker")
check("decision: what is docker", intent in (IntentType.DECISION, IntentType.TEXT, IntentType.TOOL_CALL), f"got {intent}")
# ========================
print("\n=== VERIFIER TESTS ===")
# ========================
v = Verifier()
# Valid tool call
r = v.verify('{"name": "execute_local", "arguments": {"command": "df -h"}}')
check("verify: valid tool call", r.valid and r.output_type == "tool_call")
# Tool call in markdown
r = v.verify('```json\n{"name": "read_file", "arguments": {"path": "/etc/hostname"}}\n```')
check("verify: markdown tool call", r.valid and r.output_type == "tool_call")
# Invalid tool name
r = v.verify('{"name": "invalid_tool", "arguments": {"x": 1}}')
check("verify: invalid tool name", r.output_type in ("text", "error"))
# Missing required param
r = v.verify('{"name": "execute_local", "arguments": {}}')
check("verify: missing required param", not r.valid)
# Pure text
r = v.verify("Docker is a containerization platform.")
check("verify: pure text", r.valid and r.output_type == "text")
# Empty
r = v.verify("")
check("verify: empty output", not r.valid)
# ========================
print("\n=== SCHEMA TESTS ===")
# ========================
ok, err = validate_tool_call("execute_local", {"command": "ls"})
check("schema: valid execute_local", ok, err)
ok, err = validate_tool_call("execute_local", {})
check("schema: missing command", not ok)
ok, err = validate_tool_call("read_file", {"path": "/tmp/test"})
check("schema: valid read_file", ok, err)
ok, err = validate_tool_call("diagnose", {"target": "local"})
check("schema: valid diagnose", ok, err)
ok, err = validate_tool_call("diagnose", {"target": "invalid"})
check("schema: invalid enum", not ok)
# Extract tool calls from messy text
calls = extract_tool_calls_from_text(
'I will check the disk. {"name": "execute_local", "arguments": {"command": "df -h"}} Done.'
)
check("extract: tool in text", len(calls) == 1 and calls[0]["name"] == "execute_local")
# ========================
print("\n=== GUIDED TESTS ===")
# ========================
parsed = parse_guided_output('{"name": "execute_local", "arguments": {"command": "ls"}}', "tool_call")
check("guided: tool call parse", parsed["type"] == "tool_call" and parsed["name"] == "execute_local")
parsed = parse_guided_output('{"action": "text_response", "text_response": "hello"}', "decision")
check("guided: decision text", parsed["type"] == "text" and parsed["response"] == "hello")
parsed = parse_guided_output(
'{"action": "tool_call", "tool_call": {"name": "execute_local", "arguments": {"command": "ls"}}}',
"decision"
)
check("guided: decision tool", parsed["type"] == "tool_call" and parsed["name"] == "execute_local")
# ========================
print("\n=== MEMORY TESTS ===")
# ========================
import tempfile, os
tmp = os.path.join(tempfile.mkdtemp(), "test_memory.json")
mem = EngineMemory(path=tmp)
check("memory: seeded", len(mem.examples) > 0)
mem.add("check cpu usage", '{"name": "execute_local", "arguments": {"command": "top -bn1"}}',
intent="tool_call", tool_name="execute_local")
check("memory: add", len(mem.examples) > 0)
results = mem.retrieve("show cpu", intent="tool_call", top_k=3)
check("memory: retrieve", len(results) > 0, f"got {len(results)}")
stats = mem.get_stats()
check("memory: stats", stats["total"] > 0)
# Cleanup
os.unlink(tmp)
# ========================
print(f"\n=== RESULTS: {passed} passed, {failed} failed ===\n")
# ========================
# LIVE TESTS (optional -- only if model is available)
# ========================
if ollama_host or model:
from .pipeline import Pipeline, EngineConfig
config = EngineConfig(
ollama_host=ollama_host,
model_name=model or "tokioai-coder",
verbose=verbose,
)
print("=== LIVE PIPELINE TESTS ===")
print(f"Model: {config.model_name} @ {config.ollama_host}\n")
engine = Pipeline(config)
live_passed = 0
live_failed = 0
test_cases = [
("check disk space", "tool_call", "execute_local"),
("read /etc/hostname", "tool_call", "read_file"),
("find TODO in python files", "tool_call", "search_files"),
("what is SQL injection", "text", None),
("hi", "text", None),
("list running processes", "tool_call", "execute_local"),
("check system health", "tool_call", "diagnose"),
("explain what docker is", "text", None),
]
for prompt, expected_type, expected_tool in test_cases:
try:
result = engine.process(prompt)
type_match = result.output_type == expected_type
tool_match = True
if expected_tool and result.tool_calls:
tool_match = result.tool_calls[0]["name"] == expected_tool
ok = type_match and tool_match
detail = ""
if result.output_type == "tool_call" and result.tool_calls:
detail = f"-> {result.tool_calls[0]['name']}({list(result.tool_calls[0]['arguments'].keys())})"
else:
detail = f"-> text({len(result.text)} chars)"
if ok:
live_passed += 1
print(f" \033[32mPASS\033[0m [{result.latency_ms:4d}ms] [{result.attempts}att] {prompt:40s} {detail}")
else:
live_failed += 1
print(f" \033[31mFAIL\033[0m [{result.latency_ms:4d}ms] [{result.attempts}att] {prompt:40s} {detail} (expected {expected_type}/{expected_tool})")
except Exception as e:
live_failed += 1
print(f" \033[31mERROR\033[0m {prompt}: {e}")
engine.reset()
print(f"\n=== LIVE RESULTS: {live_passed} passed, {live_failed} failed ===\n")
passed += live_passed
failed += live_failed
return passed, failed
|