Chat completions

import base64
import json
import os

import requests

API = "https://api.abstract.inc/v1"
MODEL = "allenai/molmoact2:yam"
CAMERAS = ["top", "left", "right"]


def image_part(path):
    with open(path, "rb") as file:
        data = base64.b64encode(file.read()).decode()
    return {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{data}"}}


body = {
    "model": MODEL,
    "stream": True,
    "messages": [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "put the red cup in the bin"},
                *[image_part(f"{camera}.jpg") for camera in CAMERAS],
            ],
        }
    ],
    "robot_state": {"joint_position": [0.0] * 12, "gripper_position": [1.0, 1.0]},
}

response = requests.post(
    f"{API}/chat/completions",
    headers={"Authorization": f"Bearer {os.environ['ABSTRACT_API_KEY']}"},
    json=body,
    stream=True,
    timeout=300,
)
if not response.ok:
    raise SystemExit(f"HTTP {response.status_code}: {response.text}")

chunk = usage = None
for line in response.iter_lines(decode_unicode=True):
    if not line.startswith("data: "):
        continue  # blank separators and ": processing" keep-alives
    data = line.removeprefix("data: ")
    if data == "[DONE]":
        break
    event = json.loads(data)
    if "error" in event:
        raise SystemExit(f"model failed: {event['error']}")
    delta = event["choices"][0]["delta"]
    if "content" in delta:
        chunk = json.loads(delta["content"])
    usage = event.get("usage", usage)

print(f"{len(chunk['actions'])} actions, usage {usage}")

import base64
import json
import os
import uuid

import requests

API = "https://api.abstract.inc/v1"
HEADERS = {"Authorization": f"Bearer {os.environ['ABSTRACT_API_KEY']}"}
MODEL = "abstract/dreamzero-astra:yam"
CAMERAS = ["top", "left", "right"]
INSTRUCTION = "fold the towel"


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, 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 image_part(jpeg):
    data = base64.b64encode(jpeg).decode()
    return {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{data}"}}


def request_chunk(session_id, frames):
    content = [{"type": "text", "text": INSTRUCTION}]
    for camera in CAMERAS:  # camera-major, oldest frame first
        content += [image_part(jpeg) for jpeg in frames[camera]]
    body = {
        "model": MODEL,
        "session_id": session_id,
        "messages": [{"role": "user", "content": content}],
        "robot_state": read_state(),
    }
    response = requests.post(f"{API}/chat/completions", headers=HEADERS, json=body, timeout=300)
    # HTTP 410 means the session expired: start again with a new session_id.
    if not response.ok:
        raise SystemExit(f"HTTP {response.status_code}: {response.text}")
    return json.loads(response.json()["choices"][0]["message"]["content"])


session_id = str(uuid.uuid4())  # a new id starts a new episode
frames = {camera: [capture(camera)] for camera in CAMERAS}  # 1 frame per camera opens it
for _ in range(3):
    chunk = request_chunk(session_id, frames)
    print(f"step {chunk['step']}: {len(chunk['actions'])} actions")
    frames = execute(chunk["actions"], chunk["dt"], chunk["next_frames_per_camera"])