#!/usr/bin/env python3 """ GBridge Mac Agent — Local MacBook compute proxy for HF Space. Runs on your MacBook. Connects to the HF Space via WebSocket. Receives chat prompts from your iPhone (via the Space) and answers using local Hugging Face models via Ollama or transformers. Usage: export HF_SPACE_URL="https://your-space.hf.space" export SESSION_ID="sess_xxx" python scripts/mac_agent.py Or with command-line args: python scripts/mac_agent.py --space https://your-space.hf.space --session sess_xxx --model ollama/llama3 """ import argparse import asyncio import json import os import sys import time from typing import Optional import websockets class MacAgent: def __init__(self, space_url: str, session_id: str, node_id: str, model_name: str): self.space_url = space_url.rstrip("/") self.session_id = session_id self.node_id = node_id self.model_name = model_name self.ws = None self.reconnect_delay = 3 self.use_ollama = model_name.startswith("ollama/") self.ollama_model = model_name.replace("ollama/", "") def _ws_url(self) -> str: base = self.space_url.replace("https://", "wss://").replace("http://", "ws://") return f"{base}/ws/gbridge/worker/{self.session_id}/{self.node_id}" async def run(self): while True: try: await self._connect() except Exception as e: print(f"[agent] connection error: {e}") print(f"[agent] reconnecting in {self.reconnect_delay}s...") await asyncio.sleep(self.reconnect_delay) self.reconnect_delay = min(self.reconnect_delay + 2, 30) async def _connect(self): url = self._ws_url() print(f"[agent] connecting to {url}") async with websockets.connect(url) as ws: self.ws = ws self.reconnect_delay = 3 print("[agent] connected") await ws.send(json.dumps({"op": "worker_hello", "runtime_type": "ollama", "device_public_key": None})) await ws.send(json.dumps({"op": "capabilities", "capabilities": [ {"capability_name": "iphone.text.echo.private", "trust_level": "trusted_device", "local_only": False, "requires_attestation": False}, {"capability_name": "iphone.privacy.redact.local", "trust_level": "trusted_device", "local_only": False, "requires_attestation": False} ]})) async for message in ws: await self._handle_message(json.loads(message)) async def _handle_message(self, msg: dict): op = msg.get("op") if op == "worker_welcome": print(f"[agent] registered as worker {msg.get('worker_id')}") elif op == "job_offer": await self._run_inference(msg) elif op == "heartbeat": await self._send_heartbeat() async def _send_heartbeat(self): if self.ws: await self.ws.send(json.dumps({ "op": "heartbeat", "timestamp": time.time(), "device_type": "mac", "cpu_usage": 10.0, "ram_used_mb": 4000.0, "ram_total_mb": 16000.0, "thermal_state": "nominal", "battery_level": 100.0, "tokens_per_second": 0.0 })) async def _run_inference(self, msg: dict): job_id = msg["job_id"] payload = msg.get("payload", {}) prompt = payload.get("text", "") max_tokens = payload.get("max_tokens", 512) temperature = payload.get("temperature", 0.7) print(f"[agent] job {job_id}: prompt={prompt[:60]}...") started = time.time() # Accept the job await self.ws.send(json.dumps({"op": "job_accept", "job_id": job_id})) try: if self.use_ollama: full_text = await self._run_ollama(prompt, max_tokens, temperature) else: full_text = await self._run_transformers(prompt, max_tokens, temperature) latency_ms = int((time.time() - started) * 1000) await self._send_complete( job_id, output={"text": full_text}, latency_ms=latency_ms, ) print(f"[agent] job {job_id} complete ({latency_ms}ms)") except Exception as e: print(f"[agent] inference error: {e}") await self._send_complete( job_id, output={"text": f"Error: {e}"}, latency_ms=int((time.time() - started) * 1000), ) async def _run_ollama(self, prompt: str, max_tokens: int, temperature: float) -> str: import aiohttp async with aiohttp.ClientSession() as session: async with session.post( "http://localhost:11434/api/generate", json={ "model": self.ollama_model, "prompt": prompt, "stream": False, "options": { "num_predict": max_tokens, "temperature": temperature } } ) as resp: data = await resp.json() return data.get("response", "") async def _run_transformers(self, prompt: str, max_tokens: int, temperature: float) -> str: from transformers import AutoModelForCausalLM, AutoTokenizer import torch print(f"[agent] loading model {self.model_name}...") model = AutoModelForCausalLM.from_pretrained(self.model_name, device_map="auto") tokenizer = AutoTokenizer.from_pretrained(self.model_name) inputs = tokenizer(prompt, return_tensors="pt").to(model.device) with torch.no_grad(): outputs = model.generate( **inputs, max_new_tokens=max_tokens, temperature=temperature, do_sample=True ) full_text = tokenizer.decode(outputs[0], skip_special_tokens=True) return full_text[len(prompt):] async def _send_complete(self, job_id: str, output: dict, latency_ms: int): if self.ws: await self.ws.send(json.dumps({ "op": "job_result", "job_id": job_id, "output": output, "latency_ms": latency_ms, "input_hash": None, "output_hash": None, "device_signature": None })) def main(): parser = argparse.ArgumentParser(description="GBridge Mac Agent") parser.add_argument("--space", default=os.getenv("HF_SPACE_URL"), help="HF Space public URL") parser.add_argument("--session", default=os.getenv("SESSION_ID"), help="Session ID") parser.add_argument("--node", default=os.getenv("NODE_ID", f"mac_{os.uname().nodename}"), help="Node ID") parser.add_argument("--model", default=os.getenv("MODEL", "ollama/llama3"), help="Model name (ollama/llama3 or hf/model-name)") args = parser.parse_args() if not args.space: print("Error: --space or HF_SPACE_URL required") sys.exit(1) if not args.session: print("Error: --session or SESSION_ID required") sys.exit(1) agent = MacAgent(args.space, args.session, args.node, args.model) asyncio.run(agent.run()) if __name__ == "__main__": main()