Python SDK: AI Gateway chat

completion = client.ai.chat({
    "model": "cubic",
    "messages": [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What is Square Cloud?"},
    ],
    "max_tokens": 512,
})

print(completion["choices"][0]["message"]["content"])
print(completion["usage"]["total_tokens"])

from squarecloud import SquareCloudAPIError

try:
    client.ai.chat({"messages": [{"role": "user", "content": "Hi"}]})
except SquareCloudAPIError as error:
    if error.code == "server_overloaded":
        ...  # safe to retry yourself
    else:
        raise