""" 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