Text Generation
Transformers
English
Spanish
iot
smart-home
home-assistant
tool-calling
agent
cybersecurity
devops
ollama
inference-pipeline
autonomous-agent
raspberry-pi
qwen3-coder
Instructions to use MrMoz33/tokioai-coder-iot with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use MrMoz33/tokioai-coder-iot with Transformers:
# Use a pipeline as a high-level helper from transformers import pipeline pipe = pipeline("text-generation", model="MrMoz33/tokioai-coder-iot")# pip install -U transformers accelerate # Load model directly from transformers import AutoModel model = AutoModel.from_pretrained("MrMoz33/tokioai-coder-iot", device_map="auto") - Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- vLLM
How to use MrMoz33/tokioai-coder-iot with vLLM:
Install from pip and serve model
# Install vLLM from pip: pip install vllm # Start the vLLM server: vllm serve "MrMoz33/tokioai-coder-iot" # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:8000/v1/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "MrMoz33/tokioai-coder-iot", "prompt": "Once upon a time,", "max_tokens": 512, "temperature": 0.5 }'Use Docker
docker model run hf.co/MrMoz33/tokioai-coder-iot
- SGLang
How to use MrMoz33/tokioai-coder-iot with SGLang:
Install from pip and serve model
# Install SGLang from pip: pip install sglang # Start the SGLang server: python3 -m sglang.launch_server \ --model-path "MrMoz33/tokioai-coder-iot" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "MrMoz33/tokioai-coder-iot", "prompt": "Once upon a time,", "max_tokens": 512, "temperature": 0.5 }'Use Docker images
docker run --gpus all \ --shm-size 32g \ -p 30000:30000 \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=<secret>" \ --ipc=host \ lmsysorg/sglang:latest \ python3 -m sglang.launch_server \ --model-path "MrMoz33/tokioai-coder-iot" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "MrMoz33/tokioai-coder-iot", "prompt": "Once upon a time,", "max_tokens": 512, "temperature": 0.5 }' - Docker Model Runner
How to use MrMoz33/tokioai-coder-iot with Docker Model Runner:
docker model run hf.co/MrMoz33/tokioai-coder-iot
TokioAI-Coder IoT Package v1.0 -- Engine + IoT System Prompt + Tools
Browse files- LICENSE +21 -0
- MANIFEST.in +5 -0
- README.md +242 -0
- config/example.env +42 -0
- engine/__init__.py +13 -0
- engine/__main__.py +224 -0
- engine/executor.py +331 -0
- engine/guided.py +228 -0
- engine/memory.py +353 -0
- engine/pipeline.py +476 -0
- engine/router.py +382 -0
- engine/schemas.py +497 -0
- engine/server.py +864 -0
- engine/test_engine.py +220 -0
- engine/verifier.py +186 -0
- examples/basic_usage.py +72 -0
- examples/custom_tools.py +36 -0
- examples/openai_client.py +51 -0
- prompts/iot_system_prompt.md +216 -0
- pyproject.toml +20 -0
- setup.py +52 -0
- tools/tool_definitions.json +266 -0
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 |
+
]
|