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 AI21Clientfrom ai21.models.chat import ChatMessageimport agentops# Initialize clientsagentops.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.
from ai21 import AI21Clientfrom ai21.models.chat import ChatMessageimport agentopsagentops.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')
from ai21 import AI21Clientfrom ai21.models.chat import ChatMessageimport agentopsagentops.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')
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 AI21Clientimport agentopsagentops.init(<INSERT YOUR API KEY HERE>)client = AI21Client(api_key="your-api-key")CONTEXT ="""In 2020and2021, enormous QE β approximately $4.4 trillion,or18%, of 2021 grossdomestic product (GDP) β and enormous fiscal stimulus (which has been andalways will be inflationary) β approximately $5 trillion,or21%, of 2021 GDPβ stabilized markets and allowed companies to raise enormous amounts ofcapital. In addition, this infusion of capital saved many small businesses andput more than $2.5 trillion in the hands of consumers and almost $1 trillion intostate and local coffers. These actions led to a rapid decline in unemployment,dropping from15% to under 4%in20 months β the magnitude and speed of which were bothunprecedented. Additionally, the economy grew 7%in2021 despite the arrival ofthe Delta and Omicron variants and the global supply chain shortages, which werelargely fueled by the dramatic upswing in consumer spending and the shift inthat 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.