bunny-lab/Borealis
2
0
Fork 0
Code Issues 10 Pull Requests Actions Packages Projects 1 Releases Wiki Activity

Redesigned Agent Architecture

Browse Source
This commit is contained in:
Nicole Rappe
2025-04-15 20:27:23 -06:00
parent 17e0d20063
commit 8d67e847e9
9 changed files with 487 additions and 334 deletions
Download Patch File Download Diff File Expand all files Collapse all files

167
Data/Agent/borealis-agent.py
View File

@ -8,31 +8,28 @@ import threading
import socketio
from io import BytesIO
import socket
import os
from PyQt5 import QtCore, QtGui, QtWidgets
from PIL import ImageGrab
# ---------------- Configuration ----------------
SERVER_URL = "http://localhost:5000" # WebSocket-enabled Internal URL
#SERVER_URL = "https://borealis.bunny-lab.io" # WebSocket-enabled Public URL"
SERVER_URL = os.environ.get("SERVER_URL", "http://localhost:5000")
HOSTNAME = socket.gethostname().lower()
RANDOM_SUFFIX = uuid.uuid4().hex[:8]
AGENT_ID = f"{HOSTNAME}-agent-{RANDOM_SUFFIX}"
# ---------------- State ----------------
# ---------------- App State ----------------
app_instance = None
region_widget = None
current_interval = 1000
config_ready = threading.Event()
overlay_visible = True
overlay_widgets = {}
region_launchers = {}
running_roles = {}
running_threads = {}
LAST_CONFIG = {}
# WebSocket client setup
# ---------------- Socket Setup ----------------
sio = socketio.Client()
# ---------------- WebSocket Handlers ----------------
@sio.event
def connect():
print(f"[WebSocket] Agent ID: {AGENT_ID} connected to Borealis.")
@ -45,30 +42,18 @@ def disconnect():
@sio.on('agent_config')
def on_agent_config(config):
global current_interval, overlay_visible, LAST_CONFIG
print("[PROVISIONED] Received new configuration from Borealis.")
if config != LAST_CONFIG:
print("[PROVISIONED] Received new configuration from Borealis.")
x = config.get("x", 100)
y = config.get("y", 100)
w = config.get("w", 300)
h = config.get("h", 200)
current_interval = config.get("interval", 1000)
overlay_visible = config.get("visible", True)
roles = config.get("roles", [])
stop_all_roles()
for role in roles:
start_role_thread(role)
if not region_widget:
region_launcher.trigger.emit(x, y, w, h)
else:
region_widget.setGeometry(x, y, w, h)
region_widget.setVisible(overlay_visible)
LAST_CONFIG = config
config_ready.set()
# ---------------- Region Overlay ----------------
# ---------------- Overlay Class ----------------
class ScreenshotRegion(QtWidgets.QWidget):
def __init__(self, x=100, y=100, w=300, h=200):
def __init__(self, node_id, x=100, y=100, w=300, h=200):
super().__init__()
self.node_id = node_id
self.setGeometry(x, y, w, h)
self.setWindowFlags(QtCore.Qt.FramelessWindowHint | QtCore.Qt.WindowStaysOnTopHint)
self.setAttribute(QtCore.Qt.WA_TranslucentBackground)
@ -78,7 +63,7 @@ class ScreenshotRegion(QtWidgets.QWidget):
self.setVisible(True)
self.label = QtWidgets.QLabel(self)
self.label.setText(AGENT_ID)
self.label.setText(f"{node_id[:8]}")
self.label.setStyleSheet("color: lime; background: transparent; font-size: 10px;")
self.label.move(8, 4)
@ -87,7 +72,6 @@ class ScreenshotRegion(QtWidgets.QWidget):
def paintEvent(self, event):
painter = QtGui.QPainter(self)
painter.setRenderHint(QtGui.QPainter.Antialiasing)
painter.setBrush(QtCore.Qt.transparent)
painter.setPen(QtGui.QPen(QtGui.QColor(0, 255, 0), 2))
painter.drawRect(self.rect())
@ -102,8 +86,8 @@ class ScreenshotRegion(QtWidgets.QWidget):
def mousePressEvent(self, event):
if event.button() == QtCore.Qt.LeftButton:
if (event.pos().x() > self.width() - self.resize_handle_size and
event.pos().y() > self.height() - self.resize_handle_size):
if event.pos().x() > self.width() - self.resize_handle_size and \
event.pos().y() > self.height() - self.resize_handle_size:
self.resizing = True
else:
self.drag_offset = event.globalPos() - self.frameGeometry().topLeft()
@ -124,58 +108,91 @@ class ScreenshotRegion(QtWidgets.QWidget):
geo = self.geometry()
return geo.x(), geo.y(), geo.width(), geo.height()
# ---------------- Screenshot Capture ----------------
def capture_loop():
config_ready.wait()
while region_widget is None:
time.sleep(0.2)
while True:
if overlay_visible:
x, y, w, h = region_widget.get_geometry()
try:
img = ImageGrab.grab(bbox=(x, y, x + w, y + h))
buffer = BytesIO()
img.save(buffer, format="PNG")
encoded_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
sio.emit('screenshot', {
'agent_id': AGENT_ID,
'image_base64': encoded_image
})
except Exception as e:
print(f"[ERROR] Screenshot error: {e}")
time.sleep(current_interval / 1000)
# ---------------- UI Launcher ----------------
# ---------------- Region UI Handler ----------------
class RegionLauncher(QtCore.QObject):
trigger = QtCore.pyqtSignal(int, int, int, int)
def __init__(self):
def __init__(self, node_id):
super().__init__()
self.node_id = node_id
self.trigger.connect(self.handle)
def handle(self, x, y, w, h):
launch_region(x, y, w, h)
print(f"[Overlay] Launching overlay for {self.node_id} at ({x},{y},{w},{h})")
if self.node_id in overlay_widgets:
return
widget = ScreenshotRegion(self.node_id, x, y, w, h)
overlay_widgets[self.node_id] = widget
widget.show()
region_launcher = None
# ---------------- Role Management ----------------
def stop_all_roles():
for node_id, thread in running_threads.items():
if thread and thread.is_alive():
print(f"[Role] Terminating previous task: {node_id}")
running_roles.clear()
running_threads.clear()
def launch_region(x, y, w, h):
global region_widget
if region_widget:
def start_role_thread(role_cfg):
role = role_cfg.get("role")
node_id = role_cfg.get("node_id")
if not role or not node_id:
print("[ERROR] Invalid role configuration (missing role or node_id).")
return
region_widget = ScreenshotRegion(x, y, w, h)
region_widget.show()
if role == "screenshot":
thread = threading.Thread(target=run_screenshot_loop, args=(node_id, role_cfg), daemon=True)
else:
print(f"[SKIP] Unknown role: {role}")
return
running_roles[node_id] = role_cfg
running_threads[node_id] = thread
thread.start()
print(f"[Role] Started task: {role} ({node_id})")
# ---------------- Screenshot Role Loop ----------------
def run_screenshot_loop(node_id, cfg):
interval = cfg.get("interval", 1000)
visible = cfg.get("visible", True)
x = cfg.get("x", 100)
y = cfg.get("y", 100)
w = cfg.get("w", 300)
h = cfg.get("h", 200)
if node_id not in region_launchers:
launcher = RegionLauncher(node_id)
region_launchers[node_id] = launcher
launcher.trigger.emit(x, y, w, h)
widget = overlay_widgets.get(node_id)
if widget:
widget.setGeometry(x, y, w, h)
widget.setVisible(visible)
while True:
try:
if node_id in overlay_widgets:
widget = overlay_widgets[node_id]
x, y, w, h = widget.get_geometry()
print(f"[Capture] Screenshot task {node_id} at ({x},{y},{w},{h})")
img = ImageGrab.grab(bbox=(x, y, x + w, y + h))
buffer = BytesIO()
img.save(buffer, format="PNG")
encoded = base64.b64encode(buffer.getvalue()).decode("utf-8")
sio.emit("agent_screenshot_task", {
"agent_id": AGENT_ID,
"node_id": node_id,
"image_base64": encoded
})
except Exception as e:
print(f"[ERROR] Screenshot task {node_id} failed: {e}")
time.sleep(interval / 1000)
# ---------------- Main ----------------
if __name__ == "__main__":
app_instance = QtWidgets.QApplication(sys.argv)
region_launcher = RegionLauncher()
sio.connect(SERVER_URL, transports=['websocket'])
threading.Thread(target=capture_loop, daemon=True).start()
sio.connect(SERVER_URL, transports=["websocket"])
sys.exit(app_instance.exec_())