WebSocket Real-Time Updates

import asyncio
import json
import websockets

API_KEY = "sk_live_your_key_here"
WS_URL = f"wss://alterlab.io/api/v1/ws/jobs?api_key={API_KEY}"

async def listen_for_updates(job_id: str):
    """Connect to WebSocket and stream job updates."""
    async with websockets.connect(WS_URL) as ws:
        # Wait for connection confirmation
        msg = json.loads(await ws.recv())
        print(f"Connected: {msg}")

        # Subscribe to job
        await ws.send(json.dumps({
            "action": "subscribe",
            "job_id": job_id
        }))

        # Listen for events
        async for raw in ws:
            event = json.loads(raw)

            if event["type"] == "job_update":
                print(f"Job {event['job_id']}: {event['status']}")

                if event["status"] == "succeeded":
                    print(f"Result: {event['result']}")
                    return event["result"]

                if event["status"] == "failed":
                    raise Exception(f"Job failed: {event['error']}")

            elif event["type"] == "heartbeat":
                pass  # Connection is alive

            elif event["type"] == "error":
                print(f"Error: {event['message']}")

# Usage
# result = asyncio.run(listen_for_updates("your-job-id"))

import asyncio
import json
import websockets

API_KEY = "sk_live_your_key_here"
WS_URL = f"wss://alterlab.io/api/v1/ws/jobs?api_key={API_KEY}"

async def robust_listener(job_ids: list[str], on_update):
    """WebSocket listener with exponential backoff reconnection."""
    delay = 1
    max_delay = 30

    while True:
        try:
            async with websockets.connect(WS_URL) as ws:
                delay = 1  # Reset on successful connect

                # Wait for connected event
                await ws.recv()

                # Subscribe to all jobs
                for job_id in job_ids:
                    await ws.send(json.dumps({
                        "action": "subscribe",
                        "job_id": job_id
                    }))

                # Listen for events
                async for raw in ws:
                    event = json.loads(raw)
                    if event["type"] == "job_update":
                        on_update(event)

        except (websockets.ConnectionClosed, OSError) as e:
            print(f"Disconnected: {e}. Retrying in {delay}s...")
            await asyncio.sleep(delay)
            delay = min(delay * 2, max_delay)