TokioAI-Coder IoT Package v1.0 -- Engine + IoT System Prompt + Tools
4b97c3b verified about 1 month ago
37.4 kB
"""
TokioAI Engine v3.0 -- OpenAI-Compatible API Server
SERVER-SIDE MULTI-STEP INVESTIGATION (the Opus architecture)
The key breakthrough: when the model calls a read-only tool (execute_local,
read_file, search_files, diagnose), the ENGINE executes it directly and feeds
the result back to the model -- all within a SINGLE API request. This means:
1. User asks "is Home Assistant running?"
2. Model calls execute_local(systemctl status home-assistant)
3. Engine EXECUTES it -> "not found"
4. Engine feeds result back to model
5. Model tries execute_local(docker ps | grep -i home)
6. Engine EXECUTES it -> "not found"
7. Model tries execute_local(ss -tlnp | grep 8123)
8. Engine EXECUTES it -> port 8123 listening
9. Model tries execute_local(curl -s http://localhost:8123/api/)
10. Engine EXECUTES it -> HA API responds
11. Model answers: "Yes, HA is running on port 8123 via Docker"
All 10 steps happen in ONE API call. The CLI sees only the final answer.
For WRITE tools (write_file, edit_file), the Engine delegates to the CLI.
Endpoints:
  POST /v1/chat/completions  -- Standard chat completion (with tool support)
  GET  /v1/models            -- List available models
  GET  /health               -- Health check
  GET  /stats                -- Engine statistics
"""

import json
import os
import sys
import time
import uuid
import argparse
import re
from typing import Dict, List, Optional, Tuple

from .pipeline import Pipeline, EngineConfig, PipelineResult
from .router import classify_intent
from .executor import execute_tool, is_safe_tool, SAFE_TOOLS, DELEGATE_TOOLS


def has_fake_tool_execution(text: str) -> bool:
    """Detect if text contains hallucinated/fake tool execution.
    
    ANY bash/shell code block with a system command means the model is 
    writing commands in text instead of using real tool_calls. This is 
    ALWAYS wrong -- the model should call execute_local, not print commands.
    """
    if not text:
        return False
    t = text.lower()
    # Pattern: ANY code block with a system command = fake tool execution
    # The model should be using tool_calls, not writing bash blocks
    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):
        return True
    if 'ejecutemos' in t or 'let me run' in t or 'voy a ejecutar' in t or 'voy a verificar' in t:
        if '```' in t:
            return True
    # Pattern: any code block immediately followed by "Result:" with fake data
    if re.search(r'```[^`]*```\s*\n*result\s*:', t):
        return True
    return False


