AI21 Python SDK

4.3.0 · active · verified Tue Mar 24

Official Python SDK for AI21 Labs API. Provides access to Jamba models (chat completions, streaming) and legacy J2 Jurassic models (completions). Current version: 4.3.0 (Mar 2026). SDK has two distinct API surfaces — modern Jamba chat completions API and legacy J2/Jurassic API — with different ChatMessage signatures. Jamba models are current; J2 models are legacy.

Warnings

breaking Two ChatMessage classes exist: ai21.models.chat.ChatMessage (Jamba, uses content=) and ai21.models.ChatMessage (J2/legacy, uses text= and RoleType). Mixing them causes TypeError.
Fix: For Jamba models use: from ai21.models.chat import ChatMessage with content= and string roles. For J2 legacy: from ai21.models import ChatMessage with text= and RoleType.
breaking Old J2/Jurassic API uses client.chat.create() with system= as separate param. New Jamba API uses client.chat.completions.create() with system message inside messages list. Completely different method name.
Fix: For Jamba: client.chat.completions.create(messages=[...]). For J2: client.chat.create(system=..., messages=[...]).
gotcha LLMs trained pre-2024 have no knowledge of Jamba models. Will generate J2 Jurassic patterns (model='j2-ultra', RoleType, text= parameter) which are legacy and will be deprecated.
Fix: Use Jamba models: 'jamba-large', 'jamba-mini'. Use content= not text=. Use string roles not RoleType enum.
gotcha API key env var is AI21_API_KEY. SDK reads it automatically if not passed explicitly to AI21Client().
Fix: export AI21_API_KEY=your_key or pass api_key= to AI21Client().
deprecated J2 Jurassic models (j2-ultra, j2-mid, j2-light) are legacy. Jamba is the current model family. New features only on Jamba.
Fix: Migrate to jamba-large or jamba-mini for new projects.

Install

pip install ai21 Python

Imports

AI21Client (Jamba — current)

from ai21 import AI21Client
from ai21.models.chat import ChatMessage

client = AI21Client()  # reads AI21_API_KEY from env

response = client.chat.completions.create(
    model='jamba-large',
    messages=[
        ChatMessage(content='You are a helpful assistant.', role='system'),
        ChatMessage(content='What is AI21 Labs?', role='user')
    ],
    max_tokens=200
)
print(response.choices[0].message.content)

Two ChatMessage APIs exist in the same package. New Jamba API uses content= parameter and string roles ('user', 'assistant'). Old J2 API uses text= parameter and RoleType enum. Wrong ChatMessage for wrong API raises TypeError or ValidationError.

AsyncAI21Client

from ai21 import AsyncAI21Client
from ai21.models.chat import ChatMessage
import asyncio

client = AsyncAI21Client()

async def main():
    response = await client.chat.completions.create(
        model='jamba-mini',
        messages=[ChatMessage(content='Hello', role='user')]
    )
    print(response.choices[0].message.content)

asyncio.run(main())

Use AsyncAI21Client for async. Same import path, different class.

Quickstart

Minimal AI21 Jamba chat completion using ai21 SDK 4.x.

# pip install ai21
from ai21 import AI21Client
from ai21.models.chat import ChatMessage
import os

client = AI21Client()  # reads AI21_API_KEY env var

response = client.chat.completions.create(
    model='jamba-large',
    messages=[
        ChatMessage(content='You are a helpful assistant.', role='system'),
        ChatMessage(content='Explain large language models briefly.', role='user')
    ],
    max_tokens=150,
    temperature=0.7
)
print(response.choices[0].message.content)

view raw JSON →