agent.py

#!/usr/bin/env python3
"""
OpenWinBot LLM agent — closes the perception-action loop using Claude.

The agent subscribes to win-observer (PUB), formats the semantic UI state into
a structured prompt, calls Claude with tool definitions, and relays each tool
call to win-actuator (REP). It loops until Claude emits a final answer or
the maximum iteration count is reached.

Requires:
    pip install anthropic pyzmq

Environment:
    ANTHROPIC_API_KEY must be set.

Standalone usage:
    python agent.py --window "Calculator" --task "Compute 9 + 3"
    python agent.py --window "Notepad"    --task "Type 'Hello World'"
    python agent.py --window "Paint"      --task "Describe the toolbar buttons"

Imported usage (from test_suite.py):
    from agent import run_agent
    result = run_agent(window="Calculator", task="Compute 9 + 3")
    print(result["summary"])
"""

import argparse
import json
import sys
import time
from typing import Any

import zmq
import anthropic

# ── Defaults ──────────────────────────────────────────────────────────────────

DEFAULT_OBS  = "tcp://localhost:5555"
DEFAULT_ACT  = "tcp://localhost:5556"
DEFAULT_MODEL = "claude-sonnet-4-6"
RECV_TIMEOUT_MS = 8_000

# ── Tool definitions (sent to Claude) ────────────────────────────────────────

TOOLS = [
    {
        "name": "click",
        "description": (
            "Click at an absolute screen position. "
            "Use this when you have the exact pixel coordinates. "
            "For named elements prefer click_id — it is more reliable."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "x": {"type": "integer", "description": "Screen X pixel"},
                "y": {"type": "integer", "description": "Screen Y pixel"},
            },
            "required": ["x", "y"],
        },
    },
    {
        "name": "click_id",
        "description": (
            "Click a UI element by its stable ID from the state message. "
            "The actuator resolves the centre automatically from its local state cache. "
            "This is the preferred way to click named elements — copy the id field exactly."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "id": {"type": "string", "description": "Element id (16-char hex from the state message)"},
            },
            "required": ["id"],
        },
    },
    {
        "name": "type_text",
        "description": "Type a string of text into the focused window using Unicode keystrokes.",
        "input_schema": {
            "type": "object",
            "properties": {
                "text": {"type": "string", "description": "Text to type"},
            },
            "required": ["text"],
        },
    },
    {
        "name": "read_state",
        "description": (
            "Fetch the latest UI state from win-observer. "
            "Call this after every action to observe its effect before deciding what to do next."
        ),
        "input_schema": {"type": "object", "properties": {}, "required": []},
    },
]

# ── State formatting ──────────────────────────────────────────────────────────

def format_state(state: dict) -> str:
    """
    Render a StateMessage dict as a compact, human-readable block suitable
    for inclusion in an LLM prompt. Separates informational from actionable
    elements and includes element IDs so Claude can use click_id directly.
    """
    elements = state.get("elements", [])
    window   = state.get("window", "?")
    ts       = state.get("timestamp", 0)

    info   = [e for e in elements if e.get("role") == "informational"]
    action = [e for e in elements if e.get("role") == "actionable"]
    other  = [e for e in elements if e.get("role") not in ("informational", "actionable")]

    lines = [f"Window: {window!r}  (timestamp={ts})"]

    if info:
        lines.append("\n── DISPLAY / STATUS ──────────────────────────────")
        for e in info:
            flag = " [DISABLED]" if not e.get("is_enabled", True) else ""
            lines.append(f"  {e['type']:<14} {e.get('name', '')!r}{flag}")

    if action:
        lines.append("\n── ACTIONABLE ELEMENTS ───────────────────────────")
        lines.append(f"  {'ID':>8}  {'TYPE':<14}  {'NAME':<28}  CENTER")
        for e in action:
            r   = e["rect"]
            cx  = r["x"] + r["w"] // 2
            cy  = r["y"] + r["h"] // 2
            dis = " [DISABLED]" if not e.get("is_enabled", True) else ""
            lines.append(
                f"  {e['id'][:8]}  {e['type']:<14}  {e.get('name', ''):<28}  ({cx},{cy}){dis}"
            )

    if other and other:
        lines.append(f"\n  (+{len(other)} structural elements omitted)")

    return "\n".join(lines)


# ── ZMQ helpers ───────────────────────────────────────────────────────────────

def recv_full_state(sub: zmq.Socket, max_attempts: int = 10) -> dict:
    """
    Read from the observer SUB socket until we get a 'full' state message.
    Delta messages (if --delta is active on the observer) are skipped here
    because the agent always wants a complete snapshot.
    """
    for _ in range(max_attempts):
        try:
            raw  = sub.recv()
            data = json.loads(raw)
            if data.get("type", "full") == "full":
                return data
        except zmq.Again:
            raise TimeoutError(
                f"Timed out waiting for a state frame from the observer.\n"
                f"Is win-observer running and publishing to {DEFAULT_OBS}?"
            )
    raise RuntimeError("Received only delta messages — restart observer without --delta "
                       "or wait for the next periodic full-frame broadcast.")


def send_action(req: zmq.Socket, action: dict) -> dict:
    req.send_string(json.dumps(action))
    raw = req.recv()
    return json.loads(raw)


# ── Tool execution ────────────────────────────────────────────────────────────

