Skip to main content

AI21

LiteLLM supports j2-light, j2-mid and j2-ultra from AI21

tip

We support ALL AI21 models, just set model=ai21/<any-model-on-ai21> as a prefix when sending litellm requests. See all litellm supported AI21 models here

API KEYS

import os 
os.environ["AI21_API_KEY"] = "your-api-key"

LiteLLM Python SDK Usage

Sample Usage

from litellm import completion 

# set env variable
os.environ["AI21_API_KEY"] = "your-api-key"

messages = [{"role": "user", "content": "Write me a poem about the blue sky"}]

completion(model="j2-light", messages=messages)

LiteLLM Proxy Server Usage

Here's how to call a ai21 model with the LiteLLM Proxy Server

  1. Modify the config.yaml

    model_list:
    - model_name: my-model
    litellm_params:
    model: ai21/<your-model-name> # add ai21/ prefix to route as ai21 provider
    api_key: api-key # api key to send your model
  1. Start the proxy

    $ litellm --config /path/to/config.yaml
  2. Send Request to LiteLLM Proxy Server

    import openai
    client = openai.OpenAI(
    api_key="sk-1234", # pass litellm proxy key, if you're using virtual keys
    base_url="http://0.0.0.0:4000" # litellm-proxy-base url
    )

    response = client.chat.completions.create(
    model="my-model",
    messages = [
    {
    "role": "user",
    "content": "what llm are you"
    }
    ],
    )

    print(response)

Supported OpenAI Parameters

paramtypeAI21 equivalent
toolsOptional[list]tools
response_formatOptional[dict]response_format
max_tokensOptional[int]max_tokens
temperatureOptional[float]temperature
top_pOptional[float]top_p
stopOptional[Union[str, list]]stop
nOptional[int]n
streamOptional[bool]stream
seedOptional[int]seed
tool_choiceOptional[str]tool_choice
userOptional[str]user

Supported AI21 Parameters

paramtypeAI21 equivalent
documentsOptional[List[Dict]]documents

Passing AI21 Specific Parameters - documents

LiteLLM allows you to pass all AI21 specific parameters to the litellm.completion function. Here is an example of how to pass the documents parameter to the litellm.completion function.

response = await litellm.acompletion(
model="jamba-1.5-large",
messages=[{"role": "user", "content": "what does the document say"}],
documents = [
{
"content": "hello world",
"metadata": {
"source": "google",
"author": "ishaan"
}
}
]
)

tip

We support ALL AI21 models, just set model=ai21/<any-model-on-ai21> as a prefix when sending litellm requests See all litellm supported AI21 models here

AI21 Models

Model NameFunction CallRequired OS Variables
jamba-1.5-minicompletion('jamba-1.5-mini', messages)os.environ['AI21_API_KEY']
jamba-1.5-largecompletion('jamba-1.5-large', messages)os.environ['AI21_API_KEY']
j2-lightcompletion('j2-light', messages)os.environ['AI21_API_KEY']
j2-midcompletion('j2-mid', messages)os.environ['AI21_API_KEY']
j2-ultracompletion('j2-ultra', messages)os.environ['AI21_API_KEY']