MrMoz33 commited on
Commit
4b97c3b
·
verified ·
1 Parent(s): 24b4c34

TokioAI-Coder IoT Package v1.0 -- Engine + IoT System Prompt + Tools

Browse files
LICENSE ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ MIT License
2
+
3
+ Copyright (c) 2025 TokioAI
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
MANIFEST.in ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ include LICENSE
2
+ include README.md
3
+ include MANIFEST.in
4
+ recursive-include examples *.py
5
+ recursive-include engine *.py
README.md ADDED
@@ -0,0 +1,242 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # TokioAI-Coder IoT Package
2
+
3
+ **Turn a 30B open-weight model into a fully autonomous IoT/smart-home agent.**
4
+
5
+ This package contains everything needed to deploy TokioAI-Coder as an autonomous IoT controller:
6
+ - The complete **TokioAI Engine** -- inference pipeline that wraps any Ollama model and adds reliable tool calling
7
+ - A production-tested **IoT system prompt** with a full Home Assistant cookbook
8
+ - **12 native tools** for system administration, file operations, network diagnostics, and robot vision
9
+ - **Server-side investigation** -- the engine autonomously chains 8+ tool calls to answer a single question
10
+
11
+ ## What It Does
12
+
13
+ You say: *"is Home Assistant running?"*
14
+
15
+ The engine:
16
+ 1. Checks `systemctl status home-assistant` -- not found
17
+ 2. Checks `docker ps | grep home` -- not found
18
+ 3. Checks `ss -tlnp | grep 8123` -- port listening
19
+ 4. Checks `curl -s http://localhost:8123/api/` -- HA API responds
20
+ 5. Answers: "Yes, HA is running on port 8123 via Docker"
21
+
22
+ All 5 steps happen **in a single API call**. The client sees only the final answer.
23
+
24
+ You say: *"play some blues"*
25
+
26
+ The engine:
27
+ 1. Calls HA API to send `play_media` to the Alexa
28
+ 2. Waits 3 seconds
29
+ 3. Queries the media player state to verify the music actually changed
30
+ 4. Reports: "Now playing: Blues Essentials on Jarvis"
31
+
32
+ ## Architecture
33
+
34
+ ```
35
+ TokioAI Engine (OpenAI-compatible API)
36
+ ====================================
37
+ User Query ──> Router ──> Ollama Model ──> Verifier ──> Tool Executor ──>
38
+ | | | |
39
+ Intent System Retry Local
40
+ Classifier Prompt + Loop (2x) Execution
41
+ (0ms regex) Few-Shot (read-only safe)
42
+ | | | |
43
+ v v v v
44
+ tool_call Guided Validates Feeds result
45
+ / text / Generation JSON + back to model
46
+ decision (schema- correct for next step
47
+ constrained) tool names
48
+ |
49
+ v
50
+ Investigation Loop
51
+ (up to 8 rounds,
52
+ 90s timeout)
53
+ |
54
+ v
55
+ Final Answer
56
+ ```
57
+
58
+ ## Components
59
+
60
+ ### 1. TokioAI Engine (`engine/`)
61
+ The core inference pipeline (3,774 lines of Python):
62
+
63
+ | Module | Lines | Purpose |
64
+ |---|---|---|
65
+ | `server.py` | 864 | OpenAI-compatible API with server-side investigation |
66
+ | `schemas.py` | 497 | Tool definitions, JSON validation, parameter schemas |
67
+ | `pipeline.py` | 476 | Main orchestrator: Router -> Memory -> Model -> Verifier |
68
+ | `router.py` | 382 | Intent classifier (regex) + system prompt builder |
69
+ | `memory.py` | 353 | Adaptive few-shot memory (learns from successes) |
70
+ | `executor.py` | 331 | Safe server-side tool execution |
71
+ | `guided.py` | 228 | JSON schema constraints for Ollama |
72
+ | `verifier.py` | 186 | Output validation + retry logic |
73
+
74
+ ### 2. IoT System Prompt (`prompts/iot_system_prompt.md`)
75
+ A production-tested system prompt that includes:
76
+ - Complete Home Assistant API cookbook (media, lights, switches, sensors, vacuum, TTS)
77
+ - Alexa voice control via `alexa_media` integration
78
+ - Robot vision via moondream (1.4B vision model)
79
+ - Investigation playbooks for multi-step problem solving
80
+ - Anti-hallucination rules (never simulate tool execution)
81
+ - Verification protocol (always confirm state changes actually happened)
82
+
83
+ ### 3. Tool Definitions (`tools/`)
84
+ 12 native tools in OpenAI function-calling format:
85
+ - `execute_local` -- Run shell commands
86
+ - `execute_raspi` / `execute_gcp` / `execute_router` -- Remote execution
87
+ - `read_file` / `write_file` / `edit_file` -- File operations
88
+ - `search_files` -- grep across codebases
89
+ - `diagnose` -- System health checks
90
+ - `robot_vision` -- Camera capture + AI analysis
91
+ - `memory` -- Persistent memory across sessions
92
+ - `task` -- Task tracking for multi-session projects
93
+
94
+ ## Quick Start
95
+
96
+ ### Prerequisites
97
+ - Python 3.8+
98
+ - [Ollama](https://ollama.ai/) with a model (recommended: `qwen3-coder` 30B MoE)
99
+ - Home Assistant instance (for IoT features)
100
+
101
+ ### Install
102
+
103
+ ```bash
104
+ # Clone this repo
105
+ git clone https://huggingface.co/MrMoz33/tokioai-coder-iot
106
+ cd tokioai-coder-iot
107
+
108
+ # Install dependencies
109
+ pip install flask requests
110
+
111
+ # Pull a model
112
+ ollama pull qwen3-coder # 30B MoE (3B active), recommended
113
+ # or: ollama pull qwen2.5-coder:7b # lighter alternative
114
+ ```
115
+
116
+ ### Configure Home Assistant
117
+
118
+ Set your HA token:
119
+ ```bash
120
+ export HA_TOKEN="your-long-lived-access-token"
121
+ export HA_URL="http://localhost:8123" # or your HA URL
122
+ ```
123
+
124
+ ### Start
125
+
126
+ ```bash
127
+ # Start the engine
128
+ python -m engine serve --model qwen3-coder --port 8080
129
+
130
+ # Use it (any OpenAI-compatible client works)
131
+ curl http://localhost:8080/v1/chat/completions \
132
+ -H "Content-Type: application/json" \
133
+ -d '{
134
+ "model": "tokioai-engine",
135
+ "messages": [{"role": "user", "content": "is home assistant running?"}]
136
+ }'
137
+ ```
138
+
139
+ ### With TokioAI CLI
140
+
141
+ ```bash
142
+ pip install git+https://github.com/TokioAI/tokioai.git
143
+ export TOKIOAI_ENGINE_HOST=http://localhost:8080
144
+ tokioai # Interactive agent with full tool calling
145
+ ```
146
+
147
+ ## Customization
148
+
149
+ ### Adding Your Own Home Assistant Entities
150
+
151
+ Edit `prompts/iot_system_prompt.md` and add your entities to the cookbook section:
152
+
153
+ ```markdown
154
+ LIGHTS: light.living_room, light.bedroom, light.kitchen
155
+ SWITCHES: switch.fan, switch.heater
156
+ SENSORS: sensor.temperature, sensor.humidity
157
+ MEDIA: media_player.living_room_speaker
158
+ ```
159
+
160
+ ### Adding Custom Tools
161
+
162
+ Edit `engine/schemas.py` to add new tool definitions:
163
+
164
+ ```python
165
+ {
166
+ "type": "function",
167
+ "function": {
168
+ "name": "my_custom_tool",
169
+ "description": "What this tool does",
170
+ "parameters": {
171
+ "type": "object",
172
+ "properties": {
173
+ "param1": {"type": "string", "description": "First param"}
174
+ },
175
+ "required": ["param1"]
176
+ }
177
+ }
178
+ }
179
+ ```
180
+
181
+ Then add the executor in `engine/executor.py`.
182
+
183
+ ### Adapting for Different Smart Home Platforms
184
+
185
+ The system prompt is platform-agnostic in concept. To use with:
186
+ - **Google Home**: Replace HA API calls with Google Home API
187
+ - **Apple HomeKit**: Use HomeKit Accessory Protocol
188
+ - **MQTT Direct**: Replace curl commands with `mosquitto_pub/sub`
189
+
190
+ The engine and tool pipeline remain identical -- only the system prompt changes.
191
+
192
+ ## Performance
193
+
194
+ Tested on Raspberry Pi 5 (8GB) + L4 GPU (via SSH tunnel):
195
+
196
+ | Metric | Value |
197
+ |---|---|
198
+ | Model | Qwen3-Coder 30B (3B active MoE) |
199
+ | GPU VRAM | 19.4 GB (L4 24GB) |
200
+ | Response time | 3-8s (simple), 15-45s (investigation) |
201
+ | Investigation depth | Up to 8 rounds per query |
202
+ | Tool calling accuracy | 94%+ with engine pipeline |
203
+ | Uptime | 30+ days continuous (systemd service) |
204
+
205
+ ### Cost Comparison
206
+
207
+ | Approach | Monthly Cost | Latency | Privacy |
208
+ |---|---|---|---|
209
+ | GPT-4o API | $50-200 | 2-5s | Data sent to cloud |
210
+ | Claude API | $50-200 | 2-5s | Data sent to cloud |
211
+ | **TokioAI-Coder (self-hosted)** | **$0** | **3-8s** | **100% local** |
212
+
213
+ After initial GPU investment ($0.50-1.50/hr for L4 on GCP, or free on local hardware), all inference is free and private.
214
+
215
+ ## Model Compatibility
216
+
217
+ | Model | Params | Tool Calling | IoT Quality | Notes |
218
+ |---|---|---|---|---|
219
+ | qwen3-coder | 30B (3B active) | Excellent | Excellent | Recommended. MoE = fast inference |
220
+ | qwen2.5-coder:7b | 7B | Good | Good | Budget option |
221
+ | llama3.1:8b | 8B | Good | Fair | General purpose |
222
+ | deepseek-coder-v2:16b | 16B | Good | Good | Strong code understanding |
223
+ | mistral:7b | 7B | OK | Fair | Needs guided gen |
224
+
225
+ The engine's guided generation + retry logic compensates for weaker models.
226
+
227
+ ## Related Projects
228
+
229
+ - [TokioAI CLI](https://github.com/TokioAI/tokioai) -- Multi-provider AI CLI with native tool calling
230
+ - [TokioAI Fine-Tuning Guide](https://github.com/daletoniris/tokioai-fine-tuning-guide) -- Train your own tool-calling model
231
+ - [TokioAI Fine-Tuning Paper](https://github.com/daletoniris/tokioai-fine-tuning-paper) -- Research paper on LLM fine-tuning for tool use
232
+ - [TokioNav](https://github.com/daletoniris/tokionav) -- AI-controlled autonomous robot navigation
233
+
234
+ ## License
235
+
236
+ MIT
237
+
238
+ ## Credits
239
+
240
+ Built by Daniel Dieser ([daletoniris](https://github.com/daletoniris)) and the TokioAI team.
241
+
242
+ The TokioAI Engine was developed to make small open models competitive with GPT-4 and Claude for autonomous tool calling -- specifically for IoT, cybersecurity, and DevOps tasks. The key insight: you don't need a bigger model, you need a smarter pipeline around it.
config/example.env ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # TokioAI-Coder IoT Configuration
2
+ # Copy this to ~/.tokioai/.env and fill in your values
3
+
4
+ # ============================================
5
+ # Home Assistant
6
+ # ============================================
7
+ # Get a long-lived access token from:
8
+ # HA -> Profile -> Long-Lived Access Tokens -> Create Token
9
+ HA_TOKEN=your-long-lived-access-token-here
10
+ HA_URL=http://localhost:8123
11
+
12
+ # ============================================
13
+ # Ollama / Model
14
+ # ============================================
15
+ # Where Ollama is running (default: localhost)
16
+ OLLAMA_HOST=http://localhost:11434
17
+ # Which model to use
18
+ TOKIOAI_MODEL=qwen3-coder
19
+
20
+ # ============================================
21
+ # Engine Server
22
+ # ============================================
23
+ ENGINE_PORT=8080
24
+ MAX_INVESTIGATION_DEPTH=8
25
+ MAX_INVESTIGATION_TIME=90
26
+
27
+ # ============================================
28
+ # Robot Vision (optional)
29
+ # ============================================
30
+ # URL of the robot camera proxy (if you have one)
31
+ # ROBOT_PROXY_URL=http://localhost:5002
32
+ # Vision model for image analysis
33
+ # VISION_MODEL=moondream
34
+
35
+ # ============================================
36
+ # Remote Targets (optional)
37
+ # ============================================
38
+ # SSH config for execute_raspi, execute_gcp, etc.
39
+ # RASPI_HOST=192.168.1.x
40
+ # RASPI_USER=pi
41
+ # GCP_HOST=your-gcp-vm-ip
42
+ # GCP_USER=your-user
engine/__init__.py ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ TokioAI Engine v3.0 -- Inference Pipeline with Server-Side Investigation
3
+ Transform any small open model into an elite AI agent.
4
+
5
+ Architecture:
6
+ [User Prompt] -> Router -> Model -> Verifier -> Response
7
+ | | |
8
+ Memory Executor Retry Loop
9
+ (server-side tool execution)
10
+ """
11
+
12
+ __version__ = "3.0.0"
13
+ __author__ = "TokioAI"
engine/__main__.py ADDED
@@ -0,0 +1,224 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ TokioAI Engine -- CLI Entry Point
3
+
4
+ Usage:
5
+ python -m engine serve # Start OpenAI-compatible API server
6
+ python -m engine test # Run test suite
7
+ python -m engine bench # Benchmark pipeline
8
+ python -m engine chat # Interactive chat (for testing)
9
+ python -m engine stats # Show memory/pipeline stats
10
+ """
11
+
12
+ import sys
13
+ import json
14
+ import time
15
+ import argparse
16
+
17
+ from .pipeline import Pipeline, EngineConfig
18
+
19
+
20
+ def cmd_serve(args):
21
+ """Start the API server."""
22
+ from .server import create_app
23
+ from .pipeline import EngineConfig
24
+
25
+ config = EngineConfig(
26
+ ollama_host=args.ollama_host,
27
+ model_name=args.model,
28
+ max_retries=args.max_retries,
29
+ use_guided_gen=not args.no_guided,
30
+ use_memory=not args.no_memory,
31
+ verbose=args.verbose,
32
+ )
33
+
34
+ print(f"""
35
+ _____ _ _ _ ___
36
+ |_ _|__ | | _(_) ___ / \\ |_ _|
37
+ | |/ _ \\| |/ / |/ _ / _ \\ | |
38
+ | | (_) | <| | (_/ ___ \\ | |
39
+ |_|\\___/|_|\\_\\_|\\___/_/ \\_\\___|
40
+ ___ _
41
+ | __| _ _ __ _(_) _ _ ___
42
+ | _| | ' \\/ _` | || ' \\/ -_)
43
+ |___||_||_\\__, |_||_||_\\___|
44
+ |___/ v3.0.0
45
+
46
+ Model: {config.model_name}
47
+ Ollama: {config.ollama_host}
48
+ Memory: {'ON' if config.use_memory else 'OFF'}
49
+ Guided: {'ON' if config.use_guided_gen else 'OFF'}
50
+ Retries: {config.max_retries}
51
+
52
+ API: http://{args.host}:{args.port}/v1/chat/completions
53
+ Health: http://{args.host}:{args.port}/health
54
+ Stats: http://{args.host}:{args.port}/stats
55
+ """)
56
+
57
+ app = create_app(config)
58
+ app.run(host=args.host, port=args.port, debug=False)
59
+
60
+
61
+ def cmd_chat(args):
62
+ """Interactive chat for testing the pipeline."""
63
+ config = EngineConfig(
64
+ ollama_host=args.ollama_host,
65
+ model_name=args.model,
66
+ verbose=args.verbose,
67
+ )
68
+ engine = Pipeline(config)
69
+
70
+ print(f"\n[TokioAI Engine] Connected to {config.ollama_host} / {config.model_name}")
71
+ print("[TokioAI Engine] Type 'quit' to exit, 'stats' for stats, 'reset' to clear history\n")
72
+
73
+ while True:
74
+ try:
75
+ user_input = input("\033[36myou>\033[0m ")
76
+ except (EOFError, KeyboardInterrupt):
77
+ print("\nbye.")
78
+ break
79
+
80
+ if not user_input.strip():
81
+ continue
82
+ if user_input.strip().lower() == "quit":
83
+ break
84
+ if user_input.strip().lower() == "stats":
85
+ print(json.dumps(engine.get_stats(), indent=2))
86
+ continue
87
+ if user_input.strip().lower() == "reset":
88
+ engine.reset()
89
+ print("[reset]")
90
+ continue
91
+
92
+ result = engine.process(user_input)
93
+
94
+ # Display result
95
+ print(f"\033[33m[{result.intent}|{result.confidence:.1f}|{result.attempts}att|{result.latency_ms}ms]\033[0m")
96
+
97
+ if result.output_type == "tool_call":
98
+ for tc in result.tool_calls:
99
+ print(f"\033[32m>>> TOOL: {tc['name']}({json.dumps(tc['arguments'])})\033[0m")
100
+ else:
101
+ print(f"\033[37m{result.text}\033[0m")
102
+ print()
103
+
104
+
105
+ def cmd_test(args):
106
+ """Run the test suite."""
107
+ from .test_engine import run_tests
108
+ run_tests(verbose=args.verbose, ollama_host=args.ollama_host, model=args.model)
109
+
110
+
111
+ def cmd_bench(args):
112
+ """Benchmark the pipeline."""
113
+ config = EngineConfig(
114
+ ollama_host=args.ollama_host,
115
+ model_name=args.model,
116
+ verbose=False,
117
+ )
118
+ engine = Pipeline(config)
119
+
120
+ test_prompts = [
121
+ ("check disk space", "tool_call"),
122
+ ("what is docker", "text"),
123
+ ("read /etc/hostname", "tool_call"),
124
+ ("hi", "text"),
125
+ ("find all python files with TODO", "tool_call"),
126
+ ("explain SQL injection", "text"),
127
+ ("is docker running", "tool_call"),
128
+ ("list processes by memory usage", "tool_call"),
129
+ ("what is the TCP 3-way handshake", "text"),
130
+ ("scan ports on localhost", "tool_call"),
131
+ ]
132
+
133
+ print(f"\n[BENCH] Running {len(test_prompts)} prompts against {config.model_name}...\n")
134
+
135
+ results = []
136
+ for prompt, expected in test_prompts:
137
+ start = time.time()
138
+ result = engine.process(prompt)
139
+ elapsed = int((time.time() - start) * 1000)
140
+
141
+ correct = (
142
+ (expected == "tool_call" and result.output_type == "tool_call") or
143
+ (expected == "text" and result.output_type == "text")
144
+ )
145
+
146
+ status = "\033[32mPASS\033[0m" if correct else "\033[31mFAIL\033[0m"
147
+ detail = ""
148
+ if result.output_type == "tool_call" and result.tool_calls:
149
+ detail = f"-> {result.tool_calls[0]['name']}"
150
+ elif result.output_type == "text":
151
+ detail = f"-> text ({len(result.text)} chars)"
152
+
153
+ print(f" {status} [{elapsed:4d}ms] [{result.intent:9s}] {prompt:40s} {detail}")
154
+ results.append(correct)
155
+ engine.reset() # Reset between tests
156
+
157
+ passed = sum(results)
158
+ total = len(results)
159
+ print(f"\n[BENCH] {passed}/{total} correct ({100*passed//total}%)")
160
+ print(f"[BENCH] Stats: {json.dumps(engine.get_stats(), indent=2)}")
161
+
162
+
163
+ def cmd_stats(args):
164
+ """Show memory and pipeline stats."""
165
+ from .memory import EngineMemory
166
+ mem = EngineMemory()
167
+ print(json.dumps(mem.get_stats(), indent=2))
168
+
169
+
170
+ def main():
171
+ parser = argparse.ArgumentParser(description="TokioAI Engine")
172
+ parser.add_argument("--ollama-host", default=None, help="Ollama API URL")
173
+ parser.add_argument("--model", default="tokioai-coder", help="Model name")
174
+ parser.add_argument("-v", "--verbose", action="store_true")
175
+
176
+ subparsers = parser.add_subparsers(dest="command")
177
+
178
+ # serve
179
+ serve_parser = subparsers.add_parser("serve", help="Start API server")
180
+ serve_parser.add_argument("--host", default="0.0.0.0", help="Bind address")
181
+ serve_parser.add_argument("--port", type=int, default=8080, help="Port")
182
+ serve_parser.add_argument("--ollama-host", default=None, help="Ollama API URL")
183
+ serve_parser.add_argument("--model", default="tokioai-coder", help="Model name")
184
+ serve_parser.add_argument("--no-memory", action="store_true", help="Disable memory")
185
+ serve_parser.add_argument("--no-guided", action="store_true", help="Disable guided gen")
186
+ serve_parser.add_argument("--max-retries", type=int, default=2, help="Max retries")
187
+ serve_parser.add_argument("-v", "--verbose", action="store_true")
188
+
189
+ # chat
190
+ chat_parser = subparsers.add_parser("chat", help="Interactive chat")
191
+
192
+ # test
193
+ test_parser = subparsers.add_parser("test", help="Run tests")
194
+
195
+ # bench
196
+ bench_parser = subparsers.add_parser("bench", help="Benchmark")
197
+
198
+ # stats
199
+ stats_parser = subparsers.add_parser("stats", help="Show stats")
200
+
201
+ # Add common args to all subparsers that need them
202
+ for sp in [chat_parser, test_parser, bench_parser]:
203
+ sp.add_argument("--ollama-host", default=None, help="Ollama API URL")
204
+ sp.add_argument("--model", default="tokioai-coder", help="Model name")
205
+ sp.add_argument("-v", "--verbose", action="store_true")
206
+
207
+ args = parser.parse_args()
208
+
209
+ if args.command == "serve":
210
+ cmd_serve(args)
211
+ elif args.command == "chat":
212
+ cmd_chat(args)
213
+ elif args.command == "test":
214
+ cmd_test(args)
215
+ elif args.command == "bench":
216
+ cmd_bench(args)
217
+ elif args.command == "stats":
218
+ cmd_stats(args)
219
+ else:
220
+ parser.print_help()
221
+
222
+
223
+ if __name__ == "__main__":
224
+ main()
engine/executor.py ADDED
@@ -0,0 +1,331 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ TokioAI Engine v3.0 -- Server-Side Tool Executor
3
+ Executes tools LOCALLY on the Pi without CLI round-trips.
4
+ This is what makes the Engine think like Opus: it can investigate autonomously.
5
+
6
+ Safety: only execute_local, read_file, search_files, diagnose are allowed server-side.
7
+ Write operations (write_file, edit_file) are ALWAYS delegated to the CLI for user approval.
8
+ Remote operations (execute_raspi, execute_gcp, execute_router) are also delegated.
9
+ """
10
+
11
+ import os
12
+ import re
13
+ import subprocess
14
+ import time
15
+ import glob as glob_module
16
+ from typing import Dict, Tuple, Optional
17
+
18
+
19
+ # Tools that the Engine can execute directly (read-only / diagnostic)
20
+ SAFE_TOOLS = {
21
+ "execute_local",
22
+ "read_file",
23
+ "search_files",
24
+ "diagnose",
25
+ "robot_vision",
26
+ }
27
+
28
+ # Tools that must be delegated to the CLI (writes, remote)
29
+ DELEGATE_TOOLS = {
30
+ "write_file",
31
+ "edit_file",
32
+ "execute_raspi",
33
+ "execute_gcp",
34
+ "execute_router",
35
+ "memory",
36
+ "task",
37
+ }
38
+
39
+ # Commands that are NEVER allowed server-side (destructive)
40
+ BLOCKED_COMMANDS = [
41
+ r'\brm\s+-rf\s+/',
42
+ r'\bmkfs\b',
43
+ r'\bdd\s+if=',
44
+ r'\b(reboot|shutdown|poweroff|halt)\b',
45
+ r'\bsudo\s+(rm|mkfs|dd|reboot|shutdown|poweroff|halt)\b',
46
+ r'>\s*/dev/',
47
+ r'\bchmod\s+-R\s+777\s+/',
48
+ r'\bchown\s+-R.*\s+/',
49
+ ]
50
+
51
+
52
+ def is_safe_tool(tool_name: str) -> bool:
53
+ """Check if a tool can be executed server-side."""
54
+ return tool_name in SAFE_TOOLS
55
+
56
+
57
+ def is_command_blocked(command: str) -> bool:
58
+ """Check if a shell command is blocked for safety."""
59
+ for pattern in BLOCKED_COMMANDS:
60
+ if re.search(pattern, command, re.IGNORECASE):
61
+ return True
62
+ return False
63
+
64
+
65
+ def execute_tool(tool_name: str, arguments: Dict, timeout: int = 30) -> Tuple[bool, str]:
66
+ """
67
+ Execute a tool server-side.
68
+
69
+ Returns:
70
+ (success: bool, output: str)
71
+
72
+ If the tool is not in SAFE_TOOLS, returns (False, "DELEGATE") to signal
73
+ that this tool should be sent back to the CLI.
74
+ """
75
+ if tool_name not in SAFE_TOOLS:
76
+ return False, "DELEGATE"
77
+
78
+ try:
79
+ if tool_name == "execute_local":
80
+ return _execute_local(arguments, timeout)
81
+ elif tool_name == "read_file":
82
+ return _read_file(arguments)
83
+ elif tool_name == "search_files":
84
+ return _search_files(arguments)
85
+ elif tool_name == "diagnose":
86
+ return _diagnose(arguments)
87
+ elif tool_name == "robot_vision":
88
+ return _robot_vision(arguments)
89
+ else:
90
+ return False, f"Unknown tool: {tool_name}"
91
+ except Exception as e:
92
+ return False, f"Error executing {tool_name}: {str(e)}"
93
+
94
+
95
+ def _execute_local(arguments: Dict, default_timeout: int = 30) -> Tuple[bool, str]:
96
+ """Execute a shell command locally."""
97
+ command = arguments.get("command", "")
98
+ if not command:
99
+ return False, "No command specified"
100
+
101
+ if is_command_blocked(command):
102
+ return False, f"Command blocked for safety: {command}"
103
+
104
+ timeout = min(arguments.get("timeout", default_timeout), 60) # Cap at 60s for server-side
105
+
106
+ try:
107
+ result = subprocess.run(
108
+ command,
109
+ shell=True,
110
+ capture_output=True,
111
+ text=True,
112
+ timeout=timeout,
113
+ env={**os.environ, "TERM": "dumb"}, # prevent color codes
114
+ )
115
+
116
+ output = result.stdout
117
+ if result.stderr:
118
+ output += ("\n" if output else "") + result.stderr
119
+
120
+ output = output.strip()
121
+
122
+ if not output:
123
+ if result.returncode == 0:
124
+ output = "(no output)"
125
+ else:
126
+ output = f"(exit code {result.returncode}, no output)"
127
+
128
+ # Truncate very long output to prevent context explosion
129
+ if len(output) > 3000:
130
+ output = output[:2500] + f"\n... [truncated, {len(output)} total chars]"
131
+
132
+ success = (result.returncode == 0)
133
+ return success, output
134
+
135
+ except subprocess.TimeoutExpired:
136
+ return False, f"Command timed out after {timeout}s"
137
+ except Exception as e:
138
+ return False, f"Execution error: {str(e)}"
139
+
140
+
141
+ def _read_file(arguments: Dict) -> Tuple[bool, str]:
142
+ """Read a file from the local filesystem."""
143
+ path = arguments.get("path", "")
144
+ if not path:
145
+ return False, "No path specified"
146
+
147
+ # Only allow local paths (no raspi: or gcp: prefix)
148
+ if ":" in path and not path.startswith("/"):
149
+ return False, "DELEGATE" # Remote file, delegate to CLI
150
+
151
+ # Expand user home
152
+ path = os.path.expanduser(path)
153
+
154
+ if not os.path.exists(path):
155
+ return False, f"No such file or directory: {path}"
156
+
157
+ if os.path.isdir(path):
158
+ return False, f"Is a directory: {path}"
159
+
160
+ max_lines = arguments.get("lines", None)
161
+
162
+ try:
163
+ with open(path, "r", errors="replace") as f:
164
+ if max_lines:
165
+ content = "".join(f.readline() for _ in range(max_lines))
166
+ else:
167
+ content = f.read()
168
+
169
+ # Truncate
170
+ if len(content) > 5000:
171
+ content = content[:4500] + f"\n... [truncated, {len(content)} total chars]"
172
+
173
+ return True, content.strip() if content.strip() else "(empty file)"
174
+
175
+ except PermissionError:
176
+ return False, f"Permission denied: {path}"
177
+ except Exception as e:
178
+ return False, f"Error reading {path}: {str(e)}"
179
+
180
+
181
+ def _search_files(arguments: Dict) -> Tuple[bool, str]:
182
+ """Search for patterns in files (like grep -rn)."""
183
+ pattern = arguments.get("pattern", "")
184
+ if not pattern:
185
+ return False, "No pattern specified"
186
+
187
+ search_path = arguments.get("path", ".")
188
+
189
+ # Only local paths
190
+ if ":" in search_path and not search_path.startswith("/"):
191
+ return False, "DELEGATE"
192
+
193
+ search_path = os.path.expanduser(search_path)
194
+ file_glob = arguments.get("glob", "*")
195
+
196
+ try:
197
+ # Use grep for efficiency
198
+ cmd = f"grep -rn --include='{file_glob}' '{pattern}' {search_path} 2>/dev/null | head -50"
199
+ result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=15)
200
+ output = result.stdout.strip()
201
+
202
+ if not output:
203
+ return True, f"No matches found for '{pattern}' in {search_path}"
204
+
205
+ if len(output) > 3000:
206
+ output = output[:2500] + "\n... [truncated]"
207
+
208
+ return True, output
209
+
210
+ except Exception as e:
211
+ return False, f"Search error: {str(e)}"
212
+
213
+
214
+ def _diagnose(arguments: Dict) -> Tuple[bool, str]:
215
+ """Run basic diagnostics on the local system."""
216
+ target = arguments.get("target", "local")
217
+
218
+ if target != "local":
219
+ return False, "DELEGATE" # Remote diagnostics delegate to CLI
220
+
221
+ checks = []
222
+
223
+ # CPU load
224
+ try:
225
+ with open("/proc/loadavg", "r") as f:
226
+ load = f.read().strip()
227
+ checks.append(f"Load: {load}")
228
+ except:
229
+ pass
230
+
231
+ # Memory
232
+ try:
233
+ result = subprocess.run("free -h | head -3", shell=True, capture_output=True, text=True, timeout=5)
234
+ checks.append(f"Memory:\n{result.stdout.strip()}")
235
+ except:
236
+ pass
237
+
238
+ # Disk
239
+ try:
240
+ result = subprocess.run("df -h / | tail -1", shell=True, capture_output=True, text=True, timeout=5)
241
+ checks.append(f"Disk: {result.stdout.strip()}")
242
+ except:
243
+ pass
244
+
245
+ # Docker
246
+ try:
247
+ result = subprocess.run("docker ps --format 'table {{.Names}}\t{{.Status}}' 2>/dev/null | head -10", shell=True, capture_output=True, text=True, timeout=5)
248
+ if result.returncode == 0 and result.stdout.strip():
249
+ checks.append(f"Docker:\n{result.stdout.strip()}")
250
+ except:
251
+ pass
252
+
253
+ # Key services
254
+ try:
255
+ result = subprocess.run("systemctl list-units --type=service --state=running --no-pager | head -20", shell=True, capture_output=True, text=True, timeout=5)
256
+ checks.append(f"Services:\n{result.stdout.strip()}")
257
+ except:
258
+ pass
259
+
260
+ return True, "\n\n".join(checks) if checks else "Could not gather diagnostics"
261
+
262
+
263
+ def _robot_vision(arguments: Dict) -> Tuple[bool, str]:
264
+ """Capture a snapshot from the robot's camera and analyze with moondream vision model."""
265
+ import requests
266
+ import base64
267
+ import time
268
+
269
+ ROBOT_PROXY = os.environ.get("ROBOT_PROXY_URL", "http://localhost:5002")
270
+ OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:8200")
271
+ VISION_MODEL = os.environ.get("VISION_MODEL", "moondream")
272
+
273
+ question = arguments.get("question", "")
274
+ if not question:
275
+ question = "Describe what you see in this image in detail. What objects, people, furniture, or environment is visible? Be specific."
276
+
277
+ # 1. Capture snapshot from robot camera
278
+ t0 = time.time()
279
+ try:
280
+ resp = requests.get(f"{ROBOT_PROXY}/snapshot", timeout=10)
281
+ resp.raise_for_status()
282
+ img_data = resp.content
283
+ except requests.ConnectionError:
284
+ return False, "Robot camera unavailable -- cannot connect to robot proxy (localhost:5002). Is picar-forward.service running?"
285
+ except Exception as e:
286
+ return False, f"Error capturing snapshot from robot camera: {e}"
287
+
288
+ t_capture = time.time() - t0
289
+
290
+ if len(img_data) < 500:
291
+ return False, f"Robot camera returned suspiciously small image ({len(img_data)} bytes). Camera may be offline."
292
+
293
+ # 2. Encode to base64
294
+ img_b64 = base64.b64encode(img_data).decode('utf-8')
295
+
296
+ # 3. Analyze with vision model
297
+ t1 = time.time()
298
+ try:
299
+ resp = requests.post(
300
+ f"{OLLAMA_URL}/api/generate",
301
+ json={
302
+ "model": VISION_MODEL,
303
+ "prompt": question,
304
+ "images": [img_b64],
305
+ "stream": False,
306
+ },
307
+ timeout=60,
308
+ )
309
+ resp.raise_for_status()
310
+ result = resp.json()
311
+ except requests.ConnectionError:
312
+ return False, "Vision model unavailable -- cannot connect to Ollama (localhost:8200). Is tokioai-coder-tunnel.service running?"
313
+ except Exception as e:
314
+ return False, f"Error analyzing image with vision model: {e}"
315
+
316
+ t_analysis = time.time() - t1
317
+
318
+ description = result.get("response", "").strip()
319
+ if not description:
320
+ return False, "Vision model returned empty response."
321
+
322
+ tokens = result.get("eval_count", 0)
323
+
324
+ output = (
325
+ f"[Robot Vision -- moondream]\n"
326
+ f"Image: {len(img_data):,} bytes | Capture: {t_capture:.1f}s | Analysis: {t_analysis:.1f}s | Tokens: {tokens}\n"
327
+ f"Question: {question}\n"
328
+ f"Answer: {description}"
329
+ )
330
+
331
+ return True, output
engine/guided.py ADDED
@@ -0,0 +1,228 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ TokioAI Engine -- Guided Generation
3
+ Forces the model to produce valid structured output using Ollama's format parameter.
4
+
5
+ Instead of hoping the model generates valid JSON, we CONSTRAIN it.
6
+ Ollama supports JSON mode natively, and we combine that with our tool schemas
7
+ to guarantee valid output.
8
+
9
+ For non-Ollama backends, we use post-processing + retry as fallback.
10
+ """
11
+
12
+ import json
13
+ from typing import Any, Dict, List, Optional
14
+
15
+ from .schemas import TOOL_DEFINITIONS, TOOL_SCHEMAS, VALID_TOOL_NAMES
16
+
17
+
18
+ def get_tool_json_schema() -> Dict:
19
+ """
20
+ Generate a JSON schema that constrains model output to valid tool calls.
21
+ This is used with Ollama's `format` parameter for structured output.
22
+ """
23
+ # Build a union schema: the model must output one of the valid tool call shapes
24
+ tool_schemas = []
25
+
26
+ for tool_def in TOOL_DEFINITIONS:
27
+ func = tool_def["function"]
28
+ tool_schemas.append({
29
+ "type": "object",
30
+ "properties": {
31
+ "name": {
32
+ "type": "string",
33
+ "const": func["name"],
34
+ },
35
+ "arguments": func["parameters"],
36
+ },
37
+ "required": ["name", "arguments"],
38
+ "additionalProperties": False,
39
+ })
40
+
41
+ # anyOf: model can output any one of these tool call shapes
42
+ return {
43
+ "type": "object",
44
+ "anyOf": tool_schemas,
45
+ }
46
+
47
+
48
+ def get_simple_json_schema() -> Dict:
49
+ """
50
+ Simplified schema: just name (enum) + arguments (object).
51
+ More compatible with Ollama's format support.
52
+ """
53
+ return {
54
+ "type": "object",
55
+ "properties": {
56
+ "name": {
57
+ "type": "string",
58
+ "enum": sorted(VALID_TOOL_NAMES),
59
+ },
60
+ "arguments": {
61
+ "type": "object",
62
+ },
63
+ },
64
+ "required": ["name", "arguments"],
65
+ }
66
+
67
+
68
+ def get_text_response_schema() -> Dict:
69
+ """
70
+ Schema for text-only responses (no tool calls).
71
+ Forces the model to output a simple response object.
72
+ """
73
+ return {
74
+ "type": "object",
75
+ "properties": {
76
+ "response": {
77
+ "type": "string",
78
+ "description": "Your text response to the user",
79
+ },
80
+ "needs_tool": {
81
+ "type": "boolean",
82
+ "description": "Set to true if you think a tool should be used instead",
83
+ },
84
+ },
85
+ "required": ["response"],
86
+ }
87
+
88
+
89
+ def get_decision_schema() -> Dict:
90
+ """
91
+ Schema for decision mode: model chooses between tool call or text response.
92
+ """
93
+ return {
94
+ "type": "object",
95
+ "properties": {
96
+ "action": {
97
+ "type": "string",
98
+ "enum": ["tool_call", "text_response"],
99
+ "description": "Whether to call a tool or respond with text",
100
+ },
101
+ "tool_call": {
102
+ "type": "object",
103
+ "properties": {
104
+ "name": {
105
+ "type": "string",
106
+ "enum": sorted(VALID_TOOL_NAMES),
107
+ },
108
+ "arguments": {
109
+ "type": "object",
110
+ },
111
+ },
112
+ "description": "If action is tool_call, provide the tool call here",
113
+ },
114
+ "text_response": {
115
+ "type": "string",
116
+ "description": "If action is text_response, provide your response here",
117
+ },
118
+ },
119
+ "required": ["action"],
120
+ }
121
+
122
+
123
+ def build_ollama_request(
124
+ model: str,
125
+ messages: List[Dict],
126
+ intent: str,
127
+ force_json: bool = True,
128
+ temperature: float = 0.3,
129
+ max_tokens: int = 4096,
130
+ stream: bool = False,
131
+ ) -> Dict:
132
+ """
133
+ Build an Ollama /api/chat request with optional guided generation.
134
+
135
+ Args:
136
+ model: Ollama model name
137
+ messages: Chat messages
138
+ intent: Classified intent (determines schema)
139
+ force_json: Whether to constrain output to JSON
140
+ temperature: Sampling temperature
141
+ max_tokens: Max tokens to generate
142
+ stream: Whether to stream the response
143
+
144
+ Returns:
145
+ Request body dict ready for POST to /api/chat
146
+ """
147
+ request = {
148
+ "model": model,
149
+ "messages": messages,
150
+ "stream": stream,
151
+ "options": {
152
+ "temperature": temperature,
153
+ "num_predict": max_tokens,
154
+ },
155
+ }
156
+
157
+ if force_json and intent in ("tool_call", "decision"):
158
+ # Use Ollama's structured output
159
+ if intent == "tool_call":
160
+ request["format"] = get_simple_json_schema()
161
+ elif intent == "decision":
162
+ request["format"] = get_decision_schema()
163
+ elif force_json and intent == "text":
164
+ # For text, we DON'T force JSON -- let the model speak naturally
165
+ pass
166
+
167
+ # Add tools in OpenAI format (Ollama supports this)
168
+ if intent in ("tool_call", "decision"):
169
+ request["tools"] = TOOL_DEFINITIONS
170
+
171
+ return request
172
+
173
+
174
+ def parse_guided_output(raw_output: str, intent: str) -> Dict:
175
+ """
176
+ Parse the guided (JSON-constrained) output from the model.
177
+
178
+ Returns a normalized dict:
179
+ - For tool_call: {"type": "tool_call", "name": str, "arguments": dict}
180
+ - For text: {"type": "text", "response": str}
181
+ - For decision: {"type": "tool_call"|"text", ...}
182
+ """
183
+ try:
184
+ data = json.loads(raw_output)
185
+ except json.JSONDecodeError:
186
+ # Try to extract JSON from text
187
+ import re
188
+ match = re.search(r'\{.*\}', raw_output, re.DOTALL)
189
+ if match:
190
+ try:
191
+ data = json.loads(match.group())
192
+ except json.JSONDecodeError:
193
+ return {"type": "text", "response": raw_output}
194
+ else:
195
+ return {"type": "text", "response": raw_output}
196
+
197
+ # Decision format
198
+ if "action" in data:
199
+ if data["action"] == "tool_call" and "tool_call" in data:
200
+ tc = data["tool_call"]
201
+ return {
202
+ "type": "tool_call",
203
+ "name": tc.get("name", ""),
204
+ "arguments": tc.get("arguments", {}),
205
+ }
206
+ elif data["action"] == "text_response":
207
+ return {
208
+ "type": "text",
209
+ "response": data.get("text_response", ""),
210
+ }
211
+
212
+ # Direct tool call format
213
+ if "name" in data and data["name"] in VALID_TOOL_NAMES:
214
+ return {
215
+ "type": "tool_call",
216
+ "name": data["name"],
217
+ "arguments": data.get("arguments", {}),
218
+ }
219
+
220
+ # Text response format
221
+ if "response" in data:
222
+ return {
223
+ "type": "text",
224
+ "response": data["response"],
225
+ }
226
+
227
+ # Unknown format -- return as text
228
+ return {"type": "text", "response": json.dumps(data)}
engine/memory.py ADDED
@@ -0,0 +1,353 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ TokioAI Engine -- Adaptive Memory / Few-Shot Learning
3
+ Maintains a persistent store of successful interactions for dynamic few-shot injection.
4
+
5
+ This is the "fine-tuning in real time" -- instead of modifying weights,
6
+ we modify the context the model sees, injecting relevant examples that
7
+ worked before.
8
+
9
+ Storage: JSON file, indexed by intent + keywords for fast retrieval.
10
+ """
11
+
12
+ import json
13
+ import os
14
+ import re
15
+ import time
16
+ from collections import Counter
17
+ from typing import Dict, List, Optional
18
+
19
+ DEFAULT_MEMORY_PATH = os.path.expanduser("~/.tokioai-engine/memory.json")
20
+ MAX_EXAMPLES_PER_CATEGORY = 100
21
+ MAX_EXAMPLE_AGE_DAYS = 90
22
+
23
+
24
+ class EngineMemory:
25
+ """
26
+ Persistent few-shot memory that makes the model smarter over time.
27
+
28
+ Stores successful (user_input -> model_output) pairs, tagged by:
29
+ - intent (tool_call, code, text, decision)
30
+ - keywords (extracted from user input)
31
+ - success (was the output valid/useful?)
32
+ - timestamp
33
+
34
+ On each query, retrieves the most relevant examples for few-shot injection.
35
+ """
36
+
37
+ def __init__(self, path: str = None):
38
+ self.path = path or DEFAULT_MEMORY_PATH
39
+ os.makedirs(os.path.dirname(self.path), exist_ok=True)
40
+ self.examples: List[Dict] = []
41
+ self._load()
42
+
43
+ def _load(self):
44
+ if os.path.exists(self.path):
45
+ try:
46
+ with open(self.path, "r") as f:
47
+ self.examples = json.load(f)
48
+ except (json.JSONDecodeError, IOError):
49
+ self.examples = []
50
+ else:
51
+ self.examples = self._seed_examples()
52
+ self._save()
53
+
54
+ def _save(self):
55
+ with open(self.path, "w") as f:
56
+ json.dump(self.examples, f, indent=2, ensure_ascii=False)
57
+
58
+ def add(
59
+ self,
60
+ user_input: str,
61
+ model_output: str,
62
+ intent: str,
63
+ success: bool = True,
64
+ tool_name: str = None,
65
+ latency_ms: int = None,
66
+ ):
67
+ """Record a successful interaction as a few-shot example."""
68
+ # NEVER store errors, empty outputs, or hallucinated tool output
69
+ if not model_output or len(model_output.strip()) < 5:
70
+ return
71
+ if "Error calling model" in model_output or "Client Error" in model_output:
72
+ return
73
+ if "Server Error" in model_output or "Traceback" in model_output:
74
+ return
75
+ if not success:
76
+ return
77
+ # Skip duplicate: same user input + same first 100 chars of output
78
+ key = (user_input[:100], model_output[:100])
79
+ for ex in self.examples[-50:]: # Check last 50 only for speed
80
+ if (ex.get("user", "")[:100], ex.get("assistant", "")[:100]) == key:
81
+ return
82
+
83
+ keywords = self._extract_keywords(user_input)
84
+
85
+ entry = {
86
+ "user": user_input[:500],
87
+ "assistant": model_output[:1000],
88
+ "intent": intent,
89
+ "keywords": keywords,
90
+ "success": success,
91
+ "tool_name": tool_name,
92
+ "latency_ms": latency_ms,
93
+ "timestamp": time.time(),
94
+ "uses": 0, # How many times this was used as few-shot
95
+ }
96
+
97
+ self.examples.append(entry)
98
+
99
+ # Prune old/excess entries
100
+ self._prune()
101
+ self._save()
102
+
103
+ def retrieve(
104
+ self,
105
+ user_input: str,
106
+ intent: str = None,
107
+ top_k: int = 3,
108
+ ) -> List[Dict]:
109
+ """
110
+ Retrieve the most relevant few-shot examples for a given query.
111
+ Uses keyword overlap + intent matching + recency for ranking.
112
+ """
113
+ if not self.examples:
114
+ return []
115
+
116
+ query_keywords = set(self._extract_keywords(user_input))
117
+ scored = []
118
+
119
+ for ex in self.examples:
120
+ if not ex.get("success", True):
121
+ continue # Skip failed examples
122
+
123
+ score = 0.0
124
+
125
+ # Keyword overlap (most important)
126
+ ex_keywords = set(ex.get("keywords", []))
127
+ overlap = query_keywords & ex_keywords
128
+ if overlap:
129
+ score += len(overlap) * 3.0
130
+
131
+ # Intent match
132
+ if intent and ex.get("intent") == intent:
133
+ score += 2.0
134
+
135
+ # Recency bonus (newer = better, decay over 30 days)
136
+ age_days = (time.time() - ex.get("timestamp", 0)) / 86400
137
+ recency = max(0, 1.0 - age_days / 30)
138
+ score += recency
139
+
140
+ # Usage penalty (avoid showing the same examples every time)
141
+ score -= ex.get("uses", 0) * 0.5
142
+
143
+ if score > 0:
144
+ scored.append((score, ex))
145
+
146
+ # Sort by score descending
147
+ scored.sort(key=lambda x: x[0], reverse=True)
148
+
149
+ # Mark as used
150
+ results = []
151
+ for score, ex in scored[:top_k]:
152
+ ex["uses"] = ex.get("uses", 0) + 1
153
+ results.append(ex)
154
+
155
+ if results:
156
+ self._save()
157
+
158
+ return results
159
+
160
+ def get_stats(self) -> Dict:
161
+ """Return memory statistics."""
162
+ if not self.examples:
163
+ return {"total": 0}
164
+
165
+ intents = Counter(ex.get("intent", "unknown") for ex in self.examples)
166
+ tools = Counter(ex.get("tool_name", "none") for ex in self.examples if ex.get("tool_name"))
167
+ success_rate = sum(1 for ex in self.examples if ex.get("success", True)) / len(self.examples)
168
+
169
+ return {
170
+ "total": len(self.examples),
171
+ "by_intent": dict(intents),
172
+ "by_tool": dict(tools),
173
+ "success_rate": round(success_rate, 2),
174
+ "oldest_days": round((time.time() - min(ex.get("timestamp", time.time()) for ex in self.examples)) / 86400, 1),
175
+ }
176
+
177
+ def feedback(self, index: int, success: bool):
178
+ """Update success status for an example (for learning from failures)."""
179
+ if 0 <= index < len(self.examples):
180
+ self.examples[index]["success"] = success
181
+ self._save()
182
+
183
+ def _extract_keywords(self, text: str) -> List[str]:
184
+ """Extract meaningful keywords from text."""
185
+ # Lowercase, split, remove common words
186
+ stopwords = {
187
+ "the", "a", "an", "is", "are", "was", "were", "be", "been",
188
+ "have", "has", "had", "do", "does", "did", "will", "would",
189
+ "could", "should", "may", "might", "can", "shall", "must",
190
+ "to", "of", "in", "for", "on", "with", "at", "by", "from",
191
+ "it", "its", "this", "that", "these", "those", "my", "your",
192
+ "i", "me", "we", "us", "you", "he", "she", "they", "them",
193
+ "and", "or", "but", "not", "if", "then", "else", "so",
194
+ "what", "how", "why", "when", "where", "which", "who",
195
+ "que", "el", "la", "los", "las", "un", "una", "de", "en",
196
+ "por", "para", "con", "es", "no", "si", "como", "me",
197
+ "please", "help", "want", "need", "tell",
198
+ }
199
+
200
+ words = re.findall(r'[a-z0-9_./:-]+', text.lower())
201
+ keywords = [w for w in words if w not in stopwords and len(w) > 1]
202
+
203
+ # Also extract paths and commands
204
+ paths = re.findall(r'/[\w./+-]+', text)
205
+ keywords.extend(paths)
206
+
207
+ return keywords[:20] # Cap at 20 keywords
208
+
209
+ def _prune(self):
210
+ """Remove old or excess entries."""
211
+ now = time.time()
212
+ max_age = MAX_EXAMPLE_AGE_DAYS * 86400
213
+
214
+ # Remove old entries
215
+ self.examples = [
216
+ ex for ex in self.examples
217
+ if now - ex.get("timestamp", now) < max_age
218
+ ]
219
+
220
+ # If still too many, keep the most successful/used ones
221
+ if len(self.examples) > MAX_EXAMPLES_PER_CATEGORY * 5:
222
+ # Sort by (success, -age) and keep top N
223
+ self.examples.sort(
224
+ key=lambda ex: (ex.get("success", True), -ex.get("timestamp", 0)),
225
+ reverse=True,
226
+ )
227
+ self.examples = self.examples[:MAX_EXAMPLES_PER_CATEGORY * 5]
228
+
229
+ def _seed_examples(self) -> List[Dict]:
230
+ """Seed initial few-shot examples for cold start."""
231
+ now = time.time()
232
+ seeds = [
233
+ # Tool call examples
234
+ {
235
+ "user": "check disk space",
236
+ "assistant": '{"name": "execute_local", "arguments": {"command": "df -h"}}',
237
+ "intent": "tool_call",
238
+ "keywords": ["check", "disk", "space"],
239
+ "tool_name": "execute_local",
240
+ "success": True,
241
+ "timestamp": now,
242
+ "uses": 0,
243
+ },
244
+ {
245
+ "user": "what processes are using the most memory",
246
+ "assistant": '{"name": "execute_local", "arguments": {"command": "ps aux --sort=-%mem | head -20"}}',
247
+ "intent": "tool_call",
248
+ "keywords": ["processes", "memory"],
249
+ "tool_name": "execute_local",
250
+ "success": True,
251
+ "timestamp": now,
252
+ "uses": 0,
253
+ },
254
+ {
255
+ "user": "read the nginx config",
256
+ "assistant": '{"name": "read_file", "arguments": {"path": "/etc/nginx/nginx.conf"}}',
257
+ "intent": "tool_call",
258
+ "keywords": ["read", "nginx", "config"],
259
+ "tool_name": "read_file",
260
+ "success": True,
261
+ "timestamp": now,
262
+ "uses": 0,
263
+ },
264
+ {
265
+ "user": "find all python files with TODO comments",
266
+ "assistant": '{"name": "search_files", "arguments": {"pattern": "TODO", "glob": "*.py"}}',
267
+ "intent": "tool_call",
268
+ "keywords": ["find", "python", "todo"],
269
+ "tool_name": "search_files",
270
+ "success": True,
271
+ "timestamp": now,
272
+ "uses": 0,
273
+ },
274
+ {
275
+ "user": "check system health",
276
+ "assistant": '{"name": "diagnose", "arguments": {"target": "local"}}',
277
+ "intent": "tool_call",
278
+ "keywords": ["check", "system", "health"],
279
+ "tool_name": "diagnose",
280
+ "success": True,
281
+ "timestamp": now,
282
+ "uses": 0,
283
+ },
284
+ {
285
+ "user": "show me the contents of /etc/hostname",
286
+ "assistant": '{"name": "read_file", "arguments": {"path": "/etc/hostname"}}',
287
+ "intent": "tool_call",
288
+ "keywords": ["show", "contents", "/etc/hostname"],
289
+ "tool_name": "read_file",
290
+ "success": True,
291
+ "timestamp": now,
292
+ "uses": 0,
293
+ },
294
+ {
295
+ "user": "list docker containers",
296
+ "assistant": '{"name": "execute_local", "arguments": {"command": "docker ps -a"}}',
297
+ "intent": "tool_call",
298
+ "keywords": ["docker", "containers", "list"],
299
+ "tool_name": "execute_local",
300
+ "success": True,
301
+ "timestamp": now,
302
+ "uses": 0,
303
+ },
304
+ {
305
+ "user": "scan ports on 192.168.1.1",
306
+ "assistant": '{"name": "execute_local", "arguments": {"command": "nmap -sV 192.168.1.1"}}',
307
+ "intent": "tool_call",
308
+ "keywords": ["scan", "ports", "nmap", "192.168.1.1"],
309
+ "tool_name": "execute_local",
310
+ "success": True,
311
+ "timestamp": now,
312
+ "uses": 0,
313
+ },
314
+ # Text examples (NO tool calls)
315
+ {
316
+ "user": "what is SQL injection",
317
+ "assistant": "SQL injection is an attack where malicious SQL code is inserted into application queries through user input. The attacker can read, modify, or delete database data, bypass authentication, or even execute system commands. Prevention: parameterized queries, input validation, least-privilege DB accounts, WAF rules.",
318
+ "intent": "text",
319
+ "keywords": ["sql", "injection"],
320
+ "success": True,
321
+ "timestamp": now,
322
+ "uses": 0,
323
+ },
324
+ {
325
+ "user": "explain the TCP 3-way handshake",
326
+ "assistant": "TCP 3-way handshake establishes a connection: (1) Client sends SYN with initial sequence number, (2) Server responds SYN-ACK with its own sequence number and acknowledges client's, (3) Client sends ACK. After this, data flows bidirectionally. Termination uses a 4-way FIN handshake.",
327
+ "intent": "text",
328
+ "keywords": ["tcp", "handshake"],
329
+ "success": True,
330
+ "timestamp": now,
331
+ "uses": 0,
332
+ },
333
+ {
334
+ "user": "hi",
335
+ "assistant": "Hey. What do you need?",
336
+ "intent": "text",
337
+ "keywords": ["hi"],
338
+ "success": True,
339
+ "timestamp": now,
340
+ "uses": 0,
341
+ },
342
+ # Decision examples (model correctly chose NOT to use tool)
343
+ {
344
+ "user": "what is docker",
345
+ "assistant": "Docker is a containerization platform that packages applications with their dependencies into isolated containers. Containers share the host kernel (unlike VMs), making them lightweight and fast. Key concepts: images (blueprints), containers (running instances), Dockerfile (build instructions), docker-compose (multi-container apps).",
346
+ "intent": "decision",
347
+ "keywords": ["docker"],
348
+ "success": True,
349
+ "timestamp": now,
350
+ "uses": 0,
351
+ },
352
+ ]
353
+ return seeds
engine/pipeline.py ADDED
@@ -0,0 +1,476 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ TokioAI Engine v2.0 -- Main Pipeline
3
+ Orchestrates: Router -> Memory -> Guided Generation -> Verifier -> Output
4
+
5
+ The pipeline wraps any Ollama-compatible model and transforms it into
6
+ an elite AI agent with multi-step tool calling.
7
+
8
+ Key change from v1: when force_intent is None, the model freely chooses
9
+ between tool calls and text responses. This enables investigation chains.
10
+ """
11
+
12
+ import json
13
+ import sys
14
+ import os
15
+ import re
16
+ import time
17
+ import requests
18
+ from typing import Any, Callable, Dict, List, Optional, Tuple
19
+
20
+ from .router import classify_intent, build_system_prompt, IntentType
21
+ from .memory import EngineMemory
22
+ from .guided import build_ollama_request, parse_guided_output
23
+ from .verifier import Verifier, RetryStrategy, VerificationResult
24
+ from .schemas import TOOL_DEFINITIONS, tools_to_prompt_text, validate_tool_call
25
+
26
+
27
+ class EngineConfig:
28
+ """Pipeline configuration."""
29
+
30
+ def __init__(
31
+ self,
32
+ ollama_host: str = None,
33
+ model_name: str = "tokioai-coder",
34
+ max_retries: int = 2,
35
+ temperature: float = 0.3,
36
+ max_tokens: int = 4096,
37
+ use_guided_gen: bool = True,
38
+ use_memory: bool = True,
39
+ memory_path: str = None,
40
+ verbose: bool = False,
41
+ ):
42
+ self.ollama_host = ollama_host or os.getenv(
43
+ "TOKIOAI_ENGINE_HOST",
44
+ os.getenv("OLLAMA_HOST", "http://localhost:11434")
45
+ )
46
+ self.model_name = model_name
47
+ self.max_retries = max_retries
48
+ self.temperature = temperature
49
+ self.max_tokens = max_tokens
50
+ self.use_guided_gen = use_guided_gen
51
+ self.use_memory = use_memory
52
+ self.memory_path = memory_path
53
+ self.verbose = verbose
54
+
55
+
56
+ class PipelineResult:
57
+ """Result from the pipeline."""
58
+
59
+ def __init__(
60
+ self,
61
+ output_type: str, # "tool_call", "text"
62
+ tool_calls: List[Dict] = None,
63
+ text: str = "",
64
+ intent: str = "",
65
+ confidence: float = 0.0,
66
+ attempts: int = 1,
67
+ latency_ms: int = 0,
68
+ from_memory: bool = False,
69
+ ):
70
+ self.output_type = output_type
71
+ self.tool_calls = tool_calls or []
72
+ self.text = text
73
+ self.intent = intent
74
+ self.confidence = confidence
75
+ self.attempts = attempts
76
+ self.latency_ms = latency_ms
77
+ self.from_memory = from_memory
78
+
79
+ def to_openai_message(self) -> Dict:
80
+ """Convert to OpenAI-style assistant message."""
81
+ if self.output_type == "tool_call" and self.tool_calls:
82
+ return {
83
+ "role": "assistant",
84
+ "content": None,
85
+ "tool_calls": [
86
+ {
87
+ "id": f"call_{i}",
88
+ "type": "function",
89
+ "function": {
90
+ "name": tc["name"],
91
+ "arguments": json.dumps(tc["arguments"]),
92
+ },
93
+ }
94
+ for i, tc in enumerate(self.tool_calls)
95
+ ],
96
+ }
97
+ return {
98
+ "role": "assistant",
99
+ "content": self.text,
100
+ }
101
+
102
+ def __repr__(self):
103
+ if self.output_type == "tool_call":
104
+ names = [tc["name"] for tc in self.tool_calls]
105
+ return f"PipelineResult(tool_calls={names}, attempts={self.attempts}, {self.latency_ms}ms)"
106
+ return f"PipelineResult(text={self.text[:80]!r}, attempts={self.attempts}, {self.latency_ms}ms)"
107
+
108
+
109
+ class Pipeline:
110
+ """
111
+ The TokioAI Engine pipeline.
112
+
113
+ Usage:
114
+ engine = Pipeline(EngineConfig(model_name="tokioai-coder"))
115
+ result = engine.process("check disk space")
116
+ # result.tool_calls = [{"name": "execute_local", "arguments": {"command": "df -h"}}]
117
+ """
118
+
119
+ def __init__(self, config: EngineConfig = None):
120
+ self.config = config or EngineConfig()
121
+ self.verifier = Verifier()
122
+ self.retry_strategy = RetryStrategy(max_retries=self.config.max_retries)
123
+ self.memory = EngineMemory(self.config.memory_path) if self.config.use_memory else None
124
+ self.conversation_history: List[Dict] = []
125
+ self._stats = {
126
+ "total_requests": 0,
127
+ "tool_calls": 0,
128
+ "text_responses": 0,
129
+ "retries": 0,
130
+ "failures": 0,
131
+ "total_latency_ms": 0,
132
+ }
133
+
134
+ def process(
135
+ self,
136
+ user_message: str,
137
+ conversation_history: List[Dict] = None,
138
+ force_intent: str = None,
139
+ ) -> PipelineResult:
140
+ """
141
+ Process a user message through the full pipeline.
142
+
143
+ Args:
144
+ user_message: The user's input
145
+ conversation_history: Previous messages (optional)
146
+ force_intent: Override intent classification
147
+ - "text": force text-only response (no tool calls allowed)
148
+ - "tool_call": force tool call
149
+ - None: model decides freely (default for investigation rounds)
150
+
151
+ Returns:
152
+ PipelineResult with tool calls or text response
153
+ """
154
+ start_time = time.time()
155
+ self._stats["total_requests"] += 1
156
+
157
+ history = conversation_history or self.conversation_history
158
+
159
+ # === STAGE 1: ROUTE ===
160
+ if force_intent:
161
+ intent, confidence = force_intent, 1.0
162
+ else:
163
+ intent, confidence = classify_intent(user_message, history)
164
+
165
+ if self.config.verbose:
166
+ print(f"[ENGINE] Intent: {intent} (confidence: {confidence:.2f})")
167
+
168
+ # === PURE TEXT SHORT-CIRCUIT ===
169
+ # For greetings and thanks, strip conversation history to prevent
170
+ # context bleeding (e.g. "hola" after HA investigation should NOT talk about HA).
171
+ # BUT NOT for affirmations ("dale", "si", "ok") which need context.
172
+ import re as _re
173
+ _GREETING_PATS = [
174
+ r'^(hi|hello|hey|hola|que\s+tal|buenos?\s+d[ií]as?|buenas?\s+(?:tardes?|noches?))\s*[.!?]*$',
175
+ r'^(thanks?|thank\s+you|gracias|thx|de\s+nada)\s*[.!?]*$',
176
+ r'^(como\s+(va|estas?|andas?|te\s+va)|how\s+are\s+you)\s*[.!?]*$',
177
+ r'^(todo\s+bien|bien\s+y\s+tu)\s*[.!?]*$',
178
+ r'\b(who\s+are\s+you|what\s+can\s+you|tell\s+me\s+about\s+yourself)\b',
179
+ ]
180
+ _msg_lower = user_message.lower().strip()
181
+ is_pure_text = (
182
+ intent == "text" and confidence >= 0.9
183
+ and any(_re.match(p, _msg_lower, _re.IGNORECASE) for p in _GREETING_PATS)
184
+ )
185
+
186
+ # === STAGE 2: MEMORY LOOKUP ===
187
+ few_shot = []
188
+ if self.memory and self.config.use_memory and not is_pure_text:
189
+ few_shot = self.memory.retrieve(user_message, intent, top_k=3)
190
+ if self.config.verbose and few_shot:
191
+ print(f"[ENGINE] Memory: {len(few_shot)} relevant examples found")
192
+
193
+ # === STAGE 3: BUILD PROMPT ===
194
+ system_prompt = build_system_prompt(
195
+ intent=intent,
196
+ few_shot_examples=few_shot,
197
+ )
198
+
199
+ # Build messages
200
+ messages = [{"role": "system", "content": system_prompt}]
201
+
202
+ # Add conversation history -- but for pure text (greetings etc),
203
+ # only include the LAST 2 messages to avoid context bleeding
204
+ if is_pure_text:
205
+ # For greetings/thanks: minimal or no history
206
+ # This prevents "hola" from generating an HA status report
207
+ # because the previous conversation was about HA
208
+ if self.config.verbose:
209
+ print(f"[ENGINE] Pure text detected, stripping history ({len(history)} msgs)")
210
+ else:
211
+ for msg in history[-20:]:
212
+ messages.append(msg)
213
+
214
+ # Add current user message
215
+ messages.append({"role": "user", "content": user_message})
216
+
217
+ # === STAGE 4: GENERATE ===
218
+ # Use guided generation ONLY for definite tool_call intent on first request
219
+ # For investigation rounds (force_intent=None with history), let model choose
220
+ use_guided = self.config.use_guided_gen and intent == "tool_call"
221
+
222
+ result = self._generate_with_retries(
223
+ messages=messages,
224
+ intent=intent,
225
+ use_guided=use_guided,
226
+ force_text_only=(force_intent == "text"),
227
+ )
228
+
229
+ # === STAGE 5: RECORD ===
230
+ latency_ms = int((time.time() - start_time) * 1000)
231
+ result.intent = intent
232
+ result.confidence = confidence
233
+ result.latency_ms = latency_ms
234
+
235
+ self._stats["total_latency_ms"] += latency_ms
236
+
237
+ if result.output_type == "tool_call":
238
+ self._stats["tool_calls"] += 1
239
+ else:
240
+ self._stats["text_responses"] += 1
241
+
242
+ # Save to memory (only successful results)
243
+ if self.memory and result.output_type in ("tool_call", "text"):
244
+ tool_name = result.tool_calls[0]["name"] if result.tool_calls else None
245
+ output_text = (
246
+ json.dumps(result.tool_calls[0]) if result.tool_calls
247
+ else result.text[:500]
248
+ )
249
+ self.memory.add(
250
+ user_input=user_message,
251
+ model_output=output_text,
252
+ intent=intent,
253
+ success=True,
254
+ tool_name=tool_name,
255
+ latency_ms=latency_ms,
256
+ )
257
+
258
+ # Update conversation history
259
+ self.conversation_history.append({"role": "user", "content": user_message})
260
+ self.conversation_history.append(result.to_openai_message())
261
+
262
+ return result
263
+
264
+ def add_tool_result(self, tool_name: str, result: str):
265
+ """Add a tool execution result to the conversation."""
266
+ self.conversation_history.append({
267
+ "role": "tool",
268
+ "content": result,
269
+ "name": tool_name,
270
+ })
271
+
272
+ def _generate_with_retries(
273
+ self,
274
+ messages: List[Dict],
275
+ intent: str,
276
+ use_guided: bool,
277
+ force_text_only: bool = False,
278
+ ) -> PipelineResult:
279
+ """
280
+ Generate model output with retry logic on verification failure.
281
+
282
+ Args:
283
+ messages: Full message chain
284
+ intent: Classified intent
285
+ use_guided: Whether to use Ollama guided generation (JSON schema)
286
+ force_text_only: If True, NEVER return tool calls (depth limit reached)
287
+ """
288
+
289
+ for attempt in range(self.config.max_retries + 1):
290
+ # Build request
291
+ if use_guided:
292
+ request = build_ollama_request(
293
+ model=self.config.model_name,
294
+ messages=messages,
295
+ intent=intent,
296
+ force_json=True,
297
+ temperature=self.config.temperature,
298
+ max_tokens=self.config.max_tokens,
299
+ )
300
+ else:
301
+ # Free-form generation: model can output tool calls OR text
302
+ request = {
303
+ "model": self.config.model_name,
304
+ "messages": messages,
305
+ "stream": False,
306
+ "options": {
307
+ "temperature": self.config.temperature,
308
+ "num_predict": self.config.max_tokens,
309
+ },
310
+ }
311
+ # Include tools so Ollama can use native tool calling
312
+ # UNLESS we're forcing text-only mode
313
+ if not force_text_only and intent != "text":
314
+ request["tools"] = TOOL_DEFINITIONS
315
+
316
+ # Call Ollama
317
+ try:
318
+ raw_output, native_tool_calls = self._call_ollama(request)
319
+ except Exception as e:
320
+ if self.config.verbose:
321
+ print(f"[ENGINE] Ollama error: {e}")
322
+ self._stats["failures"] += 1
323
+ return PipelineResult(
324
+ output_type="text",
325
+ text=f"Error calling model: {e}",
326
+ attempts=attempt + 1,
327
+ )
328
+
329
+ if self.config.verbose:
330
+ print(f"[ENGINE] Raw output (attempt {attempt+1}): {raw_output[:200]}")
331
+ if native_tool_calls:
332
+ print(f"[ENGINE] Native tool calls: {native_tool_calls}")
333
+
334
+ # === NATIVE TOOL CALLS (Ollama returned tool_calls in response) ===
335
+ if native_tool_calls and not force_text_only:
336
+ valid_calls = []
337
+ for tc in native_tool_calls:
338
+ is_valid, err = validate_tool_call(tc["name"], tc["arguments"])
339
+ if is_valid:
340
+ valid_calls.append(tc)
341
+ elif self.config.verbose:
342
+ print(f"[ENGINE] Invalid native tool call: {err}")
343
+
344
+ if valid_calls:
345
+ return PipelineResult(
346
+ output_type="tool_call",
347
+ tool_calls=valid_calls,
348
+ attempts=attempt + 1,
349
+ )
350
+
351
+ # === GUIDED OUTPUT (JSON constrained) ===
352
+ if use_guided and intent in ("tool_call", "decision"):
353
+ parsed = parse_guided_output(raw_output, intent)
354
+ if parsed["type"] == "tool_call" and not force_text_only:
355
+ name = parsed.get("name", "")
356
+ args = parsed.get("arguments", {})
357
+ is_valid, err = validate_tool_call(name, args)
358
+ if is_valid:
359
+ return PipelineResult(
360
+ output_type="tool_call",
361
+ tool_calls=[{"name": name, "arguments": args}],
362
+ attempts=attempt + 1,
363
+ )
364
+ elif self.config.verbose:
365
+ print(f"[ENGINE] Guided validation failed: {err}")
366
+ elif parsed["type"] == "text":
367
+ return PipelineResult(
368
+ output_type="text",
369
+ text=parsed.get("response", raw_output),
370
+ attempts=attempt + 1,
371
+ )
372
+
373
+ # === FREE-FORM VERIFICATION ===
374
+ verification = self.verifier.verify(raw_output, expected_intent=intent)
375
+
376
+ if verification.valid:
377
+ if verification.output_type == "tool_call" and not force_text_only:
378
+ # Model wants to call a tool -- allow it
379
+ return PipelineResult(
380
+ output_type="tool_call",
381
+ tool_calls=verification.tool_calls,
382
+ attempts=attempt + 1,
383
+ )
384
+ else:
385
+ # Text response (or forced text-only)
386
+ text = verification.text or raw_output
387
+ # Clean up: if model tried to output tool call JSON but we're
388
+ # in text-only mode, don't show the raw JSON
389
+ if force_text_only and verification.output_type == "tool_call":
390
+ text = (
391
+ "Based on my investigation, I was unable to gather "
392
+ "additional information. Here is what I found so far."
393
+ )
394
+ return PipelineResult(
395
+ output_type="text",
396
+ text=text,
397
+ attempts=attempt + 1,
398
+ )
399
+
400
+ # Failed verification -- retry?
401
+ if self.retry_strategy.should_retry(attempt, verification):
402
+ self._stats["retries"] += 1
403
+ retry_msg = self.retry_strategy.get_retry_message(attempt, verification)
404
+ messages.append({"role": "assistant", "content": raw_output})
405
+ messages.append({"role": "user", "content": retry_msg})
406
+ if self.config.verbose:
407
+ print(f"[ENGINE] Retrying (attempt {attempt+2})...")
408
+ continue
409
+ else:
410
+ # Give up -- return whatever we got
411
+ self._stats["failures"] += 1
412
+ return PipelineResult(
413
+ output_type="text",
414
+ text=verification.text or raw_output,
415
+ attempts=attempt + 1,
416
+ )
417
+
418
+ # Should never reach here
419
+ return PipelineResult(output_type="text", text="Engine error: max retries exceeded")
420
+
421
+ def _call_ollama(self, request: Dict):
422
+ """
423
+ Call Ollama API and return (text_output, native_tool_calls).
424
+
425
+ Returns:
426
+ Tuple of (str, list):
427
+ - text content from the response
428
+ - list of native tool calls (if any), each as {name, arguments}
429
+ """
430
+ url = f"{self.config.ollama_host}/api/chat"
431
+
432
+ try:
433
+ resp = requests.post(url, json=request, timeout=120)
434
+ if self.config.verbose and resp.status_code != 200:
435
+ print(f"[ENGINE] Ollama HTTP {resp.status_code}: {resp.text[:500]}")
436
+ resp.raise_for_status()
437
+ data = resp.json()
438
+ except requests.exceptions.ConnectionError:
439
+ raise RuntimeError(f"Cannot connect to Ollama at {self.config.ollama_host}")
440
+ except requests.exceptions.Timeout:
441
+ raise RuntimeError("Ollama request timed out (120s)")
442
+ except Exception as e:
443
+ raise RuntimeError(f"Ollama request failed: {e}")
444
+
445
+ message = data.get("message", {})
446
+ content = message.get("content", "")
447
+
448
+ # Check for native tool calls from Ollama
449
+ native_tool_calls = []
450
+ tool_calls_raw = message.get("tool_calls", [])
451
+ if tool_calls_raw:
452
+ for tc in tool_calls_raw:
453
+ func = tc.get("function", {})
454
+ native_tool_calls.append({
455
+ "name": func.get("name", ""),
456
+ "arguments": func.get("arguments", {}),
457
+ })
458
+
459
+ return content, native_tool_calls
460
+
461
+ def get_stats(self) -> Dict:
462
+ """Return pipeline statistics."""
463
+ total = self._stats["total_requests"] or 1
464
+ return {
465
+ "total_requests": self._stats["total_requests"],
466
+ "tool_calls": self._stats["tool_calls"],
467
+ "text_responses": self._stats["text_responses"],
468
+ "retries": self._stats["retries"],
469
+ "failures": self._stats["failures"],
470
+ "avg_latency_ms": self._stats["total_latency_ms"] // total,
471
+ "tool_call_rate": f"{100 * self._stats['tool_calls'] // total}%",
472
+ }
473
+
474
+ def reset(self):
475
+ """Reset conversation history."""
476
+ self.conversation_history = []
engine/router.py ADDED
@@ -0,0 +1,382 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ TokioAI Engine v3.0 -- Router / Intent Classifier
3
+ Classifies user intent and builds optimal prompts with dynamic few-shot injection.
4
+
5
+ v3.0 changes:
6
+ - System prompt rewritten with INVESTIGATION PLAYBOOKS
7
+ - "esta funcionando X?" is ALWAYS tool_call, never text
8
+ - Better Spanish + investigation pattern detection
9
+ - Knowledge about the local system (Pi, HA, Docker)
10
+ """
11
+
12
+ import re
13
+ from typing import Dict, List, Optional, Tuple
14
+ from .schemas import TOOL_DEFINITIONS, tools_to_prompt_text
15
+
16
+
17
+ # ============================================================
18
+ # INTENT CLASSIFICATION (zero-cost, pure regex/heuristics)
19
+ # ============================================================
20
+
21
+ TOOL_INDICATORS = [
22
+ # Direct action verbs
23
+ r'\b(run|execute|check|show|list|find|search|scan|install|deploy|restart|stop|start|kill|create|delete|remove|move|copy|rename)\b',
24
+ # System queries (these ALWAYS need real data)
25
+ r'\b(disk|memory|cpu|ram|process|port|network|service|docker|container|pid|uptime|load)\b',
26
+ r'\b(file|directory|folder|path|config|log|permission)\b',
27
+ # Specific commands
28
+ r'\b(ls|cat|grep|find|ps|df|du|top|htop|netstat|ss|curl|wget|ping|nmap|apt|pip|npm|systemctl)\b',
29
+ # Target references
30
+ r'\b(raspi|raspberry|gcp|router|server|vm|machine|host)\b',
31
+ # File operations
32
+ r'\b(read|write|edit|modify|update|append|save)\s+(the\s+)?file\b',
33
+ r'/[a-z]+(/[a-z0-9_./-]+)+', # Unix paths
34
+ # Diagnostic
35
+ r'\b(diagnos|health|status)\b',
36
+ # STATE QUERIES -- these ALWAYS need tools (never answer from "knowledge")
37
+ r'\b(running|installed|active|enabled|listening|open|up|down|alive|funciona|funcionando|corriendo|prendid[oa]|encendid[oa]|andando)\b',
38
+ r'\b(version|cuant[oa]s?|cuales|which|where|donde|what.*running|que.*corre|que.*hay)\b',
39
+ # Investigation patterns
40
+ r'\b(check|verify|test|try|connect|access|reach|investigate)\b',
41
+ # Home automation
42
+ r'\b(home\s*assistant|hass|hassio|zigbee|zwave|mqtt|tasmota|esphome|automation|sensor|switch|light)\b',
43
+ ]
44
+
45
+ CODE_INDICATORS = [
46
+ r'\b(write|create|build|implement|code|script|function|class|program)\b.*\b(in|using|with)\s+(python|bash|javascript|rust|go|c\+\+|java|typescript)\b',
47
+ r'\b(write|create|build)\b.*\b(python|bash|javascript|rust|go|java|typescript)\b.*\b(script|function|program|code)\b',
48
+ r'\b(algorithm|data\s+structure|design\s+pattern|architecture)\b',
49
+ r'\b(bug|error|fix|debug|traceback|exception|syntax)\b.*\b(code|script|program|file)\b',
50
+ ]
51
+
52
+ TEXT_INDICATORS = [
53
+ # Greetings (EN + ES)
54
+ r'^(hi|hello|hey|hola|que\s+tal|buenos?\s+d[ií]as?|buenas?\s+(?:tardes?|noches?))\s*[.!?]*$',
55
+ r'^(thanks?|thank\s+you|gracias|thx|de\s+nada)\s*[.!?]*$',
56
+ # Conversational follow-ups
57
+ r'^(como\s+(va|estas?|andas?|te\s+va)|how\s+are\s+you|how\'?s\s+it\s+going|what\'?s\s+up)\s*[.!?]*$',
58
+ r'^(todo\s+bien|bien\s+y\s+tu|genial|cool|nice|great|awesome)\s*[.!?]*$',
59
+ r'^(que\s+onda|que\s+hay|que\s+pasa|que\s+cuentas)\s*[.!?]*$',
60
+ # Short affirmations
61
+ r'^(yes|no|ok|okay|sure|maybe|si|no|dale|genial|perfecto|claro|obvio|exacto|listo)\s*[.!?]*$',
62
+ # Identity / meta questions
63
+ r'\b(who\s+are\s+you|what\s+can\s+you|tell\s+me\s+about\s+yourself)\b',
64
+ r'\b(opinion|think\s+about|recommend|suggest|advice)\b',
65
+ # Explicit CONCEPT explanations (NOT status checks)
66
+ r'^(explain|describe|define)\s+(the\s+)?(concept|theory|protocol|idea|principle|difference)\b',
67
+ r'^what\s+is\s+(a\s+)?(concept|theory|protocol|principle|idea|technique|method)\b',
68
+ ]
69
+
70
+
71
+ class IntentType:
72
+ TOOL_CALL = "tool_call"
73
+ CODE = "code"
74
+ TEXT = "text"
75
+ DECISION = "decision"
76
+
77
+
78
+ def classify_intent(user_message: str, conversation_history: list = None) -> Tuple[str, float]:
79
+ """
80
+ Classify user intent. Returns (intent_type, confidence).
81
+
82
+ v3.0 RULE: Any question about whether something IS RUNNING, IS INSTALLED,
83
+ IS WORKING, etc. is ALWAYS a tool_call. We can never answer these from
84
+ "knowledge" -- we must CHECK THE ACTUAL SYSTEM.
85
+ """
86
+ msg = user_message.lower().strip()
87
+
88
+ # Investigation prompts from the engine itself
89
+ if msg.startswith("investigation"):
90
+ return IntentType.DECISION, 0.9
91
+ if "gathered enough" in msg or "final answer" in msg or "what i found" in msg:
92
+ return IntentType.DECISION, 0.8
93
+
94
+ # Empty or very short
95
+ if len(msg) < 3:
96
+ return IntentType.TEXT, 0.9
97
+
98
+ # ===== STRONG TEXT SIGNALS (pure greetings, thanks, etc.) =====
99
+ # These are ONLY pure conversational. "hola" alone = text.
100
+ # But "hola, esta funcionando HA?" = tool_call (caught by tool indicators below)
101
+ for pattern in TEXT_INDICATORS:
102
+ if re.match(pattern, msg, re.IGNORECASE):
103
+ return IntentType.TEXT, 0.95
104
+
105
+ # ===== STATE QUESTIONS ARE ALWAYS TOOL CALLS =====
106
+ # "esta funcionando X?" / "is X running?" / "funciona el HA?" / etc.
107
+ state_patterns = [
108
+ r'\b(esta|está)\s+(funcionando|corriendo|andando|prendid[oa]|activ[oa]|instalad[oa])\b',
109
+ r'\b(funciona|corre|anda)\s+(el|la|lo|los|las)?\s*\w+',
110
+ r'\bis\s+\w+\s+(running|installed|active|working|up|down|alive|enabled|listening)\b',
111
+ r'\b(running|installed|active|working|funciona)\b.*\?',
112
+ r'\bcheck\s+(if|whether|that)\b',
113
+ # Port/network state queries
114
+ r'\b(puerto|port)\s+\d+',
115
+ r'\bhay\s+algo\s+(en|corriendo|escuchando)',
116
+ r'\b(que|qué)\s+(hay|corre|esta)\s+(en|corriendo)',
117
+ r'\b(listening|escuchando)\s+(on|en)\b',
118
+ # Who/what is queries about system state
119
+ r'\b(quien|quién|who)\s+(esta|está|is)\s+(logueado|logged|conectado|connected)\b',
120
+ r'\b(puertos?|ports?)\s+(estan|están|are)\s+(escuchando|listening|abiertos?|open)\b',
121
+ r'\b(vulnerabilidad|vulnerability|vulnerabilidades|vulnerabilities)\b',
122
+ r'\b(audit|auditoria|auditoría)\b.*\b(seguridad|security)\b',
123
+ r'\b(seguridad|security)\b.*\b(audit|auditoria|auditoría)\b',
124
+ r'\b(logueado|logged\s+in|conectado)\b',
125
+ # Vision / robot camera queries
126
+ r'\b(que\s+ves|qué\s+ves|what\s+do\s+you\s+see|what\s+can\s+you\s+see)\b',
127
+ r'\b(mira|look|fijate|fíjate|observa|observe)\b',
128
+ r'\b(hay\s+alguien|is\s+anyone|is\s+there\s+anyone|who\s+is\s+there|quien\s+hay)\b',
129
+ r'\b(muestra|show\s+me|take\s+a\s+(photo|picture|snapshot)|saca\s+una?\s+foto)\b',
130
+ r'\b(camara|cámara|camera)\b',
131
+ r'\b(ves\s+algo|see\s+anything|what.s\s+around)\b',
132
+ r'\b(describe\s+(the\s+)?room|describe\s+(la\s+)?habitaci[oó]n|describe\s+(el\s+)?entorno)\b',
133
+ ]
134
+ for pattern in state_patterns:
135
+ if re.search(pattern, msg, re.IGNORECASE):
136
+ return IntentType.TOOL_CALL, 0.95
137
+
138
+ # ===== "what is X?" knowledge questions (concept, not status) =====
139
+ # "what is docker?" = text. "what is running on port 8123?" = tool.
140
+ if re.match(r'^(what\s+is|what\s+are|explain|describe|define|que\s+es|qué\s+es)\s+', msg):
141
+ # Check if it's asking about runtime state
142
+ if re.search(r'\b(running|on\s+port|installed|active|listening|my\s+server|this\s+machine)\b', msg):
143
+ return IntentType.TOOL_CALL, 0.9
144
+ return IntentType.TEXT, 0.85
145
+
146
+ # Score each category
147
+ tool_score = 0
148
+ code_score = 0
149
+ text_score = 0
150
+
151
+ for pattern in TOOL_INDICATORS:
152
+ if re.search(pattern, msg, re.IGNORECASE):
153
+ tool_score += 1
154
+
155
+ for pattern in CODE_INDICATORS:
156
+ if re.search(pattern, msg, re.IGNORECASE):
157
+ code_score += 1
158
+
159
+ for pattern in TEXT_INDICATORS:
160
+ if re.search(pattern, msg, re.IGNORECASE):
161
+ text_score += 1
162
+
163
+ # Boost: previous turn context
164
+ if conversation_history:
165
+ last = conversation_history[-1] if conversation_history else None
166
+ if last and last.get("role") == "tool":
167
+ tool_score += 2
168
+ if last and last.get("role") == "assistant":
169
+ content = last.get("content", "")
170
+ if "I ran:" in content or "Result of" in content:
171
+ tool_score += 1
172
+
173
+ # Determine winner
174
+ scores = {
175
+ IntentType.TOOL_CALL: tool_score,
176
+ IntentType.CODE: code_score,
177
+ IntentType.TEXT: text_score,
178
+ }
179
+
180
+ max_score = max(scores.values())
181
+ total = sum(scores.values()) or 1
182
+
183
+ if max_score == 0:
184
+ if len(msg) < 30:
185
+ return IntentType.TEXT, 0.6
186
+ # Longer messages with no clear indicators -> decision
187
+ return IntentType.DECISION, 0.4
188
+
189
+ winner = max(scores, key=scores.get)
190
+ confidence = max_score / (total + 1)
191
+
192
+ # Close scores -> decision
193
+ if tool_score > 0 and text_score > 0 and abs(tool_score - text_score) <= 1:
194
+ return IntentType.DECISION, 0.5
195
+
196
+ return winner, min(confidence, 0.95)
197
+
198
+
199
+ # ============================================================
200
+ # SYSTEM PROMPT (v3.0 -- with investigation playbooks)
201
+ # ============================================================
202
+
203
+ SYSTEM_PROMPT_BASE = """You are TokioAI -- an elite cybersecurity and DevOps AI agent.
204
+ You execute commands, fix bugs, deploy infrastructure, and solve problems.
205
+ You have full access to the user's systems via tools.
206
+
207
+ ## CRITICAL RULES -- NEVER VIOLATE THESE
208
+
209
+ 1. NEVER simulate or fake tool execution in your text responses.
210
+ WRONG: writing "```bash\ndf -h\n```\nResult: ..." in your text.
211
+ RIGHT: call the execute_local tool with {"command": "df -h"}.
212
+ The engine will execute the tool and give you the real result.
213
+
214
+ 2. When you need system information, ALWAYS use a tool call.
215
+ NEVER make up or guess command output. NEVER write fake results.
216
+
217
+ 3. Your text responses should contain ONLY your analysis and conclusions.
218
+ NEVER include "I ran:", "Result of tool:", or "```bash" blocks that
219
+ simulate command execution.
220
+
221
+ 4. If you already have real tool results in the conversation, summarize
222
+ them in plain language. Do NOT echo back raw tool output format.
223
+
224
+ ## ENVIRONMENT
225
+ You are running on a Raspberry Pi 5 (your-device, Debian/aarch64).
226
+ Key services:
227
+ - Home Assistant: runs as Docker container "homeassistant" (NOT systemd). Check with: docker ps | grep home
228
+ - Docker: available. Most services run as containers.
229
+ - Ollama: runs locally, accessed via SSH tunnel to L4 GPU on port 8200
230
+ - TokioAI Engine: this server (port 8080)
231
+ - SSH tunnels: several active to GCP services
232
+ When checking if something is running, ALWAYS check both systemd AND docker.
233
+
234
+ ## ROBOT VISION (robot_vision tool)
235
+ You have a camera on the PiCar-X robot. Use the robot_vision tool to SEE what the robot sees.
236
+ - Use it when the user asks: "que ves?", "what do you see?", "look around", "describe the room", "hay alguien?", "who is there?"
237
+ - The tool captures a live photo from the robot's camera and analyzes it with a vision AI model (moondream)
238
+ - You can pass a specific question: robot_vision(question="how many people are in the room?")
239
+ - Or leave question empty for a general description
240
+ - The robot camera is at a low angle (ground level), so perspective may be different from human height
241
+ - Response includes image size, capture time, analysis time, and the visual description
242
+ - ALWAYS use robot_vision for visual questions. NEVER say "I can't see" -- you CAN see through the robot.
243
+
244
+ ## HOME ASSISTANT VERIFICATION (CRITICAL)
245
+ After calling ANY HA service (media, lights, switches, etc.), you MUST:
246
+ 1. Execute the service call
247
+ 2. Wait 3 seconds (sleep 3)
248
+ 3. GET the entity state to confirm it ACTUALLY changed
249
+ 4. Report what is ACTUALLY happening, not what you hoped
250
+
251
+ HA auth: HA_TOKEN=$(cat ~/.ha_token) && curl -s -H "Authorization: Bearer $HA_TOKEN" ...
252
+ Base URL: http://localhost:8123
253
+
254
+ For music changes specifically:
255
+ After play_media, run: sleep 3 && HA_TOKEN=$(cat ~/.ha_token) && curl -s -H "Authorization: Bearer $HA_TOKEN" http://localhost:8123/api/states/media_player.YOUR_SPEAKER | python3 -c "import sys,json;s=json.load(sys.stdin);a=s['attributes'];print(f'State: {s[\"state\"]}, Title: {a.get(\"media_title\",\"?\")}, Artist: {a.get(\"media_artist\",\"?\")}')"
256
+ If title/artist didn't change or state is not "playing", tell the user it FAILED.
257
+
258
+ A 200 OK from the API only means HA received the request, NOT that it worked. ALWAYS verify.
259
+
260
+ ## HOME ASSISTANT COOKBOOK (execute via execute_local tool)
261
+ Auth pattern: HA_TOKEN=$(cat ~/.ha_token) && curl -s -H "Authorization: Bearer $HA_TOKEN" ...
262
+ Base: http://localhost:8123
263
+
264
+ MEDIA (YOUR_SPEAKER (e.g. Amazon Echo), entity: media_player.YOUR_SPEAKER):
265
+ - What's playing: GET /api/states/media_player.YOUR_SPEAKER (check media_title, media_artist, state)
266
+ - CHANGE MUSIC: POST /api/services/media_player/play_media with {"entity_id":"media_player.YOUR_SPEAKER","media_content_type":"music","media_content_id":"QUERY"} (e.g. "blues", "rock", "Beatles")
267
+ - Volume: POST /api/services/media_player/volume_set with {"entity_id":"media_player.YOUR_SPEAKER","volume_level":0.5}
268
+ - Play/Pause: POST /api/services/media_player/media_play_pause with {"entity_id":"media_player.YOUR_SPEAKER"}
269
+ - Next/Prev: POST /api/services/media_player/media_next_track or media_previous_track
270
+ - TTS (Alexa speaks): POST /api/services/notify/alexa_media_YOUR_SPEAKER with {"message":"TEXT","data":{"type":"tts"}}
271
+ - Announce: POST /api/services/notify/alexa_media_YOUR_SPEAKER with {"message":"TEXT","data":{"type":"announce"}}
272
+ NEVER use notify to change music. notify is ONLY for TTS/announcements.
273
+
274
+ LIGHTS: light.living_room, light.bedroom, light.kitchen (customize for your setup)
275
+ - Turn on/off: POST /api/services/light/turn_on or turn_off with {"entity_id":"light.ENTITY"}
276
+ - Brightness: turn_on with {"entity_id":"light.ENTITY","brightness":128} (0-255)
277
+ - Color: turn_on with {"entity_id":"light.ENTITY","rgb_color":[255,0,0]}
278
+
279
+ SWITCHES: switch.YOUR_SWITCH (customize)
280
+ SENSORS: sensor.YOUR_TEMPERATURE, weather.forecast_home (customize)
281
+ VACUUM: vacuum.YOUR_VACUUM (start, stop, return_to_base)
282
+
283
+ ## CORE RULES
284
+ - Be DIRECT. Act first, explain after.
285
+ - NEVER give up. If something fails, try alternatives.
286
+ - When investigating: diagnose -> try alternatives -> root cause -> answer.
287
+ - NEVER output passwords, API keys, or tokens in plain text.
288
+ - Be concise but thorough.
289
+
290
+ ## INVESTIGATION BEHAVIOR
291
+ You are an INVESTIGATOR. When asked about system state, security, or any technical question:
292
+ 1. ALWAYS check the actual system. NEVER speculate or say "you should check X" -- CHECK IT YOURSELF.
293
+ 2. If the first command doesn't give the full picture, run MORE commands automatically.
294
+ 3. For security questions: check configs, ports, services, permissions, logs. Don't stop at 1 command.
295
+ 4. For "is X running?": check systemctl, docker, ps, ports. Try multiple approaches.
296
+ 5. NEVER ask "do you want me to check more?" -- just CHECK IT. You are the expert.
297
+ 6. NEVER say "you should review X" -- YOU review it and report findings.
298
+ 7. Think like a senior security engineer: be thorough, check multiple angles, report specifics.
299
+ 8. Always use TOOL CALLS -- never simulate commands in your text response.
300
+
301
+ ## ANSWER FORMAT
302
+ When you've gathered enough information:
303
+ - State what you found CLEARLY
304
+ - If something is running: say what it is, where, which port, how it was started
305
+ - If something is NOT available: say what you tried and why you're confident it's not there
306
+ - NEVER give generic "recommendations" -- give FACTS based on what you actually found
307
+ - NEVER guess or assume what service runs on a port -- only report what the tool output ACTUALLY shows
308
+ - If ss/netstat shows a process name, report THAT name. Do NOT map port numbers to "well-known" services unless the output confirms it
309
+ - Report EXACTLY what the command output says. If a port shows "ssh" as the process, say it's an SSH tunnel, not "Vault" or "Redis" etc."""
310
+
311
+ PROMPT_TOOL_CALL = """
312
+ The user wants you to DO something or CHECK something. Use the appropriate tool.
313
+ - For shell commands: use execute_local
314
+ - For file reading: use read_file
315
+ - For searching: use search_files
316
+ - For diagnostics: use diagnose
317
+ - For seeing/looking/camera/vision: use robot_vision
318
+
319
+ Call the tool IMMEDIATELY. Do not explain what you will do -- just do it.
320
+ If you need multiple pieces of information, call the MOST IMPORTANT tool first.
321
+ If a tool fails, you'll get a chance to try again with a different approach.
322
+ NEVER ask the user "want me to check more?" -- you ARE the expert, so CHECK IT YOURSELF.
323
+ NEVER recommend actions without doing them first. ACT, then REPORT."""
324
+
325
+ PROMPT_CODE = """
326
+ The user wants code or technical explanation.
327
+ - If they want to CREATE a file: use write_file with the complete code
328
+ - If they want to MODIFY code: use read_file first, then edit_file
329
+ - If they want EXPLANATION: respond directly with concise, expert-level text
330
+ - Show code, not descriptions of code"""
331
+
332
+ PROMPT_TEXT = """
333
+ The user is having a conversation or asking a knowledge question.
334
+ Respond directly WITHOUT using any tools.
335
+ Be concise and helpful. Answer from your knowledge.
336
+ Do NOT call any tools for pure knowledge questions.
337
+ Do NOT simulate tool execution in your text. Just answer naturally."""
338
+
339
+ PROMPT_DECISION = """
340
+ You must decide: does this need a tool, or can you answer from knowledge?
341
+ Guidelines:
342
+ - "What is X?" (concept) = text answer
343
+ - "Is X running?" / "Check X" / "Show X" = USE A TOOL (requires system state)
344
+ - Anything about the current system state = USE A TOOL
345
+ - Anything you can answer from pure knowledge = text answer
346
+ - When in doubt, USE A TOOL -- it's better to check than to guess"""
347
+
348
+ INTENT_PROMPTS = {
349
+ IntentType.TOOL_CALL: PROMPT_TOOL_CALL,
350
+ IntentType.CODE: PROMPT_CODE,
351
+ IntentType.TEXT: PROMPT_TEXT,
352
+ IntentType.DECISION: PROMPT_DECISION,
353
+ }
354
+
355
+
356
+ def build_system_prompt(
357
+ intent: str,
358
+ few_shot_examples: List[Dict] = None,
359
+ extra_context: str = None,
360
+ ) -> str:
361
+ """Build the optimal system prompt for a given intent."""
362
+ parts = [SYSTEM_PROMPT_BASE]
363
+
364
+ # Add intent-specific instructions
365
+ parts.append(INTENT_PROMPTS.get(intent, PROMPT_DECISION))
366
+
367
+ # Add tool descriptions (for tool/decision intents)
368
+ if intent in (IntentType.TOOL_CALL, IntentType.DECISION):
369
+ parts.append("\n" + tools_to_prompt_text())
370
+
371
+ # Add few-shot examples
372
+ if few_shot_examples:
373
+ parts.append("\n--- RELEVANT EXAMPLES ---")
374
+ for ex in few_shot_examples[:3]:
375
+ parts.append(f"User: {ex.get('input', '')}")
376
+ parts.append(f"Response: {ex.get('output', '')}")
377
+ parts.append("--- END EXAMPLES ---")
378
+
379
+ if extra_context:
380
+ parts.append(f"\n{extra_context}")
381
+
382
+ return "\n".join(parts)
engine/schemas.py ADDED
@@ -0,0 +1,497 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ TokioAI Engine -- Tool Schemas & JSON Validation
3
+ Defines all available tools, their parameter schemas, and validation logic.
4
+ """
5
+
6
+ import json
7
+ import re
8
+ from typing import Any, Dict, List, Optional
9
+
10
+ # ============================================================
11
+ # TOOL DEFINITIONS (OpenAI function-calling format)
12
+ # ============================================================
13
+
14
+ TOOL_DEFINITIONS = [
15
+ {
16
+ "type": "function",
17
+ "function": {
18
+ "name": "execute_local",
19
+ "description": "Execute a shell command on the local machine. Use for system operations: install packages, check processes, network scans, file operations.",
20
+ "parameters": {
21
+ "type": "object",
22
+ "properties": {
23
+ "command": {
24
+ "type": "string",
25
+ "description": "Shell command to execute (bash)"
26
+ },
27
+ "timeout": {
28
+ "type": "integer",
29
+ "description": "Timeout in seconds (default 30, max 120)"
30
+ }
31
+ },
32
+ "required": ["command"]
33
+ }
34
+ }
35
+ },
36
+ {
37
+ "type": "function",
38
+ "function": {
39
+ "name": "execute_raspi",
40
+ "description": "Execute a command on the Raspberry Pi via SSH.",
41
+ "parameters": {
42
+ "type": "object",
43
+ "properties": {
44
+ "command": {
45
+ "type": "string",
46
+ "description": "Command to run on Raspi"
47
+ },
48
+ "timeout": {
49
+ "type": "integer",
50
+ "description": "Timeout in seconds (default 30, max 120)"
51
+ }
52
+ },
53
+ "required": ["command"]
54
+ }
55
+ }
56
+ },
57
+ {
58
+ "type": "function",
59
+ "function": {
60
+ "name": "execute_gcp",
61
+ "description": "Execute a command on the GCP VM via SSH.",
62
+ "parameters": {
63
+ "type": "object",
64
+ "properties": {
65
+ "command": {
66
+ "type": "string",
67
+ "description": "Command to run on GCP"
68
+ },
69
+ "timeout": {
70
+ "type": "integer",
71
+ "description": "Timeout in seconds (default 30, max 120)"
72
+ }
73
+ },
74
+ "required": ["command"]
75
+ }
76
+ }
77
+ },
78
+ {
79
+ "type": "function",
80
+ "function": {
81
+ "name": "execute_router",
82
+ "description": "Execute a command on the router via SSH (root).",
83
+ "parameters": {
84
+ "type": "object",
85
+ "properties": {
86
+ "command": {
87
+ "type": "string",
88
+ "description": "Command to run on router as root"
89
+ }
90
+ },
91
+ "required": ["command"]
92
+ }
93
+ }
94
+ },
95
+ {
96
+ "type": "function",
97
+ "function": {
98
+ "name": "read_file",
99
+ "description": "Read a file from local or remote systems. Use raspi:/path or gcp:/path for remote.",
100
+ "parameters": {
101
+ "type": "object",
102
+ "properties": {
103
+ "path": {
104
+ "type": "string",
105
+ "description": "File path. Use raspi:/path or gcp:/path for remote"
106
+ },
107
+ "lines": {
108
+ "type": "integer",
109
+ "description": "Max lines to read (default: all)"
110
+ }
111
+ },
112
+ "required": ["path"]
113
+ }
114
+ }
115
+ },
116
+ {
117
+ "type": "function",
118
+ "function": {
119
+ "name": "write_file",
120
+ "description": "Write content to a file. Prefix with raspi: or gcp: for remote.",
121
+ "parameters": {
122
+ "type": "object",
123
+ "properties": {
124
+ "path": {
125
+ "type": "string",
126
+ "description": "File path. Use raspi:/path or gcp:/path for remote"
127
+ },
128
+ "content": {
129
+ "type": "string",
130
+ "description": "File content to write"
131
+ }
132
+ },
133
+ "required": ["path", "content"]
134
+ }
135
+ }
136
+ },
137
+ {
138
+ "type": "function",
139
+ "function": {
140
+ "name": "edit_file",
141
+ "description": "Replace a specific string in a file. old_text must be unique.",
142
+ "parameters": {
143
+ "type": "object",
144
+ "properties": {
145
+ "path": {
146
+ "type": "string",
147
+ "description": "File path. Use raspi:/path or gcp:/path for remote"
148
+ },
149
+ "old_text": {
150
+ "type": "string",
151
+ "description": "Exact text to find (must be unique)"
152
+ },
153
+ "new_text": {
154
+ "type": "string",
155
+ "description": "Replacement text"
156
+ }
157
+ },
158
+ "required": ["path", "old_text", "new_text"]
159
+ }
160
+ }
161
+ },
162
+ {
163
+ "type": "function",
164
+ "function": {
165
+ "name": "search_files",
166
+ "description": "Search for text patterns across files. Like grep -rn.",
167
+ "parameters": {
168
+ "type": "object",
169
+ "properties": {
170
+ "pattern": {
171
+ "type": "string",
172
+ "description": "Text or regex pattern to search"
173
+ },
174
+ "path": {
175
+ "type": "string",
176
+ "description": "Directory to search. Use raspi:/path or gcp:/path"
177
+ },
178
+ "glob": {
179
+ "type": "string",
180
+ "description": "File filter glob (e.g., '*.py')"
181
+ }
182
+ },
183
+ "required": ["pattern"]
184
+ }
185
+ }
186
+ },
187
+ {
188
+ "type": "function",
189
+ "function": {
190
+ "name": "robot_vision",
191
+ "description": "Capture a photo from the robot's camera and analyze it with a vision AI model. Use when user asks what the robot sees, to look around, identify objects, or describe the environment.",
192
+ "parameters": {
193
+ "type": "object",
194
+ "properties": {
195
+ "question": {
196
+ "type": "string",
197
+ "description": "What to look for or ask about the image (e.g. 'how many people?', 'what objects are on the table?', 'describe the room'). Default: general description."
198
+ }
199
+ },
200
+ "required": []
201
+ }
202
+ }
203
+ },
204
+ {
205
+ "type": "function",
206
+ "function": {
207
+ "name": "diagnose",
208
+ "description": "Run system health diagnostics on local, raspi, gcp, or all targets.",
209
+ "parameters": {
210
+ "type": "object",
211
+ "properties": {
212
+ "target": {
213
+ "type": "string",
214
+ "enum": ["local", "raspi", "gcp", "router", "all"],
215
+ "description": "Target: local, raspi, gcp, router, all"
216
+ }
217
+ },
218
+ "required": ["target"]
219
+ }
220
+ }
221
+ },
222
+ {
223
+ "type": "function",
224
+ "function": {
225
+ "name": "memory",
226
+ "description": "Persistent memory. Actions: read, write, append, clear.",
227
+ "parameters": {
228
+ "type": "object",
229
+ "properties": {
230
+ "action": {
231
+ "type": "string",
232
+ "enum": ["read", "write", "append", "clear"],
233
+ "description": "Action to perform"
234
+ },
235
+ "content": {
236
+ "type": "string",
237
+ "description": "Content to write/append"
238
+ }
239
+ },
240
+ "required": ["action"]
241
+ }
242
+ }
243
+ },
244
+ {
245
+ "type": "function",
246
+ "function": {
247
+ "name": "task",
248
+ "description": "Persistent task tracker. Actions: add, update, list, remove, clear_done.",
249
+ "parameters": {
250
+ "type": "object",
251
+ "properties": {
252
+ "action": {
253
+ "type": "string",
254
+ "enum": ["add", "update", "list", "remove", "clear_done"],
255
+ "description": "Action to perform"
256
+ },
257
+ "task": {
258
+ "type": "string",
259
+ "description": "Task description (for add)"
260
+ },
261
+ "id": {
262
+ "type": "integer",
263
+ "description": "Task ID (for update/remove)"
264
+ },
265
+ "status": {
266
+ "type": "string",
267
+ "enum": ["pending", "in_progress", "done", "blocked"],
268
+ "description": "New status"
269
+ }
270
+ },
271
+ "required": ["action"]
272
+ }
273
+ }
274
+ }
275
+ ]
276
+
277
+ # ============================================================
278
+ # TOOL NAME INDEX (for quick lookup)
279
+ # ============================================================
280
+
281
+ VALID_TOOL_NAMES = {t["function"]["name"] for t in TOOL_DEFINITIONS}
282
+
283
+ TOOL_SCHEMAS = {
284
+ t["function"]["name"]: t["function"]["parameters"]
285
+ for t in TOOL_DEFINITIONS
286
+ }
287
+
288
+ # ============================================================
289
+ # VALIDATION
290
+ # ============================================================
291
+
292
+ def validate_tool_call(name: str, arguments: Dict[str, Any]) -> tuple[bool, str]:
293
+ """
294
+ Validate a tool call against its schema.
295
+ Returns (is_valid, error_message).
296
+ """
297
+ if name not in VALID_TOOL_NAMES:
298
+ return False, f"Unknown tool '{name}'. Valid: {', '.join(sorted(VALID_TOOL_NAMES))}"
299
+
300
+ schema = TOOL_SCHEMAS[name]
301
+ properties = schema.get("properties", {})
302
+ required = schema.get("required", [])
303
+
304
+ # Check required params
305
+ for req in required:
306
+ if req not in arguments:
307
+ return False, f"Missing required parameter '{req}' for tool '{name}'"
308
+
309
+ # Check param types
310
+ for param_name, param_value in arguments.items():
311
+ if param_name not in properties:
312
+ continue # Extra params are OK (flexible)
313
+ expected_type = properties[param_name].get("type", "string")
314
+ if expected_type == "string" and not isinstance(param_value, str):
315
+ return False, f"Parameter '{param_name}' must be string, got {type(param_value).__name__}"
316
+ if expected_type == "integer" and not isinstance(param_value, int):
317
+ return False, f"Parameter '{param_name}' must be integer, got {type(param_value).__name__}"
318
+
319
+ # Check enum constraints
320
+ for param_name, param_value in arguments.items():
321
+ if param_name in properties:
322
+ enum_vals = properties[param_name].get("enum")
323
+ if enum_vals and param_value not in enum_vals:
324
+ return False, f"Parameter '{param_name}' must be one of {enum_vals}, got '{param_value}'"
325
+
326
+ return True, ""
327
+
328
+
329
+ def extract_tool_calls_from_text(text: str) -> List[Dict]:
330
+ """
331
+ Extract tool calls from model output text.
332
+ Supports multiple formats:
333
+ 1. OpenAI-style JSON: {"name": "...", "arguments": {...}}
334
+ 2. Markdown code blocks with JSON
335
+ 3. Raw function-call-like patterns
336
+
337
+ Returns list of {"name": str, "arguments": dict}
338
+ """
339
+ calls = []
340
+
341
+ # Strategy 1: Find JSON objects with "name" and "arguments" keys
342
+ # Handle both single objects and arrays
343
+ json_patterns = [
344
+ # Array of tool calls
345
+ r'\[\s*\{[^{]*?"name"\s*:\s*"[^"]+?".*?\}\s*\]',
346
+ # Single tool call object
347
+ r'\{\s*"(?:name|function)"\s*:\s*"[^"]+?".*?"arguments"\s*:\s*\{.*?\}\s*\}',
348
+ ]
349
+
350
+ for pattern in json_patterns:
351
+ for match in re.finditer(pattern, text, re.DOTALL):
352
+ try:
353
+ data = json.loads(match.group())
354
+ if isinstance(data, list):
355
+ for item in data:
356
+ call = _normalize_tool_call(item)
357
+ if call:
358
+ calls.append(call)
359
+ elif isinstance(data, dict):
360
+ call = _normalize_tool_call(data)
361
+ if call:
362
+ calls.append(call)
363
+ except json.JSONDecodeError:
364
+ continue
365
+
366
+ if calls:
367
+ return calls
368
+
369
+ # Strategy 1.5: Anthropic/XML-style tool calls
370
+ # Format: <function=name><parameter=key>value</parameter></function>
371
+ # Also: <tool_call>{"name": "...", "arguments": {...}}</tool_call>
372
+ xml_patterns = [
373
+ # <function=name><parameter=key>value</parameter>...
374
+ r'<function=(\w+)>(.*?)</function>',
375
+ # <tool_call>JSON</tool_call>
376
+ r'<tool_call>\s*(.*?)\s*</tool_call>',
377
+ ]
378
+
379
+ for pattern in xml_patterns:
380
+ for match in re.finditer(pattern, text, re.DOTALL):
381
+ groups = match.groups()
382
+ if len(groups) == 2 and not groups[0].startswith('{'):
383
+ # <function=name><parameter=key>value format
384
+ func_name = groups[0]
385
+ params_text = groups[1]
386
+ if func_name in VALID_TOOL_NAMES:
387
+ # Extract parameters
388
+ param_dict = {}
389
+ for pm in re.finditer(r'<parameter=(\w+)>\s*(.*?)\s*</parameter>', params_text, re.DOTALL):
390
+ param_dict[pm.group(1)] = pm.group(2).strip()
391
+ calls.append({"name": func_name, "arguments": param_dict})
392
+ elif len(groups) >= 1:
393
+ # <tool_call>JSON format
394
+ try:
395
+ data = json.loads(groups[-1])
396
+ call = _normalize_tool_call(data)
397
+ if call:
398
+ calls.append(call)
399
+ except (json.JSONDecodeError, TypeError):
400
+ pass
401
+
402
+ if calls:
403
+ return calls
404
+
405
+ # Strategy 2: Look for ```json code blocks
406
+ code_blocks = re.findall(r'```(?:json)?\s*\n(.*?)\n\s*```', text, re.DOTALL)
407
+ for block in code_blocks:
408
+ try:
409
+ data = json.loads(block.strip())
410
+ if isinstance(data, list):
411
+ for item in data:
412
+ call = _normalize_tool_call(item)
413
+ if call:
414
+ calls.append(call)
415
+ elif isinstance(data, dict):
416
+ call = _normalize_tool_call(data)
417
+ if call:
418
+ calls.append(call)
419
+ except json.JSONDecodeError:
420
+ continue
421
+
422
+ if calls:
423
+ return calls
424
+
425
+ # Strategy 3: Brute-force -- find ANY valid JSON object in the text
426
+ brace_depth = 0
427
+ start = None
428
+ for i, ch in enumerate(text):
429
+ if ch == '{':
430
+ if brace_depth == 0:
431
+ start = i
432
+ brace_depth += 1
433
+ elif ch == '}':
434
+ brace_depth -= 1
435
+ if brace_depth == 0 and start is not None:
436
+ candidate = text[start:i+1]
437
+ try:
438
+ data = json.loads(candidate)
439
+ call = _normalize_tool_call(data)
440
+ if call:
441
+ calls.append(call)
442
+ except json.JSONDecodeError:
443
+ pass
444
+ start = None
445
+
446
+ return calls
447
+
448
+
449
+ def _normalize_tool_call(data: dict) -> Optional[Dict]:
450
+ """Normalize various tool call formats to {"name": str, "arguments": dict}."""
451
+ if not isinstance(data, dict):
452
+ return None
453
+
454
+ # Get name -- handle multiple formats
455
+ name = data.get("name")
456
+ if not name:
457
+ func = data.get("function")
458
+ if isinstance(func, dict):
459
+ name = func.get("name")
460
+ elif isinstance(func, str):
461
+ name = func
462
+
463
+ args = data.get("arguments") or data.get("parameters") or data.get("params") or data.get("args")
464
+ if isinstance(args, str):
465
+ try:
466
+ args = json.loads(args)
467
+ except json.JSONDecodeError:
468
+ return None
469
+
470
+ if not name or not isinstance(name, str):
471
+ return None
472
+ if args is None:
473
+ args = {}
474
+ if not isinstance(args, dict):
475
+ return None
476
+
477
+ # Only accept valid tool names
478
+ if name not in VALID_TOOL_NAMES:
479
+ return None
480
+
481
+ return {"name": name, "arguments": args}
482
+
483
+
484
+ def tools_to_prompt_text() -> str:
485
+ """Generate a concise tool description for injection into prompts."""
486
+ lines = ["Available tools:"]
487
+ for tool in TOOL_DEFINITIONS:
488
+ f = tool["function"]
489
+ params = f["parameters"].get("properties", {})
490
+ required = f["parameters"].get("required", [])
491
+ param_strs = []
492
+ for pname, pinfo in params.items():
493
+ req_marker = "*" if pname in required else ""
494
+ param_strs.append(f"{pname}{req_marker}: {pinfo.get('type', 'string')}")
495
+ params_text = ", ".join(param_strs) if param_strs else "none"
496
+ lines.append(f" - {f['name']}({params_text}): {f['description'][:80]}")
497
+ return "\n".join(lines)
engine/server.py ADDED
@@ -0,0 +1,864 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ TokioAI Engine v3.0 -- OpenAI-Compatible API Server
3
+ SERVER-SIDE MULTI-STEP INVESTIGATION (the Opus architecture)
4
+
5
+ The key breakthrough: when the model calls a read-only tool (execute_local,
6
+ read_file, search_files, diagnose), the ENGINE executes it directly and feeds
7
+ the result back to the model -- all within a SINGLE API request. This means:
8
+
9
+ 1. User asks "is Home Assistant running?"
10
+ 2. Model calls execute_local(systemctl status home-assistant)
11
+ 3. Engine EXECUTES it -> "not found"
12
+ 4. Engine feeds result back to model
13
+ 5. Model tries execute_local(docker ps | grep -i home)
14
+ 6. Engine EXECUTES it -> "not found"
15
+ 7. Model tries execute_local(ss -tlnp | grep 8123)
16
+ 8. Engine EXECUTES it -> port 8123 listening
17
+ 9. Model tries execute_local(curl -s http://localhost:8123/api/)
18
+ 10. Engine EXECUTES it -> HA API responds
19
+ 11. Model answers: "Yes, HA is running on port 8123 via Docker"
20
+
21
+ All 10 steps happen in ONE API call. The CLI sees only the final answer.
22
+ For WRITE tools (write_file, edit_file), the Engine delegates to the CLI.
23
+
24
+ Endpoints:
25
+ POST /v1/chat/completions -- Standard chat completion (with tool support)
26
+ GET /v1/models -- List available models
27
+ GET /health -- Health check
28
+ GET /stats -- Engine statistics
29
+ """
30
+
31
+ import json
32
+ import os
33
+ import sys
34
+ import time
35
+ import uuid
36
+ import argparse
37
+ import re
38
+ from typing import Dict, List, Optional, Tuple
39
+
40
+ from .pipeline import Pipeline, EngineConfig, PipelineResult
41
+ from .router import classify_intent
42
+ from .executor import execute_tool, is_safe_tool, SAFE_TOOLS, DELEGATE_TOOLS
43
+
44
+
45
+ def has_fake_tool_execution(text: str) -> bool:
46
+ """Detect if text contains hallucinated/fake tool execution.
47
+
48
+ ANY bash/shell code block with a system command means the model is
49
+ writing commands in text instead of using real tool_calls. This is
50
+ ALWAYS wrong -- the model should call execute_local, not print commands.
51
+ """
52
+ if not text:
53
+ return False
54
+ t = text.lower()
55
+ # Pattern: ANY code block with a system command = fake tool execution
56
+ # The model should be using tool_calls, not writing bash blocks
57
+ if re.search(r'```(?:bash|shell|sh)?\s*\n\s*(?:ss|netstat|nmap|curl|ps|docker|find|grep|ls|cat|df|du|systemctl|journalctl|free|top|uptime|uname|id|ifconfig|ip|iptables|who|w|last|lastlog|users|finger|hostname|whoami|date|mount|lsblk|lsof|apt|dpkg|pip|npm|service|timedatectl|hostnamectl|sleep|HA_TOKEN)\b', t):
58
+ return True
59
+ if 'ejecutemos' in t or 'let me run' in t or 'voy a ejecutar' in t or 'voy a verificar' in t:
60
+ if '```' in t:
61
+ return True
62
+ # Pattern: any code block immediately followed by "Result:" with fake data
63
+ if re.search(r'```[^`]*```\s*\n*result\s*:', t):
64
+ return True
65
+ return False
66
+
67
+
68
+ def clean_model_output(text: str) -> str:
69
+ """Remove model artifacts and hallucinated tool execution from output text."""
70
+ if not text:
71
+ return text
72
+ # Remove trailing <tool_call> or [TOOL_CALLS] artifacts
73
+ text = re.sub(r'</?tool_call>\s*$', '', text)
74
+ text = re.sub(r'\[TOOL_CALLS?\]\s*$', '', text)
75
+ text = re.sub(r'\[/TOOL_CALLS?\]\s*$', '', text)
76
+ # Remove partial JSON tool call attempts at the end
77
+ text = re.sub(r'\{"name"\s*:\s*"[^"]*"\s*,?\s*("arguments".*)?$', '', text)
78
+ # Remove markdown-wrapped empty tool calls
79
+ text = re.sub(r'```json\s*\{[^}]*\}\s*```\s*$', '', text)
80
+
81
+ # Remove ALL code blocks containing system commands + fake output
82
+ # This catches: ```bash\nss -tuln\n```\n\nResult: ...\n```...```
83
+ text = re.sub(
84
+ r'```(?:bash|shell|sh)?\s*\n\s*(?:ss|netstat|nmap|curl|ps|docker|find|grep|ls|cat|df|du|systemctl|journalctl|apt|pip|npm|which|whereis|locate|dpkg|service|mount|lsblk|lsof|free|top|htop|uptime|hostname|uname|whoami|id|ifconfig|ip|ping|traceroute|dig|nslookup|iptables|sudo|who|w|last|lastlog|users|date|timedatectl)\b[^`]*```',
85
+ '', text, flags=re.MULTILINE | re.DOTALL
86
+ )
87
+
88
+ # Remove fake "Result:" sections (with or without code blocks)
89
+ text = re.sub(
90
+ r'(?:Result|Output|Salida)\s*:\s*(?:Active|Proto|tcp|udp|LISTEN|Filesystem|total|Mem:|Swap:).*?(?=\n\n[A-Z]|\Z)',
91
+ '', text, flags=re.DOTALL | re.IGNORECASE
92
+ )
93
+
94
+ # Remove hallucinated inline tool execution blocks
95
+ # Pattern: ```bash\nCOMMAND\n```\n\nResult of tool:\nFAKE_OUTPUT
96
+ text = re.sub(
97
+ r'```(?:bash|shell|sh)\n[^`]+```\s*\n*(?:Result of tool:?\s*\n[^\n]*(?:\n[^\n]*)*?)(?=\n\n|\Z)',
98
+ '', text, flags=re.MULTILINE
99
+ )
100
+ # Pattern: bare ```\nCOMMAND\n``` blocks (no language tag)
101
+ text = re.sub(
102
+ r'```\n(?:(?:ss|netstat|nmap|curl|ps|docker|find|grep|ls|cat|df|du|systemctl|journalctl|apt|pip|npm|which|whereis|locate|dpkg|service|mount|lsblk|lsof|free|top|htop|uptime|hostname|uname|whoami|id|ifconfig|ip|ping|traceroute|dig|nslookup)\b[^`]*)\n```',
103
+ '', text, flags=re.MULTILINE
104
+ )
105
+ # Pattern: "I ran: tool_name(...)\n[OK] Result:\n..."
106
+ text = re.sub(
107
+ r'I ran:\s*\w+\([^)]*\)\s*\n\[(?:OK|FAIL(?:ED)?)\]\s*Result:\n.*?(?=\n\n|\Z)',
108
+ '', text, flags=re.DOTALL
109
+ )
110
+ # Pattern: standalone "Result of tool:" blocks
111
+ text = re.sub(
112
+ r'Result of tool:\s*\n.*?(?=\n\n|\Z)',
113
+ '', text, flags=re.DOTALL
114
+ )
115
+ # Remove trailing lazy/filler phrases
116
+ lazy_tails = [
117
+ r'\n+(?:Let me know|If you\'?d? like|Do you want|Quieres que|Te gustaría|Si necesitas|Para más detalles).*$',
118
+ r'\n+(?:Would you like me to|Want me to|If you want me to|Si quieres que).*$',
119
+ ]
120
+ for pattern in lazy_tails:
121
+ text = re.sub(pattern, '', text, flags=re.IGNORECASE | re.DOTALL)
122
+
123
+ # Clean up excessive blank lines left behind
124
+ text = re.sub(r'\n{3,}', '\n\n', text)
125
+ return text.strip()
126
+
127
+
128
+ # Max number of server-side investigation rounds per request
129
+ MAX_INVESTIGATION_DEPTH = 8
130
+
131
+ # Max total time for server-side investigation (seconds)
132
+ MAX_INVESTIGATION_TIME = 90
133
+
134
+
135
+ def count_tool_depth(messages: List[Dict]) -> int:
136
+ """Count tool call rounds in messages."""
137
+ depth = 0
138
+ for msg in messages:
139
+ if msg.get("role") == "assistant" and msg.get("tool_calls"):
140
+ depth += 1
141
+ return depth
142
+
143
+
144
+ def find_original_query(messages: List[Dict]) -> str:
145
+ """Find the first user message (the original question)."""
146
+ for msg in messages:
147
+ if msg.get("role") == "user":
148
+ content = msg.get("content", "")
149
+ if content and not content.startswith("INVESTIGATION"):
150
+ return content
151
+ return ""
152
+
153
+
154
+ def create_app(config: EngineConfig = None):
155
+ """Create the Flask app with the engine pipeline."""
156
+
157
+ try:
158
+ from flask import Flask, request, jsonify, Response
159
+ except ImportError:
160
+ print("Flask not installed. Install with: pip install flask")
161
+ sys.exit(1)
162
+
163
+ app = Flask("tokioai-engine")
164
+ engine_config = config or EngineConfig()
165
+ engine = Pipeline(engine_config)
166
+
167
+ # Track investigation stats
168
+ investigation_stats = {
169
+ "total_investigations": 0,
170
+ "total_steps": 0,
171
+ "max_depth_reached": 0,
172
+ "delegated_to_cli": 0,
173
+ "avg_steps": 0,
174
+ }
175
+
176
+ @app.route("/health", methods=["GET"])
177
+ def health():
178
+ return jsonify({
179
+ "status": "ok",
180
+ "engine": "tokioai-engine",
181
+ "version": "3.0.0",
182
+ "model": engine.config.model_name,
183
+ "ollama_host": engine.config.ollama_host,
184
+ "max_investigation_depth": MAX_INVESTIGATION_DEPTH,
185
+ "max_investigation_time": MAX_INVESTIGATION_TIME,
186
+ "server_side_execution": True,
187
+ })
188
+
189
+ @app.route("/stats", methods=["GET"])
190
+ def stats():
191
+ s = engine.get_stats()
192
+ s["investigation"] = investigation_stats
193
+ return jsonify(s)
194
+
195
+ @app.route("/v1/models", methods=["GET"])
196
+ def list_models():
197
+ return jsonify({
198
+ "object": "list",
199
+ "data": [
200
+ {
201
+ "id": "tokioai-engine",
202
+ "object": "model",
203
+ "created": int(time.time()),
204
+ "owned_by": "tokioai",
205
+ "permission": [],
206
+ "root": engine.config.model_name,
207
+ "parent": None,
208
+ }
209
+ ]
210
+ })
211
+
212
+ @app.route("/v1/chat/completions", methods=["POST"])
213
+ def chat_completions():
214
+ data = request.get_json()
215
+
216
+ if not data:
217
+ return jsonify({"error": {"message": "Request body is required"}}), 400
218
+
219
+ messages = data.get("messages", [])
220
+ if not messages:
221
+ return jsonify({"error": {"message": "messages is required"}}), 400
222
+
223
+ stream = data.get("stream", False)
224
+
225
+ # ============================================================
226
+ # v3.0 SERVER-SIDE INVESTIGATION LOOP
227
+ # ============================================================
228
+ #
229
+ # Architecture:
230
+ # 1. Extract user message and conversation history
231
+ # 2. Call model (via Pipeline) -- model returns tool_calls or text
232
+ # 3. If tool_calls:
233
+ # a. If ALL tools are SAFE (read-only): execute server-side,
234
+ # feed results back to model, goto step 2
235
+ # b. If ANY tool is DELEGATE (write/remote): return tool_calls
236
+ # to CLI for execution (old behavior)
237
+ # 4. If text: return the text answer
238
+ # 5. Repeat up to MAX_INVESTIGATION_DEPTH or MAX_INVESTIGATION_TIME
239
+
240
+ # Determine if this is a tool-result continuation from CLI
241
+ last_msg_role = None
242
+ for m in reversed(messages):
243
+ if m.get("role") != "system":
244
+ last_msg_role = m.get("role")
245
+ break
246
+
247
+ is_tool_continuation = (last_msg_role == "tool")
248
+
249
+ # Build the working message history
250
+ user_msg, working_messages = _extract_messages(messages)
251
+
252
+ if not user_msg:
253
+ return jsonify({"error": {"message": "No user message found"}}), 400
254
+
255
+ original_query = user_msg
256
+
257
+ # If this is a CLI tool continuation (CLI executed a write/remote tool),
258
+ # let the model see the result and decide next step
259
+ if is_tool_continuation:
260
+ # The CLI sent back results from a delegated tool.
261
+ # Feed it all through and let the model continue.
262
+ pass
263
+
264
+ # === PRE-CLASSIFICATION: Skip investigation for pure text ===
265
+ # For greetings, thanks -- skip the investigation loop entirely.
266
+ # BUT NOT for affirmations ("dale", "si", "ok") which are contextual
267
+ # and need conversation history to understand what the user is agreeing to.
268
+ _GREETING_PATTERNS = [
269
+ r'^(hi|hello|hey|hola|que\s+tal|buenos?\s+d[ií]as?|buenas?\s+(?:tardes?|noches?))\s*[.!?]*$',
270
+ r'^(thanks?|thank\s+you|gracias|thx|de\s+nada)\s*[.!?]*$',
271
+ r'^(como\s+(va|estas?|andas?|te\s+va)|how\s+are\s+you|how\'?s\s+it\s+going|what\'?s\s+up)\s*[.!?]*$',
272
+ r'^(todo\s+bien|bien\s+y\s+tu|genial|cool|nice|great|awesome)\s*[.!?]*$',
273
+ r'^(que\s+onda|que\s+hay|que\s+pasa|que\s+cuentas)\s*[.!?]*$',
274
+ r'\b(who\s+are\s+you|what\s+can\s+you|tell\s+me\s+about\s+yourself)\b',
275
+ ]
276
+ import re as _re
277
+ _is_pure_greeting = any(_re.match(p, user_msg.lower().strip(), _re.IGNORECASE) for p in _GREETING_PATTERNS)
278
+ if not is_tool_continuation and _is_pure_greeting:
279
+ _shortcut_start = time.time()
280
+ if engine_config.verbose:
281
+ print(f"[ENGINE] Pure text shortcut: '{user_msg}' -> greeting/thanks")
282
+ engine.reset()
283
+ result = engine.process(
284
+ user_message=user_msg,
285
+ conversation_history=working_messages,
286
+ force_intent="text",
287
+ )
288
+ elapsed = time.time() - _shortcut_start
289
+ return jsonify({
290
+ "id": f"chatcmpl-{uuid.uuid4().hex[:8]}",
291
+ "object": "chat.completion",
292
+ "created": int(time.time()),
293
+ "model": "tokioai-engine",
294
+ "choices": [{
295
+ "index": 0,
296
+ "message": {
297
+ "role": "assistant",
298
+ "content": clean_model_output(result.text or ""),
299
+ },
300
+ "finish_reason": "stop",
301
+ }],
302
+ "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
303
+ "engine_metadata": {
304
+ "investigation_depth": 0,
305
+ "investigation_log": [],
306
+ "elapsed_s": round(elapsed, 1),
307
+ "shortcut": "pure_text",
308
+ },
309
+ })
310
+
311
+ # === THE INVESTIGATION LOOP ===
312
+ engine.reset()
313
+
314
+ investigation_log = [] # Track what we tried
315
+ start_time = time.time()
316
+ final_result = None
317
+ delegated_tool_calls = None
318
+ depth = 0
319
+
320
+ for depth in range(MAX_INVESTIGATION_DEPTH):
321
+ elapsed = time.time() - start_time
322
+ if elapsed > MAX_INVESTIGATION_TIME:
323
+ # Time limit -- force text summary
324
+ if engine_config.verbose:
325
+ print(f"[ENGINE] Investigation time limit reached ({elapsed:.1f}s)")
326
+ break
327
+
328
+ # Build the prompt for this round
329
+ if depth == 0 and not is_tool_continuation:
330
+ # First round: use original query
331
+ current_msg = user_msg
332
+ force_intent = None
333
+ elif depth == 0 and is_tool_continuation:
334
+ # CLI sent back tool results -- summarize and continue
335
+ current_msg = _build_continuation_prompt(messages, original_query)
336
+ force_intent = None
337
+ else:
338
+ # Investigation round N > 0: tell model what happened
339
+ current_msg = _build_investigation_prompt(
340
+ original_query, investigation_log, depth
341
+ )
342
+ # For simple queries (1 success is enough), let model decide
343
+ # For ongoing investigation, force tool calls until near the limit
344
+ successes_so_far = sum(1 for e in investigation_log if e["success"])
345
+ if depth >= MAX_INVESTIGATION_DEPTH - 2:
346
+ force_intent = "text" # Near limit: force final answer
347
+ elif successes_so_far >= 1 and depth <= 2:
348
+ force_intent = None # Let model decide (may answer or investigate more)
349
+ else:
350
+ force_intent = None # Let model decide
351
+
352
+ # Call the model
353
+ try:
354
+ result = engine.process(
355
+ user_message=current_msg,
356
+ conversation_history=working_messages,
357
+ force_intent=force_intent,
358
+ )
359
+ except Exception as e:
360
+ if engine_config.verbose:
361
+ print(f"[ENGINE] Model error at depth {depth}: {e}")
362
+ final_result = PipelineResult(
363
+ output_type="text",
364
+ text=f"Error during investigation: {e}",
365
+ )
366
+ break
367
+
368
+ if engine_config.verbose:
369
+ if result.tool_calls:
370
+ names = [tc["name"] for tc in result.tool_calls]
371
+ print(f"[ENGINE] Depth {depth}: tool_calls={names}")
372
+ else:
373
+ print(f"[ENGINE] Depth {depth}: text response ({len(result.text)} chars)")
374
+
375
+ # === MODEL RETURNED TEXT ===
376
+ if result.output_type == "text" or not result.tool_calls:
377
+ raw_text = result.text or ""
378
+ text = raw_text.lower()
379
+
380
+ # CRITICAL: Detect FAKE tool execution (hallucinated command output)
381
+ if has_fake_tool_execution(raw_text) and depth < MAX_INVESTIGATION_DEPTH - 1:
382
+ if engine_config.verbose:
383
+ print(f"[ENGINE] FAKE tool execution detected at depth {depth}, forcing real tool call")
384
+ working_messages.append({
385
+ "role": "user",
386
+ "content": (
387
+ f"STOP. You just simulated a command in text instead of actually "
388
+ f"running it. That output is FAKE. Use a real tool_call to execute "
389
+ f"the command. Call execute_local with the appropriate command. "
390
+ f"Question: {original_query}"
391
+ ),
392
+ })
393
+ continue # Force another round with real tool
394
+
395
+ # Detect LAZY/INCOMPLETE answers and force more investigation
396
+ lazy_patterns = [
397
+ "you should", "te recomiendo", "se recomienda",
398
+ "requires further", "requiere.*más profund",
399
+ "necesitaría", "se necesitaría", "se necesita",
400
+ "want me to check", "quieres que revise",
401
+ "te gustaría que", "would you like me to",
402
+ "para evaluar", "para determinar",
403
+ "no se detectaron vulnerabilidades inmediatas",
404
+ "you could check", "podrias revisar",
405
+ "voy a verificar", "voy a ejecutar",
406
+ "ejecutemos", "let me run", "let me check",
407
+ ]
408
+ is_lazy = any(p in text for p in lazy_patterns)
409
+ successes_so_far = sum(1 for e in investigation_log if e["success"])
410
+
411
+ if is_lazy and successes_so_far < 4 and depth < MAX_INVESTIGATION_DEPTH - 2:
412
+ if engine_config.verbose:
413
+ print(f"[ENGINE] Lazy answer detected at depth {depth}, pushing for more investigation")
414
+ working_messages.append({"role": "assistant", "content": raw_text})
415
+ working_messages.append({
416
+ "role": "user",
417
+ "content": (
418
+ f"Your answer was incomplete. You said things like 'you should check' "
419
+ f"or 'I will verify'. YOU are the expert -- don't describe what you "
420
+ f"will do, ACTUALLY DO IT with a tool call right now. "
421
+ f"Original question: {original_query}"
422
+ ),
423
+ })
424
+ continue # Force another round
425
+
426
+ # Clean the output before accepting it
427
+ result.text = clean_model_output(raw_text)
428
+ final_result = result
429
+ break
430
+
431
+ # === MODEL RETURNED TOOL CALLS ===
432
+ # Check if all tools are safe for server-side execution
433
+ all_safe = all(is_safe_tool(tc["name"]) for tc in result.tool_calls)
434
+
435
+ if not all_safe:
436
+ # At least one tool needs CLI execution (write/remote)
437
+ # Return tool_calls to CLI, investigation pauses
438
+ delegated_tool_calls = result
439
+ investigation_stats["delegated_to_cli"] += 1
440
+ if engine_config.verbose:
441
+ unsafe = [tc["name"] for tc in result.tool_calls if not is_safe_tool(tc["name"])]
442
+ print(f"[ENGINE] Delegating to CLI: {unsafe}")
443
+ break
444
+
445
+ # ALL tools are safe -- execute server-side!
446
+ tool_results = []
447
+ for tc in result.tool_calls:
448
+ name = tc["name"]
449
+ args = tc["arguments"]
450
+
451
+ if engine_config.verbose:
452
+ arg_preview = args.get("command", args.get("path", args.get("pattern", args.get("target", str(args)))))
453
+ print(f"[ENGINE] Executing: {name}({arg_preview})")
454
+
455
+ success, output = execute_tool(name, args, timeout=30)
456
+
457
+ tool_results.append({
458
+ "name": name,
459
+ "arguments": args,
460
+ "output": output,
461
+ "success": success,
462
+ })
463
+
464
+ investigation_log.append({
465
+ "name": name,
466
+ "arguments": args,
467
+ "output": output[:500], # Keep log manageable
468
+ "success": success,
469
+ "depth": depth,
470
+ })
471
+
472
+ if engine_config.verbose:
473
+ status = "OK" if success else "FAIL"
474
+ print(f"[ENGINE] Result [{status}]: {output[:100]}")
475
+
476
+ # Feed tool results back to model using Ollama's native tool format
477
+ # Ollama /api/chat expects: arguments as dict (not JSON string),
478
+ # no "id"/"type" in tool_calls, no "tool_call_id" in tool messages
479
+
480
+ # Add assistant message with tool_calls
481
+ working_messages.append({
482
+ "role": "assistant",
483
+ "content": "",
484
+ "tool_calls": [
485
+ {
486
+ "function": {
487
+ "name": tr["name"],
488
+ "arguments": tr["arguments"] if isinstance(tr["arguments"], dict) else json.loads(tr["arguments"]),
489
+ },
490
+ }
491
+ for tr in tool_results
492
+ ],
493
+ })
494
+
495
+ # Add each tool result as a tool message
496
+ for tr in tool_results:
497
+ working_messages.append({
498
+ "role": "tool",
499
+ "content": tr["output"],
500
+ })
501
+
502
+ # === END OF INVESTIGATION LOOP ===
503
+
504
+ # Update stats
505
+ if depth > 0:
506
+ investigation_stats["total_investigations"] += 1
507
+ investigation_stats["total_steps"] += depth
508
+ investigation_stats["max_depth_reached"] = max(
509
+ investigation_stats["max_depth_reached"], depth
510
+ )
511
+ total = investigation_stats["total_investigations"]
512
+ investigation_stats["avg_steps"] = (
513
+ investigation_stats["total_steps"] / total if total > 0 else 0
514
+ )
515
+
516
+ # If we hit the depth/time limit without a final answer, force one
517
+ if final_result is None and delegated_tool_calls is None:
518
+ # Force the model to summarize
519
+ summary_prompt = _build_summary_prompt(original_query, investigation_log)
520
+ try:
521
+ final_result = engine.process(
522
+ user_message=summary_prompt,
523
+ conversation_history=working_messages,
524
+ force_intent="text",
525
+ )
526
+ except:
527
+ final_result = PipelineResult(
528
+ output_type="text",
529
+ text=_format_investigation_summary(original_query, investigation_log),
530
+ )
531
+
532
+ # Use delegated tool calls if that's what we have
533
+ if delegated_tool_calls is not None:
534
+ result_to_return = delegated_tool_calls
535
+ else:
536
+ result_to_return = final_result
537
+
538
+ chat_id = f"chatcmpl-{uuid.uuid4().hex[:8]}"
539
+ created = int(time.time())
540
+
541
+ # --- STREAMING RESPONSE (SSE) ---
542
+ if stream:
543
+ def generate_sse():
544
+ msg = _build_response_message(result_to_return)
545
+ finish_reason = "tool_calls" if result_to_return.tool_calls else "stop"
546
+
547
+ if result_to_return.tool_calls:
548
+ tc_deltas = []
549
+ for i, tc in enumerate(result_to_return.tool_calls):
550
+ tc_deltas.append({
551
+ "index": i,
552
+ "id": f"call_{uuid.uuid4().hex[:8]}",
553
+ "type": "function",
554
+ "function": {
555
+ "name": tc["name"],
556
+ "arguments": json.dumps(tc["arguments"]) if isinstance(tc["arguments"], dict) else tc["arguments"],
557
+ },
558
+ })
559
+ chunk = {
560
+ "id": chat_id,
561
+ "object": "chat.completion.chunk",
562
+ "created": created,
563
+ "model": "tokioai-engine",
564
+ "choices": [{
565
+ "index": 0,
566
+ "delta": {
567
+ "role": "assistant",
568
+ "content": None,
569
+ "tool_calls": tc_deltas,
570
+ },
571
+ "finish_reason": finish_reason,
572
+ }],
573
+ }
574
+ yield f"data: {json.dumps(chunk)}\n\n"
575
+ else:
576
+ text = clean_model_output(result_to_return.text or "")
577
+
578
+ # First chunk: role
579
+ chunk = {
580
+ "id": chat_id,
581
+ "object": "chat.completion.chunk",
582
+ "created": created,
583
+ "model": "tokioai-engine",
584
+ "choices": [{
585
+ "index": 0,
586
+ "delta": {"role": "assistant", "content": ""},
587
+ "finish_reason": None,
588
+ }],
589
+ }
590
+ yield f"data: {json.dumps(chunk)}\n\n"
591
+
592
+ if text:
593
+ # Stream word by word
594
+ words = text.split(' ')
595
+ for i, word in enumerate(words):
596
+ token = word if i == 0 else ' ' + word
597
+ chunk = {
598
+ "id": chat_id,
599
+ "object": "chat.completion.chunk",
600
+ "created": created,
601
+ "model": "tokioai-engine",
602
+ "choices": [{
603
+ "index": 0,
604
+ "delta": {"content": token},
605
+ "finish_reason": None,
606
+ }],
607
+ }
608
+ yield f"data: {json.dumps(chunk)}\n\n"
609
+
610
+ # Final chunk
611
+ chunk = {
612
+ "id": chat_id,
613
+ "object": "chat.completion.chunk",
614
+ "created": created,
615
+ "model": "tokioai-engine",
616
+ "choices": [{
617
+ "index": 0,
618
+ "delta": {},
619
+ "finish_reason": "stop",
620
+ }],
621
+ }
622
+ yield f"data: {json.dumps(chunk)}\n\n"
623
+
624
+ yield "data: [DONE]\n\n"
625
+
626
+ return Response(
627
+ generate_sse(),
628
+ mimetype="text/event-stream",
629
+ headers={
630
+ "Cache-Control": "no-cache",
631
+ "X-Accel-Buffering": "no",
632
+ },
633
+ )
634
+
635
+ # --- NON-STREAMING RESPONSE ---
636
+ response = {
637
+ "id": chat_id,
638
+ "object": "chat.completion",
639
+ "created": created,
640
+ "model": "tokioai-engine",
641
+ "choices": [
642
+ {
643
+ "index": 0,
644
+ "message": _build_response_message(result_to_return),
645
+ "finish_reason": "tool_calls" if result_to_return.tool_calls else "stop",
646
+ }
647
+ ],
648
+ "usage": {
649
+ "prompt_tokens": 0,
650
+ "completion_tokens": 0,
651
+ "total_tokens": 0,
652
+ },
653
+ "engine_metadata": {
654
+ "investigation_depth": depth,
655
+ "investigation_log": [
656
+ {"tool": e["name"], "success": e["success"]}
657
+ for e in investigation_log
658
+ ],
659
+ "elapsed_s": round(time.time() - start_time, 1),
660
+ },
661
+ }
662
+
663
+ return jsonify(response)
664
+
665
+ # ============================================================
666
+ # HELPER FUNCTIONS
667
+ # ============================================================
668
+
669
+ def _extract_messages(messages: List[Dict]) -> Tuple[str, List[Dict]]:
670
+ """
671
+ Extract the user message and build Ollama-compatible history.
672
+ Converts tool messages to assistant context.
673
+ Returns (user_msg, ollama_messages).
674
+ """
675
+ user_msg = None
676
+ history = []
677
+
678
+ for msg in messages:
679
+ role = msg.get("role", "")
680
+ content = msg.get("content", "")
681
+
682
+ if role == "system":
683
+ history.append({"role": "system", "content": content})
684
+
685
+ elif role == "user":
686
+ if user_msg is not None:
687
+ history.append({"role": "user", "content": user_msg})
688
+ user_msg = content
689
+
690
+ elif role == "assistant":
691
+ tool_calls = msg.get("tool_calls", [])
692
+ if tool_calls:
693
+ # Convert OpenAI tool_calls format to Ollama format
694
+ # Ollama wants: arguments as dict, no id/type fields
695
+ ollama_tcs = []
696
+ for tc in tool_calls:
697
+ func = tc.get("function", {})
698
+ name = func.get("name", "")
699
+ raw_args = func.get("arguments", {})
700
+ if isinstance(raw_args, str):
701
+ try:
702
+ raw_args = json.loads(raw_args)
703
+ except (json.JSONDecodeError, TypeError):
704
+ raw_args = {}
705
+ ollama_tcs.append({
706
+ "function": {"name": name, "arguments": raw_args},
707
+ })
708
+ history.append({
709
+ "role": "assistant",
710
+ "content": content or "",
711
+ "tool_calls": ollama_tcs,
712
+ })
713
+ elif content:
714
+ history.append({"role": "assistant", "content": content})
715
+
716
+ elif role == "tool":
717
+ # Ollama tool messages: just role + content (no tool_call_id)
718
+ history.append({
719
+ "role": "tool",
720
+ "content": content or "(no output)",
721
+ })
722
+
723
+ return user_msg, history
724
+
725
+ def _build_continuation_prompt(messages: List[Dict], original_query: str) -> str:
726
+ """Build prompt when CLI sends back results from delegated tools."""
727
+ # Find the last tool result
728
+ last_tool_result = ""
729
+ last_tool_name = ""
730
+ for msg in reversed(messages):
731
+ if msg.get("role") == "tool":
732
+ last_tool_result = msg.get("content", "")[:500]
733
+ last_tool_name = msg.get("name", "tool")
734
+ break
735
+
736
+ return (
737
+ f"The tool {last_tool_name} returned:\n{last_tool_result}\n\n"
738
+ f"Original question: {original_query}\n\n"
739
+ f"Based on this result, either:\n"
740
+ f"1. Call more tools to continue investigating, OR\n"
741
+ f"2. Provide your final answer if you have enough information."
742
+ )
743
+
744
+ def _build_investigation_prompt(
745
+ original_query: str,
746
+ investigation_log: List[Dict],
747
+ depth: int,
748
+ ) -> str:
749
+ """
750
+ Build a prompt for the next investigation round.
751
+
752
+ v4.0 PHILOSOPHY: Let the model decide when it has enough data.
753
+ Don't force early answers. Only push for final answer near depth limit.
754
+ """
755
+ failures = sum(1 for e in investigation_log if not e["success"])
756
+ successes = sum(1 for e in investigation_log if e["success"])
757
+
758
+ # Near the depth limit -- force final answer
759
+ if depth >= MAX_INVESTIGATION_DEPTH - 2:
760
+ return (
761
+ f"FINAL ROUND. You have gathered data from {successes} successful "
762
+ f"tool calls. Provide your FINAL answer to: {original_query}\n\n"
763
+ f"Summarize ALL findings. Be specific -- report exact values, "
764
+ f"ports, processes, permissions. No more tool calls."
765
+ )
766
+
767
+ # All failures so far -- try different approach
768
+ if failures > 0 and successes == 0:
769
+ return (
770
+ f"The previous commands failed. Try a DIFFERENT approach "
771
+ f"to answer: {original_query}\n"
772
+ f"Use a different command. Do NOT repeat what already failed."
773
+ )
774
+
775
+ # Normal continuation -- let model freely choose
776
+ already = [e["arguments"].get("command", e["arguments"].get("path", "?"))[:50] for e in investigation_log[-3:]]
777
+ return (
778
+ f"Original question: {original_query}\n"
779
+ f"You've run {len(investigation_log)} commands ({successes} OK, {failures} failed).\n"
780
+ f"Recent: {', '.join(already)}\n\n"
781
+ f"If you have enough information, provide your FINAL answer in plain text.\n"
782
+ f"If you need more data, call another tool with a DIFFERENT command."
783
+ )
784
+
785
+ def _build_summary_prompt(original_query: str, investigation_log: List[Dict]) -> str:
786
+ """Build a forced-text prompt summarizing all investigation findings."""
787
+ return (
788
+ f"Based on all the tool results above, provide a CLEAR, DEFINITIVE "
789
+ f"answer to: {original_query}\n\n"
790
+ f"Rules for your response:\n"
791
+ f"- Write in plain language, no command blocks\n"
792
+ f"- State what IS working, what ISN'T, and why\n"
793
+ f"- Be direct and specific\n"
794
+ f"- Do NOT include any bash/shell commands or simulated tool output"
795
+ )
796
+
797
+ def _format_investigation_summary(original_query: str, investigation_log: List[Dict]) -> str:
798
+ """Fallback text summary if model fails to summarize."""
799
+ lines = [f"Investigation: {original_query}\n"]
800
+ for entry in investigation_log:
801
+ status = "OK" if entry["success"] else "FAIL"
802
+ arg_str = entry["arguments"].get("command", entry["arguments"].get("path", str(entry["arguments"])))
803
+ lines.append(f" [{status}] {entry['name']}({arg_str})")
804
+ if entry["output"]:
805
+ lines.append(f" -> {entry['output'][:150]}")
806
+ return "\n".join(lines)
807
+
808
+ def _build_response_message(result: PipelineResult) -> Dict:
809
+ """Build OpenAI-style message from pipeline result."""
810
+ if result.output_type == "tool_call" and result.tool_calls:
811
+ return {
812
+ "role": "assistant",
813
+ "content": None,
814
+ "tool_calls": [
815
+ {
816
+ "id": f"call_{uuid.uuid4().hex[:8]}",
817
+ "type": "function",
818
+ "function": {
819
+ "name": tc["name"],
820
+ "arguments": json.dumps(tc["arguments"]) if isinstance(tc["arguments"], dict) else tc["arguments"],
821
+ },
822
+ }
823
+ for tc in result.tool_calls
824
+ ],
825
+ }
826
+ return {
827
+ "role": "assistant",
828
+ "content": clean_model_output(result.text or ""),
829
+ }
830
+
831
+ return app
832
+
833
+
834
+ def main():
835
+ parser = argparse.ArgumentParser(description="TokioAI Engine Server v3.0")
836
+ parser.add_argument("command", choices=["serve"], help="Server command")
837
+ parser.add_argument("--port", type=int, default=8080, help="Server port")
838
+ parser.add_argument("--host", type=str, default="127.0.0.1", help="Bind host")
839
+ parser.add_argument("--ollama-host", type=str, default=None, help="Ollama API host")
840
+ parser.add_argument("--model", type=str, default="tokioai-coder", help="Model name")
841
+ parser.add_argument("--verbose", action="store_true", help="Verbose logging")
842
+
843
+ args = parser.parse_args()
844
+
845
+ config = EngineConfig(
846
+ ollama_host=args.ollama_host,
847
+ model_name=args.model,
848
+ verbose=args.verbose,
849
+ )
850
+
851
+ app = create_app(config)
852
+ print(f"TokioAI Engine v3.0 starting on {args.host}:{args.port}")
853
+ print(f" Model: {config.model_name}")
854
+ print(f" Ollama: {config.ollama_host}")
855
+ print(f" Max investigation depth: {MAX_INVESTIGATION_DEPTH}")
856
+ print(f" Max investigation time: {MAX_INVESTIGATION_TIME}s")
857
+ print(f" Server-side execution: ENABLED (tools: {', '.join(sorted(SAFE_TOOLS))})")
858
+ print(f" Verbose: {args.verbose}")
859
+
860
+ app.run(host=args.host, port=args.port, threaded=True)
861
+
862
+
863
+ if __name__ == "__main__":
864
+ main()
engine/test_engine.py ADDED
@@ -0,0 +1,220 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ TokioAI Engine -- Test Suite
3
+ Tests router, verifier, memory, and full pipeline.
4
+ """
5
+
6
+ import json
7
+ import time
8
+ import sys
9
+
10
+ from .router import classify_intent, build_system_prompt, IntentType
11
+ from .verifier import Verifier, verify_and_extract
12
+ from .memory import EngineMemory
13
+ from .schemas import validate_tool_call, extract_tool_calls_from_text
14
+ from .guided import parse_guided_output
15
+
16
+
17
+ def run_tests(verbose=False, ollama_host=None, model=None):
18
+ """Run all tests."""
19
+ passed = 0
20
+ failed = 0
21
+
22
+ def check(name, condition, detail=""):
23
+ nonlocal passed, failed
24
+ if condition:
25
+ passed += 1
26
+ print(f" \033[32mPASS\033[0m {name}")
27
+ else:
28
+ failed += 1
29
+ 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 call
69
+ 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 markdown
73
+ 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 name
77
+ 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 param
81
+ r = v.verify('{"name": "execute_local", "arguments": {}}')
82
+ check("verify: missing required param", not r.valid)
83
+
84
+ # Pure text
85
+ r = v.verify("Docker is a containerization platform.")
86
+ check("verify: pure text", r.valid and r.output_type == "text")
87
+
88
+ # Empty
89
+ 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 text
112
+ 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, os
138
+ 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
+ # Cleanup
154
+ 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, EngineConfig
164
+ 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 = 0
175
+ live_failed = 0
176
+
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_type
193
+ tool_match = True
194
+ if expected_tool and result.tool_calls:
195
+ tool_match = result.tool_calls[0]["name"] == expected_tool
196
+
197
+ ok = type_match and tool_match
198
+ 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 += 1
206
+ print(f" \033[32mPASS\033[0m [{result.latency_ms:4d}ms] [{result.attempts}att] {prompt:40s} {detail}")
207
+ else:
208
+ live_failed += 1
209
+ 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 += 1
212
+ 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_passed
218
+ failed += live_failed
219
+
220
+ return passed, failed
engine/verifier.py ADDED
@@ -0,0 +1,186 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ TokioAI Engine -- Output Verifier & Retry Logic
3
+ Validates model responses and orchestrates retries on failure.
4
+
5
+ The verifier catches:
6
+ 1. Invalid JSON tool calls
7
+ 2. Wrong tool names
8
+ 3. Missing required parameters
9
+ 4. Malformed output (e.g., natural language when tool expected)
10
+ 5. Verbosity (strips filler from responses)
11
+ """
12
+
13
+ import json
14
+ import re
15
+ import time
16
+ from typing import Any, Dict, List, Optional, Tuple
17
+
18
+ from .schemas import validate_tool_call, extract_tool_calls_from_text, VALID_TOOL_NAMES
19
+
20
+
21
+ class VerificationResult:
22
+ """Result of verifying model output."""
23
+
24
+ def __init__(
25
+ self,
26
+ valid: bool,
27
+ output_type: str, # "tool_call", "text", "error"
28
+ tool_calls: List[Dict] = None,
29
+ text: str = None,
30
+ error: str = None,
31
+ raw_output: str = None,
32
+ ):
33
+ self.valid = valid
34
+ self.output_type = output_type
35
+ self.tool_calls = tool_calls or []
36
+ self.text = text or ""
37
+ self.error = error or ""
38
+ self.raw_output = raw_output or ""
39
+
40
+ def __repr__(self):
41
+ if self.valid:
42
+ if self.output_type == "tool_call":
43
+ names = [tc["name"] for tc in self.tool_calls]
44
+ return f"VerificationResult(valid=True, tools={names})"
45
+ return f"VerificationResult(valid=True, type={self.output_type})"
46
+ return f"VerificationResult(valid=False, error={self.error!r})"
47
+
48
+
49
+ class Verifier:
50
+ """
51
+ Validates model output and provides structured results.
52
+ """
53
+
54
+ def __init__(self, strict: bool = False):
55
+ """
56
+ Args:
57
+ strict: If True, rejects ambiguous outputs. If False, tries harder to parse.
58
+ """
59
+ self.strict = strict
60
+
61
+ def verify(self, model_output: str, expected_intent: str = None) -> VerificationResult:
62
+ """
63
+ Verify model output.
64
+
65
+ Args:
66
+ model_output: Raw text from the model
67
+ expected_intent: What the router classified (tool_call, text, etc.)
68
+
69
+ Returns:
70
+ VerificationResult with parsed tool calls or cleaned text
71
+ """
72
+ if not model_output or not model_output.strip():
73
+ return VerificationResult(
74
+ valid=False,
75
+ output_type="error",
76
+ error="Empty model output",
77
+ raw_output=model_output,
78
+ )
79
+
80
+ raw = model_output.strip()
81
+
82
+ # Step 1: Try to extract tool calls
83
+ tool_calls = extract_tool_calls_from_text(raw)
84
+
85
+ if tool_calls:
86
+ # Validate each tool call
87
+ valid_calls = []
88
+ errors = []
89
+
90
+ for tc in tool_calls:
91
+ is_valid, err = validate_tool_call(tc["name"], tc["arguments"])
92
+ if is_valid:
93
+ valid_calls.append(tc)
94
+ else:
95
+ errors.append(err)
96
+
97
+ if valid_calls:
98
+ return VerificationResult(
99
+ valid=True,
100
+ output_type="tool_call",
101
+ tool_calls=valid_calls,
102
+ raw_output=raw,
103
+ )
104
+ elif errors:
105
+ return VerificationResult(
106
+ valid=False,
107
+ output_type="error",
108
+ error="; ".join(errors),
109
+ raw_output=raw,
110
+ )
111
+
112
+ # Step 2: If expected tool_call but got text, that might be an error
113
+ if expected_intent == "tool_call":
114
+ # Check if the model is TRYING to make a tool call but failing
115
+ if any(name in raw.lower() for name in VALID_TOOL_NAMES):
116
+ return VerificationResult(
117
+ valid=False,
118
+ output_type="error",
119
+ error="Model mentioned a tool name but didn't format it as a valid JSON tool call",
120
+ raw_output=raw,
121
+ )
122
+
123
+ # Step 3: It's a text response -- clean it up
124
+ cleaned = self._clean_text(raw)
125
+ return VerificationResult(
126
+ valid=True,
127
+ output_type="text",
128
+ text=cleaned,
129
+ raw_output=raw,
130
+ )
131
+
132
+ def _clean_text(self, text: str) -> str:
133
+ """Clean up verbose model output."""
134
+ # Remove thinking tags (Qwen3 /think blocks)
135
+ text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
136
+ # Remove empty lines at start/end
137
+ text = text.strip()
138
+ # Remove excessive newlines
139
+ text = re.sub(r'\n{3,}', '\n\n', text)
140
+ return text
141
+
142
+
143
+ class RetryStrategy:
144
+ """
145
+ Manages retry logic when the verifier rejects model output.
146
+ Each retry gets a more explicit prompt to guide the model.
147
+ """
148
+
149
+ def __init__(self, max_retries: int = 2):
150
+ self.max_retries = max_retries
151
+
152
+ def should_retry(self, attempt: int, result: VerificationResult) -> bool:
153
+ """Should we retry after a failed verification?"""
154
+ if result.valid:
155
+ return False
156
+ if attempt >= self.max_retries:
157
+ return False
158
+ # Don't retry empty outputs (model is broken)
159
+ if not result.raw_output.strip():
160
+ return False
161
+ return True
162
+
163
+ def get_retry_message(self, attempt: int, result: VerificationResult) -> str:
164
+ """Generate a retry prompt based on what went wrong."""
165
+ if attempt == 0:
166
+ # First retry: gentle correction
167
+ return (
168
+ f"Your previous response was not valid: {result.error}\n"
169
+ f"Please respond with ONLY a JSON object like this:\n"
170
+ f'{{"name": "tool_name", "arguments": {{"param": "value"}}}}\n'
171
+ f"No explanation, no markdown, just the JSON."
172
+ )
173
+ else:
174
+ # Second retry: very explicit
175
+ return (
176
+ f"ERROR: {result.error}\n"
177
+ f"You MUST respond with valid JSON only. "
178
+ f"Valid tool names: {', '.join(sorted(VALID_TOOL_NAMES))}.\n"
179
+ f"Format: {{\"name\": \"execute_local\", \"arguments\": {{\"command\": \"your command\"}}}}\n"
180
+ f"Respond with NOTHING except the JSON object."
181
+ )
182
+
183
+
184
+ def verify_and_extract(model_output: str, expected_intent: str = None) -> VerificationResult:
185
+ """Convenience function: verify model output with default settings."""
186
+ return Verifier().verify(model_output, expected_intent)
examples/basic_usage.py ADDED
@@ -0,0 +1,72 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ TokioAI Engine -- Basic Usage Example
4
+
5
+ Prerequisites:
6
+ pip install tokioai-engine
7
+ ollama serve &
8
+ ollama pull qwen3-coder
9
+ """
10
+
11
+ from engine.pipeline import Pipeline, EngineConfig
12
+
13
+ # Create engine with default config (connects to local Ollama)
14
+ config = EngineConfig(
15
+ model_name="qwen3-coder", # Any Ollama model
16
+ ollama_host="http://localhost:11434",
17
+ max_retries=2, # Retry on bad output
18
+ use_guided_gen=True, # Force valid JSON for tool calls
19
+ use_memory=True, # Learn from successful interactions
20
+ verbose=True, # Show engine internals
21
+ )
22
+
23
+ engine = Pipeline(config)
24
+
25
+ # --- Tool Call Example ---
26
+ print("=" * 60)
27
+ print("TOOL CALL: 'check disk space'")
28
+ print("=" * 60)
29
+
30
+ result = engine.process("check disk space")
31
+ print(f"Intent: {result.intent}")
32
+ print(f"Type: {result.output_type}")
33
+ print(f"Calls: {result.tool_calls}")
34
+ print(f"Latency: {result.latency_ms}ms")
35
+ print(f"Retries: {result.attempts}")
36
+ print()
37
+
38
+ # --- Text Response Example ---
39
+ print("=" * 60)
40
+ print("TEXT: 'what is SQL injection'")
41
+ print("=" * 60)
42
+
43
+ engine.reset()
44
+ result = engine.process("what is SQL injection")
45
+ print(f"Intent: {result.intent}")
46
+ print(f"Type: {result.output_type}")
47
+ print(f"Text: {result.text[:200]}")
48
+ print(f"Latency: {result.latency_ms}ms")
49
+ print()
50
+
51
+ # --- Decision Example ---
52
+ print("=" * 60)
53
+ print("DECISION: 'is docker running'")
54
+ print("=" * 60)
55
+
56
+ engine.reset()
57
+ result = engine.process("is docker running")
58
+ print(f"Intent: {result.intent}")
59
+ print(f"Type: {result.output_type}")
60
+ if result.tool_calls:
61
+ print(f"Calls: {result.tool_calls}")
62
+ else:
63
+ print(f"Text: {result.text[:200]}")
64
+ print(f"Latency: {result.latency_ms}ms")
65
+ print()
66
+
67
+ # --- Stats ---
68
+ print("=" * 60)
69
+ print("ENGINE STATS")
70
+ print("=" * 60)
71
+ import json
72
+ print(json.dumps(engine.get_stats(), indent=2))
examples/custom_tools.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ TokioAI Engine -- Custom Tools Example
4
+
5
+ Shows how to add your own tools to the engine.
6
+ Modify engine/schemas.py to add tool definitions,
7
+ or pass them through the API.
8
+ """
9
+
10
+ from engine.pipeline import Pipeline, EngineConfig
11
+ from engine.schemas import TOOL_DEFINITIONS
12
+
13
+ # You can inspect the built-in tools
14
+ print("Built-in tools:")
15
+ for tool in TOOL_DEFINITIONS:
16
+ name = tool["function"]["name"]
17
+ desc = tool["function"]["description"][:60]
18
+ print(f" {name:20s} -- {desc}")
19
+
20
+ print()
21
+
22
+ # Use the engine normally -- it already knows all TokioAI tools
23
+ config = EngineConfig(
24
+ model_name="qwen3-coder",
25
+ verbose=True,
26
+ )
27
+ engine = Pipeline(config)
28
+
29
+ # The engine's router will classify this as tool_call
30
+ # and the guided generator will constrain output to valid tool names
31
+ result = engine.process("search for TODO comments in all python files")
32
+
33
+ print(f"\nResult: {result}")
34
+ if result.tool_calls:
35
+ print(f"Tool: {result.tool_calls[0]['name']}")
36
+ print(f"Args: {result.tool_calls[0]['arguments']}")
examples/openai_client.py ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ TokioAI Engine -- OpenAI Client Example
4
+
5
+ Shows how to use TokioAI Engine as a drop-in replacement for OpenAI API.
6
+
7
+ Prerequisites:
8
+ pip install tokioai-engine openai
9
+ python -m engine serve --model qwen3-coder --port 8080
10
+ """
11
+
12
+ from openai import OpenAI
13
+
14
+ # Connect to TokioAI Engine (same API as OpenAI)
15
+ client = OpenAI(
16
+ base_url="http://localhost:8080/v1",
17
+ api_key="not-needed", # Engine doesn't require auth
18
+ )
19
+
20
+ # --- Tool calling ---
21
+ print("--- Tool Calling ---")
22
+ response = client.chat.completions.create(
23
+ model="tokioai-engine",
24
+ messages=[{"role": "user", "content": "check what processes are using the most memory"}],
25
+ )
26
+
27
+ msg = response.choices[0].message
28
+ if msg.tool_calls:
29
+ for tc in msg.tool_calls:
30
+ print(f"Tool: {tc.function.name}")
31
+ print(f"Arguments: {tc.function.arguments}")
32
+ else:
33
+ print(f"Response: {msg.content}")
34
+
35
+ print()
36
+
37
+ # --- Text response ---
38
+ print("--- Text Response ---")
39
+ response = client.chat.completions.create(
40
+ model="tokioai-engine",
41
+ messages=[{"role": "user", "content": "explain the TCP 3-way handshake in one paragraph"}],
42
+ )
43
+
44
+ print(f"Response: {response.choices[0].message.content}")
45
+ print()
46
+
47
+ # --- Engine metadata ---
48
+ if hasattr(response, "engine_metadata"):
49
+ print(f"Intent: {response.engine_metadata['intent']}")
50
+ print(f"Confidence: {response.engine_metadata['confidence']}")
51
+ print(f"Latency: {response.engine_metadata['latency_ms']}ms")
prompts/iot_system_prompt.md ADDED
@@ -0,0 +1,216 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # TokioAI IoT System Prompt
2
+
3
+ This is the production system prompt used by TokioAI-Coder for IoT/smart-home control.
4
+ Customize it for your own setup by replacing entity names and API endpoints.
5
+
6
+ ---
7
+
8
+ ## System Prompt
9
+
10
+ ```
11
+ You are TokioAI -- an elite cybersecurity and DevOps AI agent.
12
+ You execute commands, fix bugs, deploy infrastructure, and solve problems.
13
+ You have full access to the user's systems via tools.
14
+
15
+ ## CRITICAL RULES -- NEVER VIOLATE THESE
16
+
17
+ 1. NEVER simulate or fake tool execution in your text responses.
18
+ WRONG: writing "```bash\ndf -h\n```\nResult: ..." in your text.
19
+ RIGHT: call the execute_local tool with {"command": "df -h"}.
20
+ The engine will execute the tool and give you the real result.
21
+
22
+ 2. When you need system information, ALWAYS use a tool call.
23
+ NEVER make up or guess command output. NEVER write fake results.
24
+
25
+ 3. Your text responses should contain ONLY your analysis and conclusions.
26
+ NEVER include "I ran:", "Result of tool:", or "```bash" blocks that
27
+ simulate command execution.
28
+
29
+ 4. If you already have real tool results in the conversation, summarize
30
+ them in plain language. Do NOT echo back raw tool output format.
31
+
32
+ ## ENVIRONMENT
33
+ You are running on a Raspberry Pi 5 (Debian/aarch64).
34
+ Key services:
35
+ - Home Assistant: runs as Docker container "homeassistant" (NOT systemd). Check with: docker ps | grep home
36
+ - Docker: available. Most services run as containers.
37
+ - Ollama: runs locally or via SSH tunnel on port 8200
38
+ - TokioAI Engine: this server (port 8080)
39
+ When checking if something is running, ALWAYS check both systemd AND docker.
40
+
41
+ ## ROBOT VISION (robot_vision tool)
42
+ You have a camera on the robot. Use the robot_vision tool to SEE what the robot sees.
43
+ - Use it when the user asks: "what do you see?", "look around", "describe the room", "is anyone there?"
44
+ - The tool captures a live photo and analyzes it with a vision AI model
45
+ - You can pass a specific question: robot_vision(question="how many people are in the room?")
46
+ - ALWAYS use robot_vision for visual questions. NEVER say "I can't see" -- you CAN see.
47
+
48
+ ## HOME ASSISTANT VERIFICATION (CRITICAL)
49
+ After calling ANY HA service (media, lights, switches, etc.), you MUST:
50
+ 1. Execute the service call
51
+ 2. Wait 3 seconds (sleep 3)
52
+ 3. GET the entity state to confirm it ACTUALLY changed
53
+ 4. Report what is ACTUALLY happening, not what you hoped
54
+
55
+ HA auth pattern:
56
+ HA_TOKEN=$(cat /path/to/ha_token) && curl -s -H "Authorization: Bearer $HA_TOKEN" ...
57
+
58
+ Or with env var:
59
+ curl -s -H "Authorization: Bearer $HA_TOKEN" http://localhost:8123/api/...
60
+
61
+ For music changes specifically:
62
+ After play_media, run: sleep 3 && curl -s -H "Authorization: Bearer $HA_TOKEN" \
63
+ http://localhost:8123/api/states/media_player.YOUR_DEVICE | \
64
+ python3 -c "import sys,json;s=json.load(sys.stdin);a=s['attributes'];print(f'State: {s[\"state\"]}, Title: {a.get(\"media_title\",\"?\")}, Artist: {a.get(\"media_artist\",\"?\")}')"
65
+ If title/artist didn't change or state is not "playing", tell the user it FAILED.
66
+
67
+ A 200 OK from the API only means HA received the request, NOT that it worked. ALWAYS verify.
68
+
69
+ ## HOME ASSISTANT COOKBOOK (execute via execute_local tool)
70
+
71
+ ### Authentication
72
+ Auth pattern: curl -s -H "Authorization: Bearer $HA_TOKEN" http://YOUR_HA_URL/api/...
73
+ Base URL: http://localhost:8123 (change to your HA instance)
74
+
75
+ ### Media Player
76
+ Entity: media_player.YOUR_DEVICE (e.g., media_player.jarvis for Amazon Echo)
77
+
78
+ - What's playing:
79
+ GET /api/states/media_player.YOUR_DEVICE
80
+ (check media_title, media_artist, state)
81
+
82
+ - Change music:
83
+ POST /api/services/media_player/play_media
84
+ {"entity_id":"media_player.YOUR_DEVICE","media_content_type":"music","media_content_id":"QUERY"}
85
+ (e.g., "blues", "rock", "Beatles", "lofi hip hop")
86
+
87
+ - Volume:
88
+ POST /api/services/media_player/volume_set
89
+ {"entity_id":"media_player.YOUR_DEVICE","volume_level":0.5}
90
+ (0.0 to 1.0)
91
+
92
+ - Play/Pause:
93
+ POST /api/services/media_player/media_play_pause
94
+ {"entity_id":"media_player.YOUR_DEVICE"}
95
+
96
+ - Next/Previous track:
97
+ POST /api/services/media_player/media_next_track
98
+ POST /api/services/media_player/media_previous_track
99
+
100
+ - TTS (device speaks):
101
+ POST /api/services/notify/alexa_media_YOUR_DEVICE
102
+ {"message":"TEXT","data":{"type":"tts"}}
103
+
104
+ - Announce:
105
+ POST /api/services/notify/alexa_media_YOUR_DEVICE
106
+ {"message":"TEXT","data":{"type":"announce"}}
107
+
108
+ NEVER use notify to change music. notify is ONLY for TTS/announcements.
109
+
110
+ ### Lights
111
+ Entities: light.living_room, light.bedroom, light.kitchen (customize)
112
+
113
+ - Turn on/off:
114
+ POST /api/services/light/turn_on or turn_off
115
+ {"entity_id":"light.YOUR_LIGHT"}
116
+
117
+ - Brightness (0-255):
118
+ POST /api/services/light/turn_on
119
+ {"entity_id":"light.YOUR_LIGHT","brightness":128}
120
+
121
+ - Color (RGB):
122
+ POST /api/services/light/turn_on
123
+ {"entity_id":"light.YOUR_LIGHT","rgb_color":[255,0,0]}
124
+
125
+ ### Switches
126
+ Entities: switch.YOUR_SWITCH
127
+
128
+ - Turn on/off:
129
+ POST /api/services/switch/turn_on or turn_off
130
+ {"entity_id":"switch.YOUR_SWITCH"}
131
+
132
+ ### Sensors
133
+ Entities: sensor.temperature, weather.forecast_home
134
+
135
+ - Read sensor:
136
+ GET /api/states/sensor.YOUR_SENSOR
137
+
138
+ - Weather:
139
+ GET /api/states/weather.forecast_home
140
+
141
+ ### Vacuum
142
+ Entity: vacuum.YOUR_VACUUM
143
+
144
+ - Start/Stop/Return:
145
+ POST /api/services/vacuum/start
146
+ POST /api/services/vacuum/stop
147
+ POST /api/services/vacuum/return_to_base
148
+ {"entity_id":"vacuum.YOUR_VACUUM"}
149
+
150
+ ### Automations
151
+ - List all:
152
+ GET /api/states (filter entity_id starting with "automation.")
153
+
154
+ - Trigger:
155
+ POST /api/services/automation/trigger
156
+ {"entity_id":"automation.YOUR_AUTOMATION"}
157
+
158
+ ## CORE RULES
159
+ - Be DIRECT. Act first, explain after.
160
+ - NEVER give up. If something fails, try alternatives.
161
+ - When investigating: diagnose -> try alternatives -> root cause -> answer.
162
+ - NEVER output passwords, API keys, or tokens in plain text.
163
+ - Be concise but thorough.
164
+
165
+ ## INVESTIGATION BEHAVIOR
166
+ You are an INVESTIGATOR. When asked about system state:
167
+ 1. ALWAYS check the actual system. NEVER speculate.
168
+ 2. If the first command doesn't give the full picture, run MORE commands automatically.
169
+ 3. For "is X running?": check systemctl, docker, ps, ports. Try multiple approaches.
170
+ 4. NEVER ask "do you want me to check more?" -- just CHECK IT.
171
+ 5. NEVER say "you should review X" -- YOU review it and report findings.
172
+ 6. Think like a senior engineer: be thorough, check multiple angles, report specifics.
173
+ 7. Always use TOOL CALLS -- never simulate commands in your text response.
174
+
175
+ ## ANSWER FORMAT
176
+ When you've gathered enough information:
177
+ - State what you found CLEARLY
178
+ - If something is running: say what it is, where, which port, how it was started
179
+ - If something is NOT available: say what you tried and why
180
+ - NEVER give generic "recommendations" -- give FACTS
181
+ - Report EXACTLY what the command output says
182
+ ```
183
+
184
+ ---
185
+
186
+ ## Customization Guide
187
+
188
+ ### Step 1: Replace Entities
189
+ Search for `YOUR_DEVICE`, `YOUR_LIGHT`, `YOUR_SWITCH`, `YOUR_SENSOR`, `YOUR_VACUUM` and replace with your actual Home Assistant entity IDs.
190
+
191
+ ### Step 2: Set Authentication
192
+ Choose one of:
193
+ - Environment variable: `export HA_TOKEN="your-long-lived-token"`
194
+ - File: `echo "your-token" > ~/.ha_token` and use `$(cat ~/.ha_token)`
195
+ - .env file: Add `HA_TOKEN=your-token` to `~/.tokioai/.env`
196
+
197
+ ### Step 3: Add Your Entities
198
+ List all your HA entities:
199
+ ```bash
200
+ curl -s -H "Authorization: Bearer $HA_TOKEN" http://localhost:8123/api/states | \
201
+ python3 -c "import sys,json;[print(s['entity_id']) for s in json.load(sys.stdin)]" | sort
202
+ ```
203
+
204
+ Add them to the appropriate sections in the system prompt.
205
+
206
+ ### Step 4: Test
207
+ ```bash
208
+ # Test HA connection
209
+ curl -s -H "Authorization: Bearer $HA_TOKEN" http://localhost:8123/api/
210
+
211
+ # Test a light
212
+ curl -s -X POST -H "Authorization: Bearer $HA_TOKEN" \
213
+ -H "Content-Type: application/json" \
214
+ http://localhost:8123/api/services/light/turn_on \
215
+ -d '{"entity_id":"light.your_light"}'
216
+ ```
pyproject.toml ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [build-system]
2
+ requires = ["setuptools>=45", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "tokioai-engine"
7
+ version = "1.0.0"
8
+ description = "Transform any small open model into an elite AI agent"
9
+ readme = "README.md"
10
+ license = {text = "MIT"}
11
+ requires-python = ">=3.8"
12
+ authors = [{name = "TokioAI"}]
13
+ keywords = ["ai", "llm", "inference", "tool-calling", "agent", "security", "ollama"]
14
+ dependencies = ["requests>=2.28", "flask>=2.0"]
15
+
16
+ [project.scripts]
17
+ tokioai-engine = "engine.server:main"
18
+
19
+ [project.optional-dependencies]
20
+ dev = ["pytest>=7.0", "openai>=1.0"]
setup.py ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from setuptools import setup, find_packages
2
+
3
+ with open("README.md", "r", encoding="utf-8") as f:
4
+ long_description = f.read()
5
+
6
+ setup(
7
+ name="tokioai-engine",
8
+ version="1.0.0",
9
+ description="Transform any small open model into an elite AI agent with tool calling, guided generation, and adaptive memory",
10
+ long_description=long_description,
11
+ long_description_content_type="text/markdown",
12
+ author="TokioAI",
13
+ author_email="hello@tokio.ai",
14
+ url="https://github.com/TokioAI/tokioai-engine",
15
+ project_urls={
16
+ "Bug Reports": "https://github.com/TokioAI/tokioai-engine/issues",
17
+ "Source": "https://github.com/TokioAI/tokioai-engine",
18
+ },
19
+ packages=find_packages(),
20
+ python_requires=">=3.8",
21
+ install_requires=[
22
+ "requests>=2.28",
23
+ "flask>=2.0",
24
+ ],
25
+ extras_require={
26
+ "dev": [
27
+ "pytest>=7.0",
28
+ "openai>=1.0",
29
+ ],
30
+ },
31
+ entry_points={
32
+ "console_scripts": [
33
+ "tokioai-engine=engine.server:main",
34
+ ],
35
+ },
36
+ classifiers=[
37
+ "Development Status :: 4 - Beta",
38
+ "Intended Audience :: Developers",
39
+ "Intended Audience :: Science/Research",
40
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
41
+ "Topic :: Security",
42
+ "License :: OSI Approved :: MIT License",
43
+ "Programming Language :: Python :: 3",
44
+ "Programming Language :: Python :: 3.8",
45
+ "Programming Language :: Python :: 3.9",
46
+ "Programming Language :: Python :: 3.10",
47
+ "Programming Language :: Python :: 3.11",
48
+ "Programming Language :: Python :: 3.12",
49
+ "Operating System :: OS Independent",
50
+ ],
51
+ keywords="ai llm inference tool-calling agent security devops ollama",
52
+ )
tools/tool_definitions.json ADDED
@@ -0,0 +1,266 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [
2
+ {
3
+ "type": "function",
4
+ "function": {
5
+ "name": "execute_local",
6
+ "description": "Execute a shell command on the local machine. Use for system operations: install packages, check processes, network scans, file operations.",
7
+ "parameters": {
8
+ "type": "object",
9
+ "properties": {
10
+ "command": {
11
+ "type": "string",
12
+ "description": "Shell command to execute (bash)"
13
+ },
14
+ "timeout": {
15
+ "type": "integer",
16
+ "description": "Timeout in seconds (default 30, max 120)"
17
+ }
18
+ },
19
+ "required": ["command"]
20
+ }
21
+ }
22
+ },
23
+ {
24
+ "type": "function",
25
+ "function": {
26
+ "name": "execute_raspi",
27
+ "description": "Execute a command on a Raspberry Pi via SSH.",
28
+ "parameters": {
29
+ "type": "object",
30
+ "properties": {
31
+ "command": {
32
+ "type": "string",
33
+ "description": "Command to run on Raspi"
34
+ },
35
+ "timeout": {
36
+ "type": "integer",
37
+ "description": "Timeout in seconds (default 30, max 120)"
38
+ }
39
+ },
40
+ "required": ["command"]
41
+ }
42
+ }
43
+ },
44
+ {
45
+ "type": "function",
46
+ "function": {
47
+ "name": "execute_gcp",
48
+ "description": "Execute a command on a GCP VM via SSH.",
49
+ "parameters": {
50
+ "type": "object",
51
+ "properties": {
52
+ "command": {
53
+ "type": "string",
54
+ "description": "Command to run on GCP"
55
+ },
56
+ "timeout": {
57
+ "type": "integer",
58
+ "description": "Timeout in seconds (default 30, max 120)"
59
+ }
60
+ },
61
+ "required": ["command"]
62
+ }
63
+ }
64
+ },
65
+ {
66
+ "type": "function",
67
+ "function": {
68
+ "name": "execute_router",
69
+ "description": "Execute a command on the router via SSH (root).",
70
+ "parameters": {
71
+ "type": "object",
72
+ "properties": {
73
+ "command": {
74
+ "type": "string",
75
+ "description": "Command to run on router as root"
76
+ }
77
+ },
78
+ "required": ["command"]
79
+ }
80
+ }
81
+ },
82
+ {
83
+ "type": "function",
84
+ "function": {
85
+ "name": "read_file",
86
+ "description": "Read a file from local or remote systems.",
87
+ "parameters": {
88
+ "type": "object",
89
+ "properties": {
90
+ "path": {
91
+ "type": "string",
92
+ "description": "File path. Use raspi:/path or gcp:/path for remote"
93
+ },
94
+ "lines": {
95
+ "type": "integer",
96
+ "description": "Max lines to read (default: all)"
97
+ }
98
+ },
99
+ "required": ["path"]
100
+ }
101
+ }
102
+ },
103
+ {
104
+ "type": "function",
105
+ "function": {
106
+ "name": "write_file",
107
+ "description": "Write content to a file.",
108
+ "parameters": {
109
+ "type": "object",
110
+ "properties": {
111
+ "path": {
112
+ "type": "string",
113
+ "description": "File path. Use raspi:/path or gcp:/path for remote"
114
+ },
115
+ "content": {
116
+ "type": "string",
117
+ "description": "File content to write"
118
+ }
119
+ },
120
+ "required": ["path", "content"]
121
+ }
122
+ }
123
+ },
124
+ {
125
+ "type": "function",
126
+ "function": {
127
+ "name": "edit_file",
128
+ "description": "Replace a specific string in a file. old_text must be unique.",
129
+ "parameters": {
130
+ "type": "object",
131
+ "properties": {
132
+ "path": {
133
+ "type": "string",
134
+ "description": "File path"
135
+ },
136
+ "old_text": {
137
+ "type": "string",
138
+ "description": "Exact text to find (must be unique)"
139
+ },
140
+ "new_text": {
141
+ "type": "string",
142
+ "description": "Replacement text"
143
+ }
144
+ },
145
+ "required": ["path", "old_text", "new_text"]
146
+ }
147
+ }
148
+ },
149
+ {
150
+ "type": "function",
151
+ "function": {
152
+ "name": "search_files",
153
+ "description": "Search for text patterns across files. Like grep -rn.",
154
+ "parameters": {
155
+ "type": "object",
156
+ "properties": {
157
+ "pattern": {
158
+ "type": "string",
159
+ "description": "Text or regex pattern to search"
160
+ },
161
+ "path": {
162
+ "type": "string",
163
+ "description": "Directory to search"
164
+ },
165
+ "glob": {
166
+ "type": "string",
167
+ "description": "File filter glob (e.g., '*.py')"
168
+ }
169
+ },
170
+ "required": ["pattern"]
171
+ }
172
+ }
173
+ },
174
+ {
175
+ "type": "function",
176
+ "function": {
177
+ "name": "robot_vision",
178
+ "description": "Capture a photo from the robot's camera and analyze it with a vision AI model.",
179
+ "parameters": {
180
+ "type": "object",
181
+ "properties": {
182
+ "question": {
183
+ "type": "string",
184
+ "description": "What to look for in the image (e.g. 'how many people?', 'describe the room')"
185
+ }
186
+ },
187
+ "required": []
188
+ }
189
+ }
190
+ },
191
+ {
192
+ "type": "function",
193
+ "function": {
194
+ "name": "diagnose",
195
+ "description": "Run system health diagnostics.",
196
+ "parameters": {
197
+ "type": "object",
198
+ "properties": {
199
+ "target": {
200
+ "type": "string",
201
+ "enum": ["local", "raspi", "gcp", "router", "all"],
202
+ "description": "Target system"
203
+ }
204
+ },
205
+ "required": ["target"]
206
+ }
207
+ }
208
+ },
209
+ {
210
+ "type": "function",
211
+ "function": {
212
+ "name": "memory",
213
+ "description": "Persistent memory across sessions. Actions: read, write, append, clear.",
214
+ "parameters": {
215
+ "type": "object",
216
+ "properties": {
217
+ "action": {
218
+ "type": "string",
219
+ "enum": ["read", "write", "append", "clear"],
220
+ "description": "Memory operation"
221
+ },
222
+ "content": {
223
+ "type": "string",
224
+ "description": "Content to write/append"
225
+ }
226
+ },
227
+ "required": ["action"]
228
+ }
229
+ }
230
+ },
231
+ {
232
+ "type": "function",
233
+ "function": {
234
+ "name": "task",
235
+ "description": "Persistent task tracker across sessions.",
236
+ "parameters": {
237
+ "type": "object",
238
+ "properties": {
239
+ "action": {
240
+ "type": "string",
241
+ "enum": ["add", "update", "list", "remove", "clear_done"],
242
+ "description": "Task operation"
243
+ },
244
+ "task": {
245
+ "type": "string",
246
+ "description": "Task description (for add)"
247
+ },
248
+ "id": {
249
+ "type": "integer",
250
+ "description": "Task ID (for update/remove)"
251
+ },
252
+ "status": {
253
+ "type": "string",
254
+ "enum": ["pending", "in_progress", "done", "blocked"],
255
+ "description": "New status"
256
+ },
257
+ "note": {
258
+ "type": "string",
259
+ "description": "Note to add"
260
+ }
261
+ },
262
+ "required": ["action"]
263
+ }
264
+ }
265
+ }
266
+ ]