AI21 | liteLLM (original) (raw)

LiteLLM supports the following AI21 models:

• jamba-1.5-mini
• jamba-1.5-large
• j2-light
• j2-mid
• j2-ultra

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="ai21/jamba-1.5-mini", 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

1. Send Request to LiteLLM Proxy Server

• OpenAI Python v1.0.0+
• curl

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​

Supported AI21 Parameters​

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.

• LiteLLM Python SDK
• LiteLLM Proxy Server

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​