AgentOps has first party support for AI21 via its Python SDK. Explore development with AI21 by visiting their docs.

Steps to integrate AI21 with AgentOps

1

Install the AgentOps SDK

pip install agentops
2

Install the AI21 SDK

We currently support v2.x.x of the AI21 SDK and plan to support v3.x.x in the future.

pip install "ai21<3.0.0"
3

Add AgentOps to your code

Make sure to call agentops.init before calling any openai, cohere, crew, etc models.

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

# Initialize clients
agentops.init(<INSERT YOUR API KEY HERE>)
client = AI21Client(api_key="your-api-key")

# Your AI21 code here...

agentops.end_session("Success")

Set your API key as an .env variable for easy access.

AGENTOPS_API_KEY=<YOUR API KEY>

Read more about environment variables in Advanced Configuration

4

Run your Agent

Execute your program and visit app.agentops.ai/drilldown to observe your Agent! πŸ•΅οΈ

After your run, AgentOps prints a clickable url to console linking directly to your session in the Dashboard

Clickable link to session

Full Examples

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

agentops.init(<INSERT YOUR API KEY HERE>)
client = AI21Client(api_key="your-api-key")

messages = [
    ChatMessage(
        content="You are a world renowned poet in the style of Edgar Allan Poe.",
        role="system",
    ),
    ChatMessage(
        content="Write me a short poem about the AI agents co-existing within the human brain.",
        role="user",
    ),
]

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

print(response.choices[0].message.content)
agentops.end_session('Success')

Streaming Examples

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

agentops.init(<INSERT YOUR API KEY HERE>)
client = AI21Client(api_key="your-api-key")

messages = [
    ChatMessage(
        content="You are a world renowned poet in the style of Edgar Allan Poe.",
        role="system",
    ),
    ChatMessage(
        content="Write me a short poem about the AI agents co-existing within the human brain.",
        role="user",
    ),
]

complete_response = ""

response = client.chat.completions.create(
    messages=messages,
    model="jamba-1.5-mini",
    stream=True,
)

for chunk in response:
  complete_response += str(chunk.choices[0].delta.content)

print(complete_response)
agentops.end_session('Success')

Task-Specific Models

AI21 v2.x.x provides specialized models called Task Specific Models for specific tasks.

Here’s an example using the contextual answers endpoint:

from ai21 import AI21Client
import agentops

agentops.init(<INSERT YOUR API KEY HERE>)
client = AI21Client(api_key="your-api-key")

CONTEXT = """
In 2020 and 2021, enormous QE β€” approximately $4.4 trillion, or 18%, of 2021 gross
domestic product (GDP) β€” and enormous fiscal stimulus (which has been and
always will be inflationary) β€” approximately $5 trillion, or 21%, of 2021 GDP
β€” stabilized markets and allowed companies to raise enormous amounts of
capital. In addition, this infusion of capital saved many small businesses and
put more than $2.5 trillion in the hands of consumers and almost $1 trillion into
state and local coffers. These actions led to a rapid decline in unemployment, 
dropping from 15% to under 4% in 20 months β€” the magnitude and speed of which were both
unprecedented. Additionally, the economy grew 7% in 2021 despite the arrival of
the Delta and Omicron variants and the global supply chain shortages, which were
largely fueled by the dramatic upswing in consumer spending and the shift in
that spend from services to goods.
"""

response = client.answer.create(
    context=CONTEXT,
    question="Did the economy shrink after the Omicron variant arrived?",
)

print(response.answer)
agentops.end_session('Success')

If you want to stream the response, you can use the stream=True flag. The streaming response is handled by the AI21 SDK.

All of these examples can be found in this notebook.