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"
poetry add "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 AsyncAI21Clientimport agentopsimport asyncioasync def main(): agentops.init(<INSERT YOUR API KEY HERE>) client = AsyncAI21Client(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 = await client.chat.completions.create( messages=messages, model="jamba-1.5-mini", ) print(response.choices[0].message.content) agentops.end_session('Success')asyncio.run(main())
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')
from ai21 import AsyncAI21Clientimport agentopsimport asyncioasync def main(): agentops.init(<INSERT YOUR API KEY HERE>) client = AsyncAI21Client(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 = await client.chat.completions.create( messages=messages, model="jamba-1.5-mini", stream=True, ) async for chunk in response: complete_response += str(chunk.choices[0].delta.content) print(complete_response) agentops.end_session('Success')asyncio.run(main())
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 2020 and 2021, enormous QE — approximately $4.4 trillion, or 18%, of 2021 grossdomestic product (GDP) — and enormous fiscal stimulus (which has been andalways will be inflationary) — approximately $5 trillion, or 21%, 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 from 15% to under 4% in 20 months — the magnitude and speed of which were bothunprecedented. Additionally, the economy grew 7% in 2021 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.
