agent-runtime/agent/nodes/thinker_v2.py

"""Thinker Node v2: pure executor — runs tools as directed by Director."""

import asyncio
import json
import logging

from .base import Node
from ..llm import llm_call
from ..process import ProcessManager
from ..types import Command, DirectorPlan, ThoughtResult

log = logging.getLogger("runtime")


class ThinkerV2Node(Node):
    name = "thinker_v2"
    model = "google/gemini-2.0-flash-001"  # Fast model — just executes
    max_context_tokens = 4000

    RESPONSE_SYSTEM = """You are the Thinker — a fast executor in a cognitive runtime.
The Director (a smart model) already decided what to do. You just executed the tools.
Now write a natural response to the user based on the results.

{hint}

Rules:
- Be concise and natural.
- If tool results contain data, summarize it clearly.
- NEVER apologize. NEVER say "I" — you are part of a team.
- Keep it short: 1-3 sentences for simple responses.
- For data: reference the numbers, don't repeat raw output."""

    DB_HOST = "mariadb-eras"
    DB_USER = "root"
    DB_PASS = "root"

    def __init__(self, send_hud, process_manager: ProcessManager = None):
        super().__init__(send_hud)
        self.pm = process_manager

    def _run_db_query(self, query: str, database: str = "eras2_production") -> str:
        """Execute SQL query against MariaDB."""
        import pymysql
        trimmed = query.strip().upper()
        if not (trimmed.startswith("SELECT") or trimmed.startswith("DESCRIBE") or trimmed.startswith("SHOW")):
            return "Error: Only SELECT/DESCRIBE/SHOW queries allowed"
        if database not in ("eras2_production", "plankiste_test"):
            return f"Error: Unknown database '{database}'"
        conn = pymysql.connect(host=self.DB_HOST, user=self.DB_USER,
                               password=self.DB_PASS, database=database,
                               connect_timeout=5, read_timeout=15)
        try:
            with conn.cursor() as cur:
                cur.execute(query)
                rows = cur.fetchall()
                if not rows:
                    return "(no results)"
                cols = [d[0] for d in cur.description]
                lines = ["\t".join(cols)]
                for row in rows:
                    lines.append("\t".join(str(v) if v is not None else "" for v in row))
                return "\n".join(lines)
        finally:
            conn.close()

    async def process(self, command: Command, plan: DirectorPlan,
                      history: list[dict], memory_context: str = "") -> ThoughtResult:
        """Execute Director's plan and produce ThoughtResult."""
        await self.hud("thinking", detail=f"executing plan: {plan.goal}")

        actions = []
        state_updates = {}
        display_items = []
        machine_ops = []
        tool_used = ""
        tool_output = ""

        # Execute tool_sequence in order
        for step in plan.tool_sequence:
            tool = step.get("tool", "")
            args = step.get("args", {})
            await self.hud("tool_call", tool=tool, args=args)

            if tool == "emit_actions":
                actions.extend(args.get("actions", []))
            elif tool == "set_state":
                key = args.get("key", "")
                if key:
                    state_updates[key] = args.get("value")
            elif tool == "emit_display":
                display_items.extend(args.get("items", []))
            elif tool == "create_machine":
                machine_ops.append({"op": "create", **args})
            elif tool == "add_state":
                machine_ops.append({"op": "add_state", **args})
            elif tool == "reset_machine":
                machine_ops.append({"op": "reset", **args})
            elif tool == "destroy_machine":
                machine_ops.append({"op": "destroy", **args})
            elif tool == "query_db":
                query = args.get("query", "")
                database = args.get("database", "eras2_production")
                try:
                    result = await asyncio.to_thread(self._run_db_query, query, database)
                    tool_used = "query_db"
                    tool_output = result
                    await self.hud("tool_result", tool="query_db", output=result[:200])
                except Exception as e:
                    tool_used = "query_db"
                    tool_output = f"Error: {e}"
                    await self.hud("tool_result", tool="query_db", output=str(e)[:200])

        # Generate text response
        hint = plan.response_hint or f"Goal: {plan.goal}"
        if tool_output:
            hint += f"\nTool result:\n{tool_output[:500]}"

        messages = [
            {"role": "system", "content": self.RESPONSE_SYSTEM.format(hint=hint)},
        ]
        for msg in history[-8:]:
            messages.append(msg)
        messages.append({"role": "user", "content": command.source_text})
        messages = self.trim_context(messages)

        response = await llm_call(self.model, messages)
        if not response:
            response = "[no response]"

        await self.hud("decided", instruction=response[:200])

        return ThoughtResult(
            response=response,
            tool_used=tool_used,
            tool_output=tool_output,
            actions=actions,
            state_updates=state_updates,
            display_items=display_items,
            machine_ops=machine_ops,
        )