| """ |
| 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() |
|
|
| |
| r = v.verify('{"name": "execute_local", "arguments": {"command": "df -h"}}') |
| check("verify: valid tool call", r.valid and r.output_type == "tool_call") |
|
|
| |
| 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") |
|
|
| |
| r = v.verify('{"name": "invalid_tool", "arguments": {"x": 1}}') |
| check("verify: invalid tool name", r.output_type in ("text", "error")) |
|
|
| |
| r = v.verify('{"name": "execute_local", "arguments": {}}') |
| check("verify: missing required param", not r.valid) |
|
|
| |
| r = v.verify("Docker is a containerization platform.") |
| check("verify: pure text", r.valid and r.output_type == "text") |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| os.unlink(tmp) |
|
|
| |
| print(f"\n=== RESULTS: {passed} passed, {failed} failed ===\n") |
|
|
| |
| |
| |
| 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 |
|
|
