MrMoz33/tokioai-coder-iot
0
1"""2TokioAI Engine -- Test Suite3Tests router, verifier, memory, and full pipeline.4"""5 6import json7import time8import sys9 10from .router import classify_intent, build_system_prompt, IntentType11from .verifier import Verifier, verify_and_extract12from .memory import EngineMemory13from .schemas import validate_tool_call, extract_tool_calls_from_text14from .guided import parse_guided_output15 16 17def run_tests(verbose=False, ollama_host=None, model=None):18 """Run all tests."""19 passed = 020 failed = 021 22 def check(name, condition, detail=""):23 nonlocal passed, failed24 if condition:25 passed += 126 print(f" \033[32mPASS\033[0m {name}")27 else:28 failed += 129 print(f" \033[31mFAIL\033[0m {name} {detail}")30 31 # ========================32 print("\n=== ROUTER TESTS ===")33 # ========================34 35 intent, conf = classify_intent("check disk space")36 check("tool: check disk space", intent == IntentType.TOOL_CALL, f"got {intent}")37 38 intent, conf = classify_intent("list docker containers")39 check("tool: docker containers", intent == IntentType.TOOL_CALL, f"got {intent}")40 41 intent, conf = classify_intent("read /etc/hostname")42 check("tool: read file", intent == IntentType.TOOL_CALL, f"got {intent}")43 44 intent, conf = classify_intent("run nmap on 192.168.1.1")45 check("tool: nmap scan", intent == IntentType.TOOL_CALL, f"got {intent}")46 47 intent, conf = classify_intent("what is SQL injection")48 check("text: SQL injection", intent in (IntentType.TEXT, IntentType.CODE), f"got {intent}")49 50 intent, conf = classify_intent("hi")51 check("text: greeting", intent == IntentType.TEXT, f"got {intent}")52 53 intent, conf = classify_intent("thanks")54 check("text: thanks", intent == IntentType.TEXT, f"got {intent}")55 56 intent, conf = classify_intent("explain the TCP handshake")57 check("text: TCP explanation", intent in (IntentType.TEXT, IntentType.CODE), f"got {intent}")58 59 intent, conf = classify_intent("what is docker")60 check("decision: what is docker", intent in (IntentType.DECISION, IntentType.TEXT, IntentType.TOOL_CALL), f"got {intent}")61 62 # ========================63 print("\n=== VERIFIER TESTS ===")64 # ========================65 66 v = Verifier()67 68 # Valid tool call69 r = v.verify('{"name": "execute_local", "arguments": {"command": "df -h"}}')70 check("verify: valid tool call", r.valid and r.output_type == "tool_call")71 72 # Tool call in markdown73 r = v.verify('```json\n{"name": "read_file", "arguments": {"path": "/etc/hostname"}}\n```')74 check("verify: markdown tool call", r.valid and r.output_type == "tool_call")75 76 # Invalid tool name77 r = v.verify('{"name": "invalid_tool", "arguments": {"x": 1}}')78 check("verify: invalid tool name", r.output_type in ("text", "error"))79 80 # Missing required param81 r = v.verify('{"name": "execute_local", "arguments": {}}')82 check("verify: missing required param", not r.valid)83 84 # Pure text85 r = v.verify("Docker is a containerization platform.")86 check("verify: pure text", r.valid and r.output_type == "text")87 88 # Empty89 r = v.verify("")90 check("verify: empty output", not r.valid)91 92 # ========================93 print("\n=== SCHEMA TESTS ===")94 # ========================95 96 ok, err = validate_tool_call("execute_local", {"command": "ls"})97 check("schema: valid execute_local", ok, err)98 99 ok, err = validate_tool_call("execute_local", {})100 check("schema: missing command", not ok)101 102 ok, err = validate_tool_call("read_file", {"path": "/tmp/test"})103 check("schema: valid read_file", ok, err)104 105 ok, err = validate_tool_call("diagnose", {"target": "local"})106 check("schema: valid diagnose", ok, err)107 108 ok, err = validate_tool_call("diagnose", {"target": "invalid"})109 check("schema: invalid enum", not ok)110 111 # Extract tool calls from messy text112 calls = extract_tool_calls_from_text(113 'I will check the disk. {"name": "execute_local", "arguments": {"command": "df -h"}} Done.'114 )115 check("extract: tool in text", len(calls) == 1 and calls[0]["name"] == "execute_local")116 117 # ========================118 print("\n=== GUIDED TESTS ===")119 # ========================120 121 parsed = parse_guided_output('{"name": "execute_local", "arguments": {"command": "ls"}}', "tool_call")122 check("guided: tool call parse", parsed["type"] == "tool_call" and parsed["name"] == "execute_local")123 124 parsed = parse_guided_output('{"action": "text_response", "text_response": "hello"}', "decision")125 check("guided: decision text", parsed["type"] == "text" and parsed["response"] == "hello")126 127 parsed = parse_guided_output(128 '{"action": "tool_call", "tool_call": {"name": "execute_local", "arguments": {"command": "ls"}}}',129 "decision"130 )131 check("guided: decision tool", parsed["type"] == "tool_call" and parsed["name"] == "execute_local")132 133 # ========================134 print("\n=== MEMORY TESTS ===")135 # ========================136 137 import tempfile, os138 tmp = os.path.join(tempfile.mkdtemp(), "test_memory.json")139 mem = EngineMemory(path=tmp)140 141 check("memory: seeded", len(mem.examples) > 0)142 143 mem.add("check cpu usage", '{"name": "execute_local", "arguments": {"command": "top -bn1"}}',144 intent="tool_call", tool_name="execute_local")145 check("memory: add", len(mem.examples) > 0)146 147 results = mem.retrieve("show cpu", intent="tool_call", top_k=3)148 check("memory: retrieve", len(results) > 0, f"got {len(results)}")149 150 stats = mem.get_stats()151 check("memory: stats", stats["total"] > 0)152 153 # Cleanup154 os.unlink(tmp)155 156 # ========================157 print(f"\n=== RESULTS: {passed} passed, {failed} failed ===\n")158 159 # ========================160 # LIVE TESTS (optional -- only if model is available)161 # ========================162 if ollama_host or model:163 from .pipeline import Pipeline, EngineConfig164 config = EngineConfig(165 ollama_host=ollama_host,166 model_name=model or "tokioai-coder",167 verbose=verbose,168 )169 170 print("=== LIVE PIPELINE TESTS ===")171 print(f"Model: {config.model_name} @ {config.ollama_host}\n")172 173 engine = Pipeline(config)174 live_passed = 0175 live_failed = 0176 177 test_cases = [178 ("check disk space", "tool_call", "execute_local"),179 ("read /etc/hostname", "tool_call", "read_file"),180 ("find TODO in python files", "tool_call", "search_files"),181 ("what is SQL injection", "text", None),182 ("hi", "text", None),183 ("list running processes", "tool_call", "execute_local"),184 ("check system health", "tool_call", "diagnose"),185 ("explain what docker is", "text", None),186 ]187 188 for prompt, expected_type, expected_tool in test_cases:189 try:190 result = engine.process(prompt)191 192 type_match = result.output_type == expected_type193 tool_match = True194 if expected_tool and result.tool_calls:195 tool_match = result.tool_calls[0]["name"] == expected_tool196 197 ok = type_match and tool_match198 detail = ""199 if result.output_type == "tool_call" and result.tool_calls:200 detail = f"-> {result.tool_calls[0]['name']}({list(result.tool_calls[0]['arguments'].keys())})"201 else:202 detail = f"-> text({len(result.text)} chars)"203 204 if ok:205 live_passed += 1206 print(f" \033[32mPASS\033[0m [{result.latency_ms:4d}ms] [{result.attempts}att] {prompt:40s} {detail}")207 else:208 live_failed += 1209 print(f" \033[31mFAIL\033[0m [{result.latency_ms:4d}ms] [{result.attempts}att] {prompt:40s} {detail} (expected {expected_type}/{expected_tool})")210 except Exception as e:211 live_failed += 1212 print(f" \033[31mERROR\033[0m {prompt}: {e}")213 214 engine.reset()215 216 print(f"\n=== LIVE RESULTS: {live_passed} passed, {live_failed} failed ===\n")217 passed += live_passed218 failed += live_failed219 220 return passed, failed221