add main.py
This commit is contained in:
parent
bef85e8e73
commit
325a8fbfb7
264
easybot_windows_agent/main.py
Normal file
264
easybot_windows_agent/main.py
Normal file
@ -0,0 +1,264 @@
|
||||
"""
|
||||
EasyBot Windows Agent — main entry point.
|
||||
|
||||
Connects to the EasyBot runtime via WebSocket, receives tasks, and executes
|
||||
them using the `windows-use` library for GUI automation.
|
||||
|
||||
Configuration via environment variables:
|
||||
EASYBOT_SERVER WebSocket URL (e.g., wss://app.example.com/ws/windows-agent)
|
||||
EASYBOT_TOKEN Authentication token (eb_wag_...)
|
||||
LOG_LEVEL DEBUG | INFO | WARNING (default: INFO)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
import signal
|
||||
import sys
|
||||
import traceback
|
||||
from typing import Any
|
||||
|
||||
import websockets
|
||||
from websockets.client import WebSocketClientProtocol
|
||||
|
||||
from . import __version__
|
||||
|
||||
logger = logging.getLogger("easybot-windows-agent")
|
||||
|
||||
# ── Configuration ────────────────────────────────────────────────────────────
|
||||
|
||||
RECONNECT_DELAY_BASE = 5 # seconds
|
||||
RECONNECT_DELAY_MAX = 120 # seconds
|
||||
HEARTBEAT_TIMEOUT = 45 # seconds — must be > server's 15s interval
|
||||
|
||||
|
||||
# ── Agent Client ─────────────────────────────────────────────────────────────
|
||||
|
||||
class AgentClient:
|
||||
def __init__(self, server_url: str, token: str):
|
||||
self.server_url = server_url
|
||||
self.token = token
|
||||
self.ws: WebSocketClientProtocol | None = None
|
||||
self.running = True
|
||||
self.current_task_id: str | None = None
|
||||
self.cancel_event = asyncio.Event()
|
||||
self.llm_config: dict[str, Any] = {}
|
||||
self.max_steps = 25
|
||||
self._reconnect_delay = RECONNECT_DELAY_BASE
|
||||
|
||||
async def run(self) -> None:
|
||||
"""Main loop: connect → handle messages → reconnect on failure."""
|
||||
while self.running:
|
||||
try:
|
||||
await self._connect_and_handle()
|
||||
except (websockets.ConnectionClosed, OSError, ConnectionRefusedError) as e:
|
||||
logger.warning("Connection lost: %s — reconnecting in %ds", e, self._reconnect_delay)
|
||||
except Exception:
|
||||
logger.exception("Unexpected error — reconnecting in %ds", self._reconnect_delay)
|
||||
|
||||
if not self.running:
|
||||
break
|
||||
await asyncio.sleep(self._reconnect_delay)
|
||||
self._reconnect_delay = min(self._reconnect_delay * 2, RECONNECT_DELAY_MAX)
|
||||
|
||||
async def _connect_and_handle(self) -> None:
|
||||
"""Single connection lifecycle."""
|
||||
extra_headers = {"Authorization": f"Bearer {self.token}"}
|
||||
|
||||
async with websockets.connect(
|
||||
self.server_url,
|
||||
additional_headers=extra_headers,
|
||||
ping_interval=None, # we handle heartbeats manually
|
||||
max_size=16 * 1024 * 1024,
|
||||
close_timeout=10,
|
||||
) as ws:
|
||||
self.ws = ws
|
||||
self._reconnect_delay = RECONNECT_DELAY_BASE
|
||||
logger.info("Connected to %s", self.server_url)
|
||||
|
||||
# Send client info
|
||||
await self._send({
|
||||
"type": "client_info",
|
||||
"hostname": platform.node(),
|
||||
"os_version": f"{platform.system()} {platform.version()}",
|
||||
"python_version": platform.python_version(),
|
||||
"client_version": __version__,
|
||||
})
|
||||
|
||||
try:
|
||||
async for raw in ws:
|
||||
msg = json.loads(raw)
|
||||
await self._handle_message(msg)
|
||||
finally:
|
||||
self.ws = None
|
||||
self.current_task_id = None
|
||||
|
||||
async def _handle_message(self, msg: dict[str, Any]) -> None:
|
||||
"""Dispatch incoming messages from the server."""
|
||||
msg_type = msg.get("type")
|
||||
|
||||
if msg_type == "ping":
|
||||
await self._send({"type": "pong"})
|
||||
|
||||
elif msg_type == "config":
|
||||
self.llm_config = msg.get("llm", {})
|
||||
self.max_steps = msg.get("max_steps", 25)
|
||||
logger.info("Config received: provider=%s model=%s max_steps=%d",
|
||||
self.llm_config.get("provider"), self.llm_config.get("model"), self.max_steps)
|
||||
|
||||
elif msg_type == "task":
|
||||
task_id = msg["task_id"]
|
||||
instruction = msg["instruction"]
|
||||
logger.info("Task received: %s — %s", task_id, instruction[:100])
|
||||
# Run task in background so we keep processing messages
|
||||
asyncio.create_task(self._execute_task(task_id, instruction))
|
||||
|
||||
elif msg_type == "cancel":
|
||||
task_id = msg.get("task_id")
|
||||
if self.current_task_id == task_id:
|
||||
logger.info("Cancel requested for task %s", task_id)
|
||||
self.cancel_event.set()
|
||||
|
||||
else:
|
||||
logger.debug("Unknown message type: %s", msg_type)
|
||||
|
||||
async def _execute_task(self, task_id: str, instruction: str) -> None:
|
||||
"""Execute a windows-use task and report results."""
|
||||
self.current_task_id = task_id
|
||||
self.cancel_event.clear()
|
||||
|
||||
try:
|
||||
await self._send({"type": "task_started", "task_id": task_id})
|
||||
|
||||
result = await self._run_windows_use(task_id, instruction)
|
||||
|
||||
await self._send({
|
||||
"type": "task_completed",
|
||||
"task_id": task_id,
|
||||
"result": result.get("result", "Task completed"),
|
||||
"steps_count": result.get("steps_count", 0),
|
||||
})
|
||||
logger.info("Task %s completed successfully", task_id)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
await self._send({
|
||||
"type": "task_error",
|
||||
"task_id": task_id,
|
||||
"error": "Task was cancelled",
|
||||
})
|
||||
except Exception as e:
|
||||
error_msg = f"{type(e).__name__}: {e}"
|
||||
logger.error("Task %s failed: %s", task_id, error_msg)
|
||||
await self._send({
|
||||
"type": "task_error",
|
||||
"task_id": task_id,
|
||||
"error": error_msg,
|
||||
})
|
||||
finally:
|
||||
self.current_task_id = None
|
||||
|
||||
async def _run_windows_use(self, task_id: str, instruction: str) -> dict[str, Any]:
|
||||
"""Run the windows-use agent with the given instruction."""
|
||||
from windows_use import Agent, AgentConfig # type: ignore
|
||||
|
||||
config_kwargs: dict[str, Any] = {}
|
||||
if self.llm_config.get("api_key"):
|
||||
config_kwargs["api_key"] = self.llm_config["api_key"]
|
||||
if self.llm_config.get("model"):
|
||||
config_kwargs["model"] = self.llm_config["model"]
|
||||
if self.llm_config.get("provider"):
|
||||
config_kwargs["provider"] = self.llm_config["provider"]
|
||||
|
||||
config = AgentConfig(
|
||||
max_steps=self.max_steps,
|
||||
**config_kwargs,
|
||||
)
|
||||
|
||||
agent = Agent(config=config)
|
||||
step_count = 0
|
||||
|
||||
# Callback for step reporting
|
||||
def on_step(step_info: dict[str, Any]) -> None:
|
||||
nonlocal step_count
|
||||
step_count += 1
|
||||
asyncio.get_event_loop().call_soon_threadsafe(
|
||||
asyncio.ensure_future,
|
||||
self._send({
|
||||
"type": "task_step",
|
||||
"task_id": task_id,
|
||||
"step_number": step_count,
|
||||
"action": step_info.get("action", "step"),
|
||||
"details": step_info.get("details", ""),
|
||||
})
|
||||
)
|
||||
|
||||
# Run agent (windows-use is sync internally — run in thread)
|
||||
loop = asyncio.get_event_loop()
|
||||
result_text = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: agent.run(instruction)
|
||||
)
|
||||
|
||||
return {
|
||||
"result": str(result_text) if result_text else "Task completed",
|
||||
"steps_count": step_count,
|
||||
}
|
||||
|
||||
async def _send(self, msg: dict[str, Any]) -> None:
|
||||
"""Send a JSON message to the server."""
|
||||
if self.ws and self.ws.open:
|
||||
try:
|
||||
await self.ws.send(json.dumps(msg))
|
||||
except websockets.ConnectionClosed:
|
||||
pass
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Signal the client to stop."""
|
||||
self.running = False
|
||||
self.cancel_event.set()
|
||||
|
||||
|
||||
# ── Entry Point ──────────────────────────────────────────────────────────────
|
||||
|
||||
def main() -> None:
|
||||
server_url = os.environ.get("EASYBOT_SERVER")
|
||||
token = os.environ.get("EASYBOT_TOKEN")
|
||||
|
||||
if not server_url or not token:
|
||||
print("Error: EASYBOT_SERVER and EASYBOT_TOKEN environment variables are required.")
|
||||
print()
|
||||
print("Usage:")
|
||||
print(" set EASYBOT_SERVER=wss://app.example.com/ws/windows-agent")
|
||||
print(" set EASYBOT_TOKEN=eb_wag_...")
|
||||
print(" easybot-windows-agent")
|
||||
sys.exit(1)
|
||||
|
||||
log_level = os.environ.get("LOG_LEVEL", "INFO").upper()
|
||||
logging.basicConfig(
|
||||
level=getattr(logging, log_level, logging.INFO),
|
||||
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
)
|
||||
|
||||
logger.info("EasyBot Windows Agent v%s starting", __version__)
|
||||
logger.info("Server: %s", server_url)
|
||||
logger.info("Token: %s...", token[:12] if len(token) > 12 else "***")
|
||||
|
||||
client = AgentClient(server_url, token)
|
||||
|
||||
# Graceful shutdown
|
||||
def handle_signal(*_: Any) -> None:
|
||||
logger.info("Shutdown signal received")
|
||||
client.stop()
|
||||
|
||||
signal.signal(signal.SIGINT, handle_signal)
|
||||
signal.signal(signal.SIGTERM, handle_signal)
|
||||
|
||||
asyncio.run(client.run())
|
||||
logger.info("Agent stopped")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Reference in New Issue
Block a user