scripts/mac_agent.py

#!/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()