def clean_model_output(text: str) -> str:
    """Remove model artifacts and hallucinated tool execution from output text."""
    if not text:
        return text
    # Remove trailing <tool_call> or [TOOL_CALLS] artifacts
    text = re.sub(r'</?tool_call>\s*$', '', text)
    text = re.sub(r'\[TOOL_CALLS?\]\s*$', '', text)
    text = re.sub(r'\[/TOOL_CALLS?\]\s*$', '', text)
    # Remove partial JSON tool call attempts at the end
    text = re.sub(r'\{"name"\s*:\s*"[^"]*"\s*,?\s*("arguments".*)?$', '', text)
    # Remove markdown-wrapped empty tool calls
    text = re.sub(r'```json\s*\{[^}]*\}\s*```\s*$', '', text)
    
    # Remove ALL code blocks containing system commands + fake output
    # This catches: ```bash\nss -tuln\n```\n\nResult: ...\n```...```
    text = re.sub(
        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[^`]*```',
        '', text, flags=re.MULTILINE | re.DOTALL
    )
    
    # Remove fake "Result:" sections (with or without code blocks)
    text = re.sub(
        r'(?:Result|Output|Salida)\s*:\s*(?:Active|Proto|tcp|udp|LISTEN|Filesystem|total|Mem:|Swap:).*?(?=\n\n[A-Z]|\Z)',
        '', text, flags=re.DOTALL | re.IGNORECASE
    )
    
    # Remove hallucinated inline tool execution blocks
    # Pattern: ```bash\nCOMMAND\n```\n\nResult of tool:\nFAKE_OUTPUT
    text = re.sub(
        r'```(?:bash|shell|sh)\n[^`]+```\s*\n*(?:Result of tool:?\s*\n[^\n]*(?:\n[^\n]*)*?)(?=\n\n|\Z)',
        '', text, flags=re.MULTILINE
    )
    # Pattern: bare ```\nCOMMAND\n``` blocks (no language tag)
    text = re.sub(
        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```',
        '', text, flags=re.MULTILINE
    )
    # Pattern: "I ran: tool_name(...)\n[OK] Result:\n..."
    text = re.sub(
        r'I ran:\s*\w+\([^)]*\)\s*\n\[(?:OK|FAIL(?:ED)?)\]\s*Result:\n.*?(?=\n\n|\Z)',
        '', text, flags=re.DOTALL
    )
    # Pattern: standalone "Result of tool:" blocks
    text = re.sub(
        r'Result of tool:\s*\n.*?(?=\n\n|\Z)',
        '', text, flags=re.DOTALL
    )
    # Remove trailing lazy/filler phrases
    lazy_tails = [
        r'\n+(?:Let me know|If you\'?d? like|Do you want|Quieres que|Te gustaría|Si necesitas|Para más detalles).*$',
        r'\n+(?:Would you like me to|Want me to|If you want me to|Si quieres que).*$',
    ]
    for pattern in lazy_tails:
        text = re.sub(pattern, '', text, flags=re.IGNORECASE | re.DOTALL)
    
    # Clean up excessive blank lines left behind
    text = re.sub(r'\n{3,}', '\n\n', text)
    return text.strip()


# Max number of server-side investigation rounds per request
MAX_INVESTIGATION_DEPTH = 8

# Max total time for server-side investigation (seconds)
MAX_INVESTIGATION_TIME = 90


def count_tool_depth(messages: List[Dict]) -> int:
    """Count tool call rounds in messages."""
    depth = 0
    for msg in messages:
        if msg.get("role") == "assistant" and msg.get("tool_calls"):
            depth += 1
    return depth


def find_original_query(messages: List[Dict]) -> str:
    """Find the first user message (the original question)."""
    for msg in messages:
        if msg.get("role") == "user":
            content = msg.get("content", "")
            if content and not content.startswith("INVESTIGATION"):
                return content
    return ""


def create_app(config: EngineConfig = None):
    """Create the Flask app with the engine pipeline."""
    
    try:
        from flask import Flask, request, jsonify, Response
    except ImportError:
        print("Flask not installed. Install with: pip install flask")
        sys.exit(1)

    app = Flask("tokioai-engine")
    engine_config = config or EngineConfig()
    engine = Pipeline(engine_config)

    # Track investigation stats
    investigation_stats = {
        "total_investigations": 0,
        "total_steps": 0,
        "max_depth_reached": 0,
        "delegated_to_cli": 0,
        "avg_steps": 0,
    }

    @app.route("/health", methods=["GET"])
    def health():
        return jsonify({
            "status": "ok",
            "engine": "tokioai-engine",
            "version": "3.0.0",
            "model": engine.config.model_name,
            "ollama_host": engine.config.ollama_host,
            "max_investigation_depth": MAX_INVESTIGATION_DEPTH,
            "max_investigation_time": MAX_INVESTIGATION_TIME,
            "server_side_execution": True,
        })

    @app.route("/stats", methods=["GET"])
    def stats():
        s = engine.get_stats()
        s["investigation"] = investigation_stats
        return jsonify(s)

    @app.route("/v1/models", methods=["GET"])
    def list_models():
        return jsonify({
            "object": "list",
            "data": [
                {
                    "id": "tokioai-engine",
                    "object": "model",
                    "created": int(time.time()),
                    "owned_by": "tokioai",
                    "permission": [],
                    "root": engine.config.model_name,
                    "parent": None,
                }
            ]
        })

    @app.route("/v1/chat/completions", methods=["POST"])
    def chat_completions():
        data = request.get_json()

        if not data:
            return jsonify({"error": {"message": "Request body is required"}}), 400

        messages = data.get("messages", [])
        if not messages:
            return jsonify({"error": {"message": "messages is required"}}), 400

        stream = data.get("stream", False)

        # ============================================================
        # v3.0 SERVER-SIDE INVESTIGATION LOOP
        # ============================================================
        #
        # Architecture:
        # 1. Extract user message and conversation history
        # 2. Call model (via Pipeline) -- model returns tool_calls or text
        # 3. If tool_calls:
        #    a. If ALL tools are SAFE (read-only): execute server-side,
        #       feed results back to model, goto step 2
        #    b. If ANY tool is DELEGATE (write/remote): return tool_calls
        #       to CLI for execution (old behavior)
        # 4. If text: return the text answer
        # 5. Repeat up to MAX_INVESTIGATION_DEPTH or MAX_INVESTIGATION_TIME
        
        # Determine if this is a tool-result continuation from CLI
        last_msg_role = None
        for m in reversed(messages):
            if m.get("role") != "system":
                last_msg_role = m.get("role")
                break
        
        is_tool_continuation = (last_msg_role == "tool")
        
        # Build the working message history
        user_msg, working_messages = _extract_messages(messages)
        
        if not user_msg:
            return jsonify({"error": {"message": "No user message found"}}), 400
        
        original_query = user_msg
        
        # If this is a CLI tool continuation (CLI executed a write/remote tool),
        # let the model see the result and decide next step
        if is_tool_continuation:
            # The CLI sent back results from a delegated tool.
            # Feed it all through and let the model continue.
            pass
        
        # === PRE-CLASSIFICATION: Skip investigation for pure text ===
        # For greetings, thanks -- skip the investigation loop entirely.
        # BUT NOT for affirmations ("dale", "si", "ok") which are contextual
        # and need conversation history to understand what the user is agreeing to.
        _GREETING_PATTERNS = [
            r'^(hi|hello|hey|hola|que\s+tal|buenos?\s+d[ií]as?|buenas?\s+(?:tardes?|noches?))\s*[.!?]*$',
            r'^(thanks?|thank\s+you|gracias|thx|de\s+nada)\s*[.!?]*$',
            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*[.!?]*$',
            r'^(todo\s+bien|bien\s+y\s+tu|genial|cool|nice|great|awesome)\s*[.!?]*$',
            r'^(que\s+onda|que\s+hay|que\s+pasa|que\s+cuentas)\s*[.!?]*$',
            r'\b(who\s+are\s+you|what\s+can\s+you|tell\s+me\s+about\s+yourself)\b',
        ]
        import re as _re
        _is_pure_greeting = any(_re.match(p, user_msg.lower().strip(), _re.IGNORECASE) for p in _GREETING_PATTERNS)
        if not is_tool_continuation and _is_pure_greeting:
                _shortcut_start = time.time()
                if engine_config.verbose:
                    print(f"[ENGINE] Pure text shortcut: '{user_msg}' -> greeting/thanks")
                engine.reset()
                result = engine.process(
                    user_message=user_msg,
                    conversation_history=working_messages,
                    force_intent="text",
                )
                elapsed = time.time() - _shortcut_start
                return jsonify({
                    "id": f"chatcmpl-{uuid.uuid4().hex[:8]}",
                    "object": "chat.completion",
                    "created": int(time.time()),
                    "model": "tokioai-engine",
                    "choices": [{
                        "index": 0,
                        "message": {
                            "role": "assistant",
                            "content": clean_model_output(result.text or ""),
                        },
                        "finish_reason": "stop",
                    }],
                    "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
                    "engine_metadata": {
                        "investigation_depth": 0,
                        "investigation_log": [],
                        "elapsed_s": round(elapsed, 1),
                        "shortcut": "pure_text",
                    },
                })

        # === THE INVESTIGATION LOOP ===
        engine.reset()
        
        investigation_log = []  # Track what we tried
        start_time = time.time()
        final_result = None
        delegated_tool_calls = None
        depth = 0
        
        for depth in range(MAX_INVESTIGATION_DEPTH):
            elapsed = time.time() - start_time
            if elapsed > MAX_INVESTIGATION_TIME:
                # Time limit -- force text summary
                if engine_config.verbose:
                    print(f"[ENGINE] Investigation time limit reached ({elapsed:.1f}s)")
                break
            
            # Build the prompt for this round
            if depth == 0 and not is_tool_continuation:
                # First round: use original query
                current_msg = user_msg
                force_intent = None
            elif depth == 0 and is_tool_continuation:
                # CLI sent back tool results -- summarize and continue
                current_msg = _build_continuation_prompt(messages, original_query)
                force_intent = None
            else:
                # Investigation round N > 0: tell model what happened
                current_msg = _build_investigation_prompt(
                    original_query, investigation_log, depth
                )
                # For simple queries (1 success is enough), let model decide
                # For ongoing investigation, force tool calls until near the limit
                successes_so_far = sum(1 for e in investigation_log if e["success"])
                if depth >= MAX_INVESTIGATION_DEPTH - 2:
                    force_intent = "text"  # Near limit: force final answer
                elif successes_so_far >= 1 and depth <= 2:
                    force_intent = None  # Let model decide (may answer or investigate more)
                else:
                    force_intent = None  # Let model decide
            
            # Call the model
            try:
                result = engine.process(
                    user_message=current_msg,
                    conversation_history=working_messages,
                    force_intent=force_intent,
                )
            except Exception as e:
                if engine_config.verbose:
                    print(f"[ENGINE] Model error at depth {depth}: {e}")
                final_result = PipelineResult(
                    output_type="text",
                    text=f"Error during investigation: {e}",
                )
                break
            
            if engine_config.verbose:
                if result.tool_calls:
                    names = [tc["name"] for tc in result.tool_calls]
                    print(f"[ENGINE] Depth {depth}: tool_calls={names}")
                else:
                    print(f"[ENGINE] Depth {depth}: text response ({len(result.text)} chars)")
            
            # === MODEL RETURNED TEXT ===
            if result.output_type == "text" or not result.tool_calls:
                raw_text = result.text or ""
                text = raw_text.lower()
                
                # CRITICAL: Detect FAKE tool execution (hallucinated command output)
                if has_fake_tool_execution(raw_text) and depth < MAX_INVESTIGATION_DEPTH - 1:
                    if engine_config.verbose:
                        print(f"[ENGINE] FAKE tool execution detected at depth {depth}, forcing real tool call")
                    working_messages.append({
                        "role": "user",
                        "content": (
                            f"STOP. You just simulated a command in text instead of actually "
                            f"running it. That output is FAKE. Use a real tool_call to execute "
                            f"the command. Call execute_local with the appropriate command. "
                            f"Question: {original_query}"
                        ),
                    })
                    continue  # Force another round with real tool
                
                # Detect LAZY/INCOMPLETE answers and force more investigation
                lazy_patterns = [
                    "you should", "te recomiendo", "se recomienda",
                    "requires further", "requiere.*más profund",
                    "necesitaría", "se necesitaría", "se necesita",
                    "want me to check", "quieres que revise",
                    "te gustaría que", "would you like me to",
                    "para evaluar", "para determinar",
                    "no se detectaron vulnerabilidades inmediatas",
                    "you could check", "podrias revisar",
                    "voy a verificar", "voy a ejecutar",
                    "ejecutemos", "let me run", "let me check",
                ]
                is_lazy = any(p in text for p in lazy_patterns)
                successes_so_far = sum(1 for e in investigation_log if e["success"])
                
                if is_lazy and successes_so_far < 4 and depth < MAX_INVESTIGATION_DEPTH - 2:
                    if engine_config.verbose:
                        print(f"[ENGINE] Lazy answer detected at depth {depth}, pushing for more investigation")
                    working_messages.append({"role": "assistant", "content": raw_text})
                    working_messages.append({
                        "role": "user",
                        "content": (
                            f"Your answer was incomplete. You said things like 'you should check' "
                            f"or 'I will verify'. YOU are the expert -- don't describe what you "
                            f"will do, ACTUALLY DO IT with a tool call right now. "
                            f"Original question: {original_query}"
                        ),
                    })
                    continue  # Force another round
                
                # Clean the output before accepting it
                result.text = clean_model_output(raw_text)
                final_result = result
                break
            
            # === MODEL RETURNED TOOL CALLS ===
            # Check if all tools are safe for server-side execution
            all_safe = all(is_safe_tool(tc["name"]) for tc in result.tool_calls)
            
            if not all_safe:
                # At least one tool needs CLI execution (write/remote)
                # Return tool_calls to CLI, investigation pauses
                delegated_tool_calls = result
                investigation_stats["delegated_to_cli"] += 1
                if engine_config.verbose:
                    unsafe = [tc["name"] for tc in result.tool_calls if not is_safe_tool(tc["name"])]
                    print(f"[ENGINE] Delegating to CLI: {unsafe}")
                break
            
            # ALL tools are safe -- execute server-side!
            tool_results = []
            for tc in result.tool_calls:
                name = tc["name"]
                args = tc["arguments"]
                
                if engine_config.verbose:
                    arg_preview = args.get("command", args.get("path", args.get("pattern", args.get("target", str(args)))))
                    print(f"[ENGINE]   Executing: {name}({arg_preview})")
                
                success, output = execute_tool(name, args, timeout=30)
                
                tool_results.append({
                    "name": name,
                    "arguments": args,
                    "output": output,
                    "success": success,
                })
                
                investigation_log.append({
                    "name": name,
                    "arguments": args,
                    "output": output[:500],  # Keep log manageable
                    "success": success,
                    "depth": depth,
                })
                
                if engine_config.verbose:
                    status = "OK" if success else "FAIL"
                    print(f"[ENGINE]   Result [{status}]: {output[:100]}")
            
            # Feed tool results back to model using Ollama's native tool format
            # Ollama /api/chat expects: arguments as dict (not JSON string),
            # no "id"/"type" in tool_calls, no "tool_call_id" in tool messages
            
            # Add assistant message with tool_calls
            working_messages.append({
                "role": "assistant",
                "content": "",
                "tool_calls": [
                    {
                        "function": {
                            "name": tr["name"],
                            "arguments": tr["arguments"] if isinstance(tr["arguments"], dict) else json.loads(tr["arguments"]),
                        },
                    }
                    for tr in tool_results
                ],
            })
            
            # Add each tool result as a tool message
            for tr in tool_results:
                working_messages.append({
                    "role": "tool",
                    "content": tr["output"],
                })
        
        # === END OF INVESTIGATION LOOP ===
        
        # Update stats
        if depth > 0:
            investigation_stats["total_investigations"] += 1
            investigation_stats["total_steps"] += depth
            investigation_stats["max_depth_reached"] = max(
                investigation_stats["max_depth_reached"], depth
            )
            total = investigation_stats["total_investigations"]
            investigation_stats["avg_steps"] = (
                investigation_stats["total_steps"] / total if total > 0 else 0
            )
        
        # If we hit the depth/time limit without a final answer, force one
        if final_result is None and delegated_tool_calls is None:
            # Force the model to summarize
            summary_prompt = _build_summary_prompt(original_query, investigation_log)
            try:
                final_result = engine.process(
                    user_message=summary_prompt,
                    conversation_history=working_messages,
                    force_intent="text",
                )
            except:
                final_result = PipelineResult(
                    output_type="text",
                    text=_format_investigation_summary(original_query, investigation_log),
                )
        
        # Use delegated tool calls if that's what we have
        if delegated_tool_calls is not None:
            result_to_return = delegated_tool_calls
        else:
            result_to_return = final_result
        
        chat_id = f"chatcmpl-{uuid.uuid4().hex[:8]}"
        created = int(time.time())

        # --- STREAMING RESPONSE (SSE) ---
        if stream:
            def generate_sse():
                msg = _build_response_message(result_to_return)
                finish_reason = "tool_calls" if result_to_return.tool_calls else "stop"

                if result_to_return.tool_calls:
                    tc_deltas = []
                    for i, tc in enumerate(result_to_return.tool_calls):
                        tc_deltas.append({
                            "index": i,
                            "id": f"call_{uuid.uuid4().hex[:8]}",
                            "type": "function",
                            "function": {
                                "name": tc["name"],
                                "arguments": json.dumps(tc["arguments"]) if isinstance(tc["arguments"], dict) else tc["arguments"],
                            },
                        })
                    chunk = {
                        "id": chat_id,
                        "object": "chat.completion.chunk",
                        "created": created,
                        "model": "tokioai-engine",
                        "choices": [{
                            "index": 0,
                            "delta": {
                                "role": "assistant",
                                "content": None,
                                "tool_calls": tc_deltas,
                            },
                            "finish_reason": finish_reason,
                        }],
                    }
                    yield f"data: {json.dumps(chunk)}\n\n"
                else:
                    text = clean_model_output(result_to_return.text or "")

                    # First chunk: role
                    chunk = {
                        "id": chat_id,
                        "object": "chat.completion.chunk",
                        "created": created,
                        "model": "tokioai-engine",
                        "choices": [{
                            "index": 0,
                            "delta": {"role": "assistant", "content": ""},
                            "finish_reason": None,
                        }],
                    }
                    yield f"data: {json.dumps(chunk)}\n\n"

                    if text:
                        # Stream word by word
                        words = text.split(' ')
                        for i, word in enumerate(words):
                            token = word if i == 0 else ' ' + word
                            chunk = {
                                "id": chat_id,
                                "object": "chat.completion.chunk",
                                "created": created,
                                "model": "tokioai-engine",
                                "choices": [{
                                    "index": 0,
                                    "delta": {"content": token},
                                    "finish_reason": None,
                                }],
                            }
                            yield f"data: {json.dumps(chunk)}\n\n"

                    # Final chunk
                    chunk = {
                        "id": chat_id,
                        "object": "chat.completion.chunk",
                        "created": created,
                        "model": "tokioai-engine",
                        "choices": [{
                            "index": 0,
                            "delta": {},
                            "finish_reason": "stop",
                        }],
                    }
                    yield f"data: {json.dumps(chunk)}\n\n"

                yield "data: [DONE]\n\n"

            return Response(
                generate_sse(),
                mimetype="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "X-Accel-Buffering": "no",
                },
            )

        # --- NON-STREAMING RESPONSE ---
        response = {
            "id": chat_id,
            "object": "chat.completion",
            "created": created,
            "model": "tokioai-engine",
            "choices": [
                {
                    "index": 0,
                    "message": _build_response_message(result_to_return),
                    "finish_reason": "tool_calls" if result_to_return.tool_calls else "stop",
                }
            ],
            "usage": {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0,
            },
            "engine_metadata": {
                "investigation_depth": depth,
                "investigation_log": [
                    {"tool": e["name"], "success": e["success"]}
                    for e in investigation_log
                ],
                "elapsed_s": round(time.time() - start_time, 1),
            },
        }

        return jsonify(response)

    # ============================================================
    # HELPER FUNCTIONS
    # ============================================================
    
    def _extract_messages(messages: List[Dict]) -> Tuple[str, List[Dict]]:
        """
        Extract the user message and build Ollama-compatible history.
        Converts tool messages to assistant context.
        Returns (user_msg, ollama_messages).
        """
        user_msg = None
        history = []
        
        for msg in messages:
            role = msg.get("role", "")
            content = msg.get("content", "")
            
            if role == "system":
                history.append({"role": "system", "content": content})
                
            elif role == "user":
                if user_msg is not None:
                    history.append({"role": "user", "content": user_msg})
                user_msg = content
                
            elif role == "assistant":
                tool_calls = msg.get("tool_calls", [])
                if tool_calls:
                    # Convert OpenAI tool_calls format to Ollama format
                    # Ollama wants: arguments as dict, no id/type fields
                    ollama_tcs = []
                    for tc in tool_calls:
                        func = tc.get("function", {})
                        name = func.get("name", "")
                        raw_args = func.get("arguments", {})
                        if isinstance(raw_args, str):
                            try:
                                raw_args = json.loads(raw_args)
                            except (json.JSONDecodeError, TypeError):
                                raw_args = {}
                        ollama_tcs.append({
                            "function": {"name": name, "arguments": raw_args},
                        })
                    history.append({
                        "role": "assistant",
                        "content": content or "",
                        "tool_calls": ollama_tcs,
                    })
                elif content:
                    history.append({"role": "assistant", "content": content})
                    
            elif role == "tool":
                # Ollama tool messages: just role + content (no tool_call_id)
                history.append({
                    "role": "tool",
                    "content": content or "(no output)",
                })
        
        return user_msg, history
    
    def _build_continuation_prompt(messages: List[Dict], original_query: str) -> str:
        """Build prompt when CLI sends back results from delegated tools."""
        # Find the last tool result
        last_tool_result = ""
        last_tool_name = ""
        for msg in reversed(messages):
            if msg.get("role") == "tool":
                last_tool_result = msg.get("content", "")[:500]
                last_tool_name = msg.get("name", "tool")
                break
        
        return (
            f"The tool {last_tool_name} returned:\n{last_tool_result}\n\n"
            f"Original question: {original_query}\n\n"
            f"Based on this result, either:\n"
            f"1. Call more tools to continue investigating, OR\n"
            f"2. Provide your final answer if you have enough information."
        )
    
    def _build_investigation_prompt(
        original_query: str,
        investigation_log: List[Dict],
        depth: int,
    ) -> str:
        """
        Build a prompt for the next investigation round.
        
        v4.0 PHILOSOPHY: Let the model decide when it has enough data.
        Don't force early answers. Only push for final answer near depth limit.
        """
        failures = sum(1 for e in investigation_log if not e["success"])
        successes = sum(1 for e in investigation_log if e["success"])
        
        # Near the depth limit -- force final answer
        if depth >= MAX_INVESTIGATION_DEPTH - 2:
            return (
                f"FINAL ROUND. You have gathered data from {successes} successful "
                f"tool calls. Provide your FINAL answer to: {original_query}\n\n"
                f"Summarize ALL findings. Be specific -- report exact values, "
                f"ports, processes, permissions. No more tool calls."
            )
        
        # All failures so far -- try different approach
        if failures > 0 and successes == 0:
            return (
                f"The previous commands failed. Try a DIFFERENT approach "
                f"to answer: {original_query}\n"
                f"Use a different command. Do NOT repeat what already failed."
            )
        
        # Normal continuation -- let model freely choose
        already = [e["arguments"].get("command", e["arguments"].get("path", "?"))[:50] for e in investigation_log[-3:]]
        return (
            f"Original question: {original_query}\n"
            f"You've run {len(investigation_log)} commands ({successes} OK, {failures} failed).\n"
            f"Recent: {', '.join(already)}\n\n"
            f"If you have enough information, provide your FINAL answer in plain text.\n"
            f"If you need more data, call another tool with a DIFFERENT command."
        )
    
    def _build_summary_prompt(original_query: str, investigation_log: List[Dict]) -> str:
        """Build a forced-text prompt summarizing all investigation findings."""
        return (
            f"Based on all the tool results above, provide a CLEAR, DEFINITIVE "
            f"answer to: {original_query}\n\n"
            f"Rules for your response:\n"
            f"- Write in plain language, no command blocks\n"
            f"- State what IS working, what ISN'T, and why\n"
            f"- Be direct and specific\n"
            f"- Do NOT include any bash/shell commands or simulated tool output"
        )
    
    def _format_investigation_summary(original_query: str, investigation_log: List[Dict]) -> str:
        """Fallback text summary if model fails to summarize."""
        lines = [f"Investigation: {original_query}\n"]
        for entry in investigation_log:
            status = "OK" if entry["success"] else "FAIL"
            arg_str = entry["arguments"].get("command", entry["arguments"].get("path", str(entry["arguments"])))
            lines.append(f"  [{status}] {entry['name']}({arg_str})")
            if entry["output"]:
                lines.append(f"         -> {entry['output'][:150]}")
        return "\n".join(lines)

    def _build_response_message(result: PipelineResult) -> Dict:
        """Build OpenAI-style message from pipeline result."""
        if result.output_type == "tool_call" and result.tool_calls:
            return {
                "role": "assistant",
                "content": None,
                "tool_calls": [
                    {
                        "id": f"call_{uuid.uuid4().hex[:8]}",
                        "type": "function",
                        "function": {
                            "name": tc["name"],
                            "arguments": json.dumps(tc["arguments"]) if isinstance(tc["arguments"], dict) else tc["arguments"],
                        },
                    }
                    for tc in result.tool_calls
                ],
            }
        return {
            "role": "assistant",
            "content": clean_model_output(result.text or ""),
        }

    return app


def main():
    parser = argparse.ArgumentParser(description="TokioAI Engine Server v3.0")
    parser.add_argument("command", choices=["serve"], help="Server command")
    parser.add_argument("--port", type=int, default=8080, help="Server port")
    parser.add_argument("--host", type=str, default="127.0.0.1", help="Bind host")
    parser.add_argument("--ollama-host", type=str, default=None, help="Ollama API host")
    parser.add_argument("--model", type=str, default="tokioai-coder", help="Model name")
    parser.add_argument("--verbose", action="store_true", help="Verbose logging")

    args = parser.parse_args()

    config = EngineConfig(
        ollama_host=args.ollama_host,
        model_name=args.model,
        verbose=args.verbose,
    )

    app = create_app(config)
    print(f"TokioAI Engine v3.0 starting on {args.host}:{args.port}")
    print(f"  Model: {config.model_name}")
    print(f"  Ollama: {config.ollama_host}")
    print(f"  Max investigation depth: {MAX_INVESTIGATION_DEPTH}")
    print(f"  Max investigation time: {MAX_INVESTIGATION_TIME}s")
    print(f"  Server-side execution: ENABLED (tools: {', '.join(sorted(SAFE_TOOLS))})")
    print(f"  Verbose: {args.verbose}")

    app.run(host=args.host, port=args.port, threaded=True)


if __name__ == "__main__":
    main()