Package: ai21
AI21 Labs Python SDKMigration from v1.3.4 and belowInv2.0.0we introduced a new SDK that is not backwards compatible with the previous version.
This version allows for Non static instances of the client, defined parameters to each resource, modelized responses and
more.Migration ExamplesInstance creation (not available in v1.3.4 and below)fromai21importAI21Clientclient=AI21Client(api_key='my_api_key')# or set api_key in environment variable - AI21_API_KEY and thenclient=AI21Client()We No longer support static methods for each resource, instead we have a client instance that has a method for each
allowing for more flexibility and better control.Completion before/afterprompt = "some prompt"- import ai21- response = ai21.Completion.execute(model="j2-light", prompt=prompt, maxTokens=2)+ from ai21 import AI21Client+ client = ai21.AI21Client()+ response = client.completion.create(model="j2-light", prompt=prompt, max_tokens=2)This applies to all resources. You would now need to create a client instance and use it to call the resource method.Tokenization and Token counting before/after- response = ai21.Tokenization.execute(text=prompt)- print(len(response)) # number of tokens+ from ai21 import AI21Client+ client = AI21Client()+ token_count = client.count_tokens(text=prompt)Key Access in Response Objects before/afterIt is no longer possible to access the response object as a dictionary. Instead, you can access the response object as an object with attributes.- import ai21- response = ai21.Summarize.execute(source="some text", sourceType="TEXT")- response["summary"]+ from ai21 import AI21Client+ from ai21.models import DocumentType+ client = AI21Client()+ response = client.summarize.create(source="some text", source_type=DocumentType.TEXT)+ response.summaryAWS Client CreationsBedrock Client creation before/after- import ai21- destination = ai21.BedrockDestination(model_id=ai21.BedrockModelID.J2_MID_V1)- response = ai21.Completion.execute(prompt=prompt, maxTokens=1000, destination=destination)+ from ai21 import AI21BedrockClient, BedrockModelID+ client = AI21BedrockClient()+ response = client.completion.create(prompt=prompt, max_tokens=1000, model_id=BedrockModelID.J2_MID_V1)SageMaker Client creation before/after- import ai21- destination = ai21.SageMakerDestination("j2-mid-test-endpoint")- response = ai21.Completion.execute(prompt=prompt, maxTokens=1000, destination=destination)+ from ai21 import AI21SageMakerClient+ client = AI21SageMakerClient(endpoint_name="j2-mid-test-endpoint")+ response = client.completion.create(prompt=prompt, max_tokens=1000)Installationpippipinstallai21UsageModel TypesWherever you are required to pass a model name as a parameter, you can use any of the available AI21
state-of-the-art models:j2-lightj2-midj2-ultrayou can read more about the modelshere.Client Instance Creationfromai21importAI21Clientclient=AI21Client(# defaults to os.enviorn.get('AI21_API_KEY')api_key='my_api_key',)response=client.completion.create(prompt="<your prompt here>",max_tokens=10,model="j2-mid",temperature=0.3,top_p=1,)print(response.completions)print(response.prompt)Token CountingBy using thecount_tokensmethod, you can estimate the billing for a given request.fromai21importAI21Clientclient=AI21Client()client.count_tokens(text="some text")# returns intFile Uploadfromai21importAI21Clientclient=AI21Client()file_id=client.library.files.create(file_path="path/to/file",path="path/to/file/in/library",labels=["label1","label2"],public_url="www.example.com",)uploaded_file=client.library.files.get(file_id)Environment VariablesYou can set several environment variables to configure the client.LoggingWe use the standard libraryloggingmodule.To enable logging, set theAI21_LOG_LEVELenvironment variable.$exportAI21_LOG_LEVEL=debugOther Important Environment VariablesAI21_API_KEY- Your API key. If not set, you must pass it to the client constructor.AI21_API_VERSION- The API version. Defaults tov1.AI21_API_HOST- The API host. Defaults tohttps://api.ai21.com/v1/.AI21_TIMEOUT_SEC- The timeout for API requests.AI21_NUM_RETRIES- The maximum number of retries for API requests. Defaults to3retries.AI21_AWS_REGION- The AWS region to use for AWS clients. Defaults tous-east-1.Error Handlingfromai21importerrorsasai21_errorsfromai21importAI21Client,AI21APIErrorfromai21.modelsimportChatMessageclient=AI21Client()system="You're a support engineer in a SaaS company"messages=[# Notice the given role does not exist and will be the reason for the raised errorChatMessage(text="Hello, I need help with a signup process.",role="Non-Existent-Role"),]try:chat_completion=client.chat.create(messages=messages,model="j2-ultra",system=system)exceptai21_errors.AI21ServerErrorase:print("Server error and could not be reached")print(e.details)exceptai21_errors.TooManyRequestsErrorase:print("A 429 status code was returned. Slow down on the requests")exceptAI21APIErrorase:print("A non 200 status code error. For more error types see ai21.errors")AWS ClientsAI21 Library provides convenient ways to interact with two AWS clients for use with AWS SageMaker and AWS Bedrock.Installationpipinstall"ai21[AWS]"This will make sure you have the required dependencies installed, includingboto3 >= 1.28.82.UsageSageMakerfromai21importAI21SageMakerClientclient=AI21SageMakerClient(endpoint_name="j2-endpoint-name")response=client.summarize.create(source="Text to summarize",source_type="TEXT",)print(response.summary)With Boto3 Sessionfromai21importAI21SageMakerClientimportboto3boto_session=boto3.Session(region_name="us-east-1")client=AI21SageMakerClient(session=boto_session,endpoint_name="j2-endpoint-name",)Bedrockfromai21importAI21BedrockClient,BedrockModelIDclient=AI21BedrockClient(region='us-east-1')# region is optional, as you can use the env variable insteadresponse=client.completion.create(prompt="Your prompt here",model_id=BedrockModelID.J2_MID_V1,max_tokens=10,)print(response.completions[0].data.text)With Boto3 Sessionfromai21importAI21BedrockClient,BedrockModelIDimportboto3boto_session=boto3.Session(region_name="us-east-1")client=AI21BedrockClient(session=boto_session,)response=client.completion.create(prompt="Your prompt here",model_id=BedrockModelID.J2_MID_V1,max_tokens=10,)