Sessions over WebSocket

import json
import os
from urllib.parse import quote

from websockets.sync.client import connect

MODEL = "abstract/dreamzero-astra:yam"
URL = f"wss://api.abstract.inc/v1/responses?model={quote(MODEL, safe='')}"
HEADERS = {"Authorization": f"Bearer {os.environ['ABSTRACT_API_KEY']}"}


def capture(camera):
    """Return the latest JPEG from one camera. Replace this with your camera driver."""
    with open(f"{camera}.jpg", "rb") as file:
        return file.read()


def read_state():
    """Return the robot's state now. Replace this with your robot driver."""
    return {"joint_position": [0.0] * 12, "gripper_position": [1.0, 1.0]}


def execute(actions, dt, cameras, frames_per_camera):
    """Run the chunk on the robot, one row every dt seconds. Replace this with your
    controller. Return frames_per_camera JPEGs per camera taken during the chunk, oldest first."""
    return {camera: [capture(camera) for _ in range(frames_per_camera)] for camera in cameras}


def receive(socket):
    message = json.loads(socket.recv())
    if message["type"] == "error":
        raise SystemExit(f"{message['error']['code']}: {message['error']['message']}")
    return message


with connect(URL, additional_headers=HEADERS, subprotocols=["abstract.responses.v1"]) as socket:
    socket.send(json.dumps({"type": "session.open", "model": MODEL, "instruction": "fold the towel"}))
    opened = receive(socket)
    cameras = opened["cameras"]
    frames = {camera: [capture(camera)] for camera in cameras}  # frames_per_camera_first is 1
    for _ in range(3):
        for index, camera in enumerate(cameras):
            for jpeg in frames[camera]:
                socket.send(bytes([index]) + jpeg)
        socket.send(json.dumps({"type": "observation", "robot_state": read_state()}))
        chunk = receive(socket)
        print(f"step {chunk['step']}: {len(chunk['actions'])} actions, usage {chunk['usage']}")
        frames = execute(chunk["actions"], chunk["dt"], cameras, opened["frames_per_camera_next"])
    socket.send(json.dumps({"type": "session.close"}))
    print(receive(socket))