def execute_tool(name: str, inputs: dict,
                 req: zmq.Socket, sub: zmq.Socket) -> Any:
    if name == "click":
        result = send_action(req, {"action": "click", "x": inputs["x"], "y": inputs["y"]})
        time.sleep(0.3)   # let the UI settle
        return result

    elif name == "click_id":
        result = send_action(req, {"action": "click_id", "id": inputs["id"]})
        time.sleep(0.3)
        return result

    elif name == "type_text":
        result = send_action(req, {"action": "type", "text": inputs["text"]})
        time.sleep(0.2)
        return result

    elif name == "read_state":
        state = recv_full_state(sub)
        return {"state": format_state(state)}

    return {"error": f"unknown tool: {name}"}


# ── System prompt ─────────────────────────────────────────────────────────────

SYSTEM_PROMPT = """\
You are a Windows desktop automation agent powered by the OpenWinBot framework.

You receive a structured view of a window's UI elements and can interact with
them using the provided tools.

Rules:
1. Only interact with elements where role is "actionable" and is_enabled is true.
2. Prefer click_id over click — paste the element's id field exactly as shown.
3. After every action call read_state to observe the result before deciding next step.
4. When the task is complete, describe exactly what you did and what the final
   state of the relevant UI element shows (e.g., what number is on the display).
5. If an element you need is not visible, say so clearly rather than guessing.
"""


# ── Main agent loop ───────────────────────────────────────────────────────────

def run_agent(
    task:      str,
    window:    str,
    obs_addr:  str  = DEFAULT_OBS,
    act_addr:  str  = DEFAULT_ACT,
    model:     str  = DEFAULT_MODEL,
    max_iters: int  = 12,
    verbose:   bool = True,
) -> dict:
    """
    Run one agent session for the given task on the given window.

    Returns:
        {"success": bool, "summary": str, "iterations": int}
    """
    client = anthropic.Anthropic()

    ctx = zmq.Context()

    sub = ctx.socket(zmq.SUB)
    sub.setsockopt(zmq.RCVTIMEO, RECV_TIMEOUT_MS)
    sub.setsockopt_string(zmq.SUBSCRIBE, "")
    sub.connect(obs_addr)

    req = ctx.socket(zmq.REQ)
    req.setsockopt(zmq.RCVTIMEO, RECV_TIMEOUT_MS)
    req.connect(act_addr)

    if verbose:
        print(f"\n{'='*60}")
        print(f"  Task   : {task}")
        print(f"  Window : {window}")
        print(f"  Model  : {model}")
        print(f"{'='*60}\n")

    try:
        # ── Initial perception ────────────────────────────────────────────────
        if verbose: print("[agent] Fetching initial state...")
        state = recv_full_state(sub)
        state_text = format_state(state)

        if verbose:
            print("[agent] State received:\n" + state_text + "\n")

        messages = [
            {
                "role": "user",
                "content": (
                    f"Task: {task}\n\n"
                    f"Current UI state:\n{state_text}\n\n"
                    "Use the tools to complete the task. "
                    "Call read_state after each action to verify progress."
                ),
            }
        ]

        # ── Agentic loop ──────────────────────────────────────────────────────
        for iteration in range(1, max_iters + 1):
            if verbose: print(f"[agent] Iteration {iteration}/{max_iters}...")

            response = client.messages.create(
                model=model,
                max_tokens=2048,
                system=SYSTEM_PROMPT,
                tools=TOOLS,
                messages=messages,
            )

            messages.append({"role": "assistant", "content": response.content})

            # ── Task complete ─────────────────────────────────────────────────
            if response.stop_reason == "end_turn":
                summary = next(
                    (b.text for b in response.content if hasattr(b, "text")), ""
                )
                if verbose:
                    print(f"\n[agent] ✓ Done after {iteration} iteration(s).")
                    print(f"[agent] Summary: {summary}\n")
                return {"success": True, "summary": summary, "iterations": iteration}

            # ── Execute tool calls ────────────────────────────────────────────
            if response.stop_reason == "tool_use":
                tool_results = []

                for block in response.content:
                    if not hasattr(block, "type") or block.type != "tool_use":
                        continue

                    if verbose:
                        print(f"[agent] Tool call: {block.name}({json.dumps(block.input)})")

                    result = execute_tool(block.name, block.input, req, sub)

                    if verbose:
                        print(f"[agent] Tool result: {json.dumps(result)}")

                    tool_results.append({
                        "type":        "tool_result",
                        "tool_use_id": block.id,
                        "content":     json.dumps(result),
                    })

                messages.append({"role": "user", "content": tool_results})
                continue

            # Unexpected stop reason
            break

        summary = f"Reached max iterations ({max_iters}) without completing task."
        if verbose: print(f"[agent] ✗ {summary}")
        return {"success": False, "summary": summary, "iterations": max_iters}

    except TimeoutError as e:
        msg = str(e)
        if verbose: print(f"[agent] ✗ Timeout: {msg}")
        return {"success": False, "summary": msg, "iterations": 0}

    finally:
        sub.close()
        req.close()
        ctx.term()


# ── CLI entry point ───────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="OpenWinBot LLM agent")
    parser.add_argument("--window", "-w", required=True, help="Target window title")
    parser.add_argument("--task",   "-t", required=True, help="Task for the agent to complete")
    parser.add_argument("--obs",          default=DEFAULT_OBS, help="Observer PUB address")
    parser.add_argument("--act",          default=DEFAULT_ACT, help="Actuator REP address")
    parser.add_argument("--model",        default=DEFAULT_MODEL)
    parser.add_argument("--max-iters",    type=int, default=12)
    args = parser.parse_args()

    result = run_agent(
        task      = args.task,
        window    = args.window,
        obs_addr  = args.obs,
        act_addr  = args.act,
        model     = args.model,
        max_iters = args.max_iters,
    )
    sys.exit(0 if result["success"] else 1)


if __name__ == "__main__":
    main()