AI21Labs/ai21-python

AI21Labs/ai21-python

Description: AI21 Python SDK

Language: Python

License: Apache-2.0

Stars: 70

Forks: 13

Open issues: 11

Created: 2023-12-10T09:20:15Z

Pushed: 2026-01-28T16:25:12Z

Default branch: main

Fork: no

Archived: no

README:

AI21 Labs Python SDK

[//]: # "Add when public" [//]: # '' [//]: # ''

---

Table of Contents

• [Examples](#examples-tldr) 🗂️
• [AI21 Official Documentation](#Documentation)
• [Installation](#Installation) 💿
• [Usage - Chat Completions](#Usage)
• [Maestro](#Maestro)
• [Agents (Beta)](#Agents-Beta)
• [Conversational RAG (Beta)](#Conversational-RAG-Beta)
• [Older Models Support Usage](#Older-Models-Support-Usage)
• [More Models](#More-Models)
• [Streaming](#Streaming)
• [Environment Variables](#Environment-Variables)
• [Error Handling](#Error-Handling)
• [Cloud Providers](#Cloud-Providers) ☁️
• [AWS](#AWS)
• [Bedrock](#Bedrock)
• [SageMaker](#SageMaker)
• [Azure](#Azure)
• [Vertex](#Vertex)

Examples (tl;dr)

If you want to quickly get a glance how to use the AI21 Python SDK and jump straight to business, you can check out the examples. Take a look at our models and see them in action! Several examples and demonstrations have been put together to show our models' functionality and capabilities.

[Check out the Examples](examples/)

Feel free to dive in, experiment, and adapt these examples to suit your needs. We believe they'll help you get up and running quickly.

Documentation

---

The full documentation for the REST API can be found on docs.ai21.com.

Installation

---

pip install ai21

Usage

---

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

client = AI21Client(
# defaults to os.enviorn.get('AI21_API_KEY')
api_key='my_api_key',
)

system = "You're a support engineer in a SaaS company"
messages = [
ChatMessage(content=system, role="system"),
ChatMessage(content="Hello, I need help with a signup process.", role="user"),
]

chat_completions = client.chat.completions.create(
messages=messages,
model="jamba-mini",
)

Async Usage

You can use the AsyncAI21Client to make asynchronous requests. There is no difference between the sync and the async client in terms of usage.

import asyncio

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

system = "You're a support engineer in a SaaS company"
messages = [
ChatMessage(content=system, role="system"),
ChatMessage(content="Hello, I need help with a signup process.", role="user"),
]

client = AsyncAI21Client(
# defaults to os.enviorn.get('AI21_API_KEY')
api_key='my_api_key',
)

async def main():
response = await client.chat.completions.create(
messages=messages,
model="jamba-mini",
)

print(response)

asyncio.run(main())

A more detailed example can be found [here](examples/studio/chat/chat_completions.py).

Chat

from ai21 import AI21Client
from ai21.models import RoleType
from ai21.models import ChatMessage

system = "You're a support engineer in a SaaS company"
messages = [
ChatMessage(text="Hello, I need help with a signup process.", role=RoleType.USER),
ChatMessage(text="Hi Alice, I can help you with that. What seems to be the problem?", role=RoleType.ASSISTANT),
ChatMessage(text="I am having trouble signing up for your product with my Google account.", role=RoleType.USER),
]

client = AI21Client()
chat_response = client.chat.create(
system=system,
messages=messages,
model="j2-ultra",
)

For a more detailed example, see the chat [examples](examples/studio/chat.py).

Completion

from ai21 import AI21Client

client = AI21Client()
completion_response = client.completion.create(
prompt="This is a test prompt",
model="j2-mid",
)

Chat Completion

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

system = "You're a support engineer in a SaaS company"
messages = [
ChatMessage(content=system, role="system"),
ChatMessage(content="Hello, I need help with a signup process.", role="user"),
ChatMessage(content="Hi Alice, I can help you with that. What seems to be the problem?", role="assistant"),
ChatMessage(content="I am having trouble signing up for your product with my Google account.", role="user"),
]

client = AI21Client()

response = client.chat.completions.create(
messages=messages,
model="jamba-large",
max_tokens=100,
temperature=0.7,
top_p=1.0,
stop=["\n"],
)

print(response)

Note that jamba-large supports async and streaming as well.

For a more detailed example, see the completion [examples](examples/studio/chat/chat_completions.py).

---

Streaming

We currently support streaming for the Chat Completions API in Jamba.

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

messages = [ChatMessage(content="What is the meaning of life?", role="user")]

client = AI21Client()

response = client.chat.completions.create(
messages=messages,
model="jamba-large",
stream=True,
)
for chunk in response:
print(chunk.choices[0].delta.content, end="")

Async Streaming

import asyncio

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

messages = [ChatMessage(content="What is the meaning of life?", role="user")]

client = AsyncAI21Client()

async def main():
response = await client.chat.completions.create(
messages=messages,
model="jamba-mini",
stream=True,
)
async for chunk in response:
print(chunk.choices[0].delta.content, end="")

asyncio.run(main())

---

Maestro

AI Planning & Orchestration System built for the enterprise. Read more here.

from ai21 import AI21Client

client = AI21Client()

run_result = client.beta.maestro.runs.create_and_poll(
input="Tell me about AI21 Maestro",
requirements=[
{
"name": "length requirement",
"description": "The length of the response should be less than 2000 characters",
},
{
"name": "source requirement",
"description": (
"Should rely on information from these websites: "
"https://www.ai21.com/, https://www.ai21.com/maestro/, "
"https://docs.ai21.com/home"
),
},
],
include=["requirements_result"] # Include additional fields in the result
)

For a more detailed example, see maestro [sync](examples/studio/maestro/run.py) and [async](examples/studio/maestro/async_run.py) examples.

---

Agents (Beta)

AI21 Agents provide a comprehensive way to create, manage, and run your Agents.

from ai21 import AI21Client…