| """ |
| TokioAI Engine -- CLI Entry Point |
| |
| Usage: |
| python -m engine serve # Start OpenAI-compatible API server |
| python -m engine test # Run test suite |
| python -m engine bench # Benchmark pipeline |
| python -m engine chat # Interactive chat (for testing) |
| python -m engine stats # Show memory/pipeline stats |
| """ |
|
|
| import sys |
| import json |
| import time |
| import argparse |
|
|
| from .pipeline import Pipeline, EngineConfig |
|
|
|
|
| def cmd_serve(args): |
| """Start the API server.""" |
| from .server import create_app |
| from .pipeline import EngineConfig |
|
|
| config = EngineConfig( |
| ollama_host=args.ollama_host, |
| model_name=args.model, |
| max_retries=args.max_retries, |
| use_guided_gen=not args.no_guided, |
| use_memory=not args.no_memory, |
| verbose=args.verbose, |
| ) |
|
|
| print(f""" |
| _____ _ _ _ ___ |
| |_ _|__ | | _(_) ___ / \\ |_ _| |
| | |/ _ \\| |/ / |/ _ / _ \\ | | |
| | | (_) | <| | (_/ ___ \\ | | |
| |_|\\___/|_|\\_\\_|\\___/_/ \\_\\___| |
| ___ _ |
| | __| _ _ __ _(_) _ _ ___ |
| | _| | ' \\/ _` | || ' \\/ -_) |
| |___||_||_\\__, |_||_||_\\___| |
| |___/ v3.0.0 |
| |
| Model: {config.model_name} |
| Ollama: {config.ollama_host} |
| Memory: {'ON' if config.use_memory else 'OFF'} |
| Guided: {'ON' if config.use_guided_gen else 'OFF'} |
| Retries: {config.max_retries} |
| |
| API: http://{args.host}:{args.port}/v1/chat/completions |
| Health: http://{args.host}:{args.port}/health |
| Stats: http://{args.host}:{args.port}/stats |
| """) |
|
|
| app = create_app(config) |
| app.run(host=args.host, port=args.port, debug=False) |
|
|
|
|
| def cmd_chat(args): |
| """Interactive chat for testing the pipeline.""" |
| config = EngineConfig( |
| ollama_host=args.ollama_host, |
| model_name=args.model, |
| verbose=args.verbose, |
| ) |
| engine = Pipeline(config) |
|
|
| print(f"\n[TokioAI Engine] Connected to {config.ollama_host} / {config.model_name}") |
| print("[TokioAI Engine] Type 'quit' to exit, 'stats' for stats, 'reset' to clear history\n") |
|
|
| while True: |
| try: |
| user_input = input("\033[36myou>\033[0m ") |
| except (EOFError, KeyboardInterrupt): |
| print("\nbye.") |
| break |
|
|
| if not user_input.strip(): |
| continue |
| if user_input.strip().lower() == "quit": |
| break |
| if user_input.strip().lower() == "stats": |
| print(json.dumps(engine.get_stats(), indent=2)) |
| continue |
| if user_input.strip().lower() == "reset": |
| engine.reset() |
| print("[reset]") |
| continue |
|
|
| result = engine.process(user_input) |
|
|
| |
| print(f"\033[33m[{result.intent}|{result.confidence:.1f}|{result.attempts}att|{result.latency_ms}ms]\033[0m") |
|
|
| if result.output_type == "tool_call": |
| for tc in result.tool_calls: |
| print(f"\033[32m>>> TOOL: {tc['name']}({json.dumps(tc['arguments'])})\033[0m") |
| else: |
| print(f"\033[37m{result.text}\033[0m") |
| print() |
|
|
|
|
| def cmd_test(args): |
| """Run the test suite.""" |
| from .test_engine import run_tests |
| run_tests(verbose=args.verbose, ollama_host=args.ollama_host, model=args.model) |
|
|
|
|
| def cmd_bench(args): |
| """Benchmark the pipeline.""" |
| config = EngineConfig( |
| ollama_host=args.ollama_host, |
| model_name=args.model, |
| verbose=False, |
| ) |
| engine = Pipeline(config) |
|
|
| test_prompts = [ |
| ("check disk space", "tool_call"), |
| ("what is docker", "text"), |
| ("read /etc/hostname", "tool_call"), |
| ("hi", "text"), |
| ("find all python files with TODO", "tool_call"), |
| ("explain SQL injection", "text"), |
| ("is docker running", "tool_call"), |
| ("list processes by memory usage", "tool_call"), |
| ("what is the TCP 3-way handshake", "text"), |
| ("scan ports on localhost", "tool_call"), |
| ] |
|
|
| print(f"\n[BENCH] Running {len(test_prompts)} prompts against {config.model_name}...\n") |
| |
| results = [] |
| for prompt, expected in test_prompts: |
| start = time.time() |
| result = engine.process(prompt) |
| elapsed = int((time.time() - start) * 1000) |
|
|
| correct = ( |
| (expected == "tool_call" and result.output_type == "tool_call") or |
| (expected == "text" and result.output_type == "text") |
| ) |
|
|
| status = "\033[32mPASS\033[0m" if correct else "\033[31mFAIL\033[0m" |
| detail = "" |
| if result.output_type == "tool_call" and result.tool_calls: |
| detail = f"-> {result.tool_calls[0]['name']}" |
| elif result.output_type == "text": |
| detail = f"-> text ({len(result.text)} chars)" |
|
|
| print(f" {status} [{elapsed:4d}ms] [{result.intent:9s}] {prompt:40s} {detail}") |
| results.append(correct) |
| engine.reset() |
|
|
| passed = sum(results) |
| total = len(results) |
| print(f"\n[BENCH] {passed}/{total} correct ({100*passed//total}%)") |
| print(f"[BENCH] Stats: {json.dumps(engine.get_stats(), indent=2)}") |
|
|
|
|
| def cmd_stats(args): |
| """Show memory and pipeline stats.""" |
| from .memory import EngineMemory |
| mem = EngineMemory() |
| print(json.dumps(mem.get_stats(), indent=2)) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="TokioAI Engine") |
| parser.add_argument("--ollama-host", default=None, help="Ollama API URL") |
| parser.add_argument("--model", default="tokioai-coder", help="Model name") |
| parser.add_argument("-v", "--verbose", action="store_true") |
| |
| subparsers = parser.add_subparsers(dest="command") |
| |
| |
| serve_parser = subparsers.add_parser("serve", help="Start API server") |
| serve_parser.add_argument("--host", default="0.0.0.0", help="Bind address") |
| serve_parser.add_argument("--port", type=int, default=8080, help="Port") |
| serve_parser.add_argument("--ollama-host", default=None, help="Ollama API URL") |
| serve_parser.add_argument("--model", default="tokioai-coder", help="Model name") |
| serve_parser.add_argument("--no-memory", action="store_true", help="Disable memory") |
| serve_parser.add_argument("--no-guided", action="store_true", help="Disable guided gen") |
| serve_parser.add_argument("--max-retries", type=int, default=2, help="Max retries") |
| serve_parser.add_argument("-v", "--verbose", action="store_true") |
| |
| |
| chat_parser = subparsers.add_parser("chat", help="Interactive chat") |
| |
| |
| test_parser = subparsers.add_parser("test", help="Run tests") |
| |
| |
| bench_parser = subparsers.add_parser("bench", help="Benchmark") |
| |
| |
| stats_parser = subparsers.add_parser("stats", help="Show stats") |
|
|
| |
| for sp in [chat_parser, test_parser, bench_parser]: |
| sp.add_argument("--ollama-host", default=None, help="Ollama API URL") |
| sp.add_argument("--model", default="tokioai-coder", help="Model name") |
| sp.add_argument("-v", "--verbose", action="store_true") |
|
|
| args = parser.parse_args() |
|
|
| if args.command == "serve": |
| cmd_serve(args) |
| elif args.command == "chat": |
| cmd_chat(args) |
| elif args.command == "test": |
| cmd_test(args) |
| elif args.command == "bench": |
| cmd_bench(args) |
| elif args.command == "stats": |
| cmd_stats(args) |
| else: |
| parser.print_help() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
