feat: support ai21 client (#800)

5 files changed, 48 insertions, 71 deletions

diff --git a/Argcfile.sh b/Argcfile.sh
index 6f6e0e9..e1af57f 100755
--- a/Argcfile.sh
+++ b/Argcfile.sh
@@ -79,20 +79,21 @@ test-server() {
 }
 
 OPENAI_COMPATIBLE_PLATFORMS=( \
-  openai,gpt-3.5-turbo,https://api.openai.com/v1 \
-  deepinfra,meta-llama/Meta-Llama-3-8B-Instruct,https://api.deepinfra.com/v1/openai \
+  openai,gpt-4o-mini,https://api.openai.com/v1 \
+  ai21,jamba-1.5-mini,https://api.ai21.com/studio/v1 \
+  deepinfra,meta-llama/Meta-Llama-3.1-8B-Instruct,https://api.deepinfra.com/v1/openai \
   deepseek,deepseek-chat,https://api.deepseek.com \
-  fireworks,accounts/fireworks/models/llama-v3-8b-instruct,https://api.fireworks.ai/inference/v1 \
+  fireworks,accounts/fireworks/models/llama-v3p1-8b-instruct,https://api.fireworks.ai/inference/v1 \
   groq,llama3-8b-8192,https://api.groq.com/openai/v1 \
-  mistral,mistral-small-latest,https://api.mistral.ai/v1 \
+  mistral,open-mistral-nemo,https://api.mistral.ai/v1 \
   moonshot,moonshot-v1-8k,https://api.moonshot.cn/v1 \
-  openrouter,meta-llama/llama-3-8b-instruct,https://openrouter.ai/api/v1 \
-  octoai,meta-llama-3-8b-instruct,https://text.octoai.run/v1 \
-  perplexity,llama-3-8b-instruct,https://api.perplexity.ai \
-  together,meta-llama/Llama-3-8b-chat-hf,https://api.together.xyz/v1 \
+  openrouter,openai/gpt-4o-mini,https://openrouter.ai/api/v1 \
+  octoai,meta-llama-3.1-8b-instruct,https://text.octoai.run/v1 \
+  perplexity,llama-3.1-8b-instruct,https://api.perplexity.ai \
+  together,meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo,https://api.together.xyz/v1 \
   zhipuai,glm-4-0520,https://open.bigmodel.cn/api/paas/v4 \
   lingyiwanwu,yi-large,https://api.lingyiwanwu.com/v1 \
-  github,meta-llama-3.1-8b-instruct,https://models.inference.ai.azure.com \
+  github,gpt-4o-mini,https://models.inference.ai.azure.com \
 )
 
 # @cmd Chat with any LLM api 
@@ -280,49 +281,6 @@ chat-vertexai() {
 -d "$(_build_body vertexai "$@")" 
 }
 
-# @cmd Chat with vertexai-claude api
-# @env require-tools gcloud
-# @env VERTEXAI_PROJECT_ID!
-# @env VERTEXAI_LOCATION!
-# @option -m --model=claude-3-haiku@20240307 $VERTEXAI_CLAUDE_MODEL
-# @flag -S --no-stream
-# @arg text~
-chat-vertexai-claude() {
-    api_key="$(gcloud auth print-access-token)"
-    url=https://$VERTEXAI_LOCATION-aiplatform.googleapis.com/v1/projects/$VERTEXAI_PROJECT_ID/locations/$VERTEXAI_LOCATION/publishers/anthropic/models/$argc_model:streamRawPredict
-    _wrapper curl -i $url \
--X POST \
--H "Authorization: Bearer $api_key" \
--H 'Content-Type: application/json' \
--d "$(_build_body vertexai-claude "$@")" 
-}
-
-# @cmd Chat with bedrock api
-# @meta require-tools aws
-# @option -m --model=mistral.mistral-7b-instruct-v0:2 $BEDROCK_MODEL
-# @env AWS_REGION=us-east-1
-chat-bedrock() {
-    file="$(mktemp)"
-    case "$argc_model" in
-        mistral.* | meta.*)
-            body='{"prompt":"'"$*"'"}'
-            ;;
-        anthropic.*)
-            body="$(_build_body bedrock-claude "$@")"
-            ;;
-        *)
-            _die "Invalid model: $argc_model"
-            ;;
-    esac
-
-    _wrapper aws bedrock-runtime invoke-model \
-        --model-id $argc_model \
-        --region $AWS_REGION \
-        --body "$(echo "$body" | base64)" \
-        "$file"
-    cat "$file"
-}
-
 # @cmd Chat with cloudflare api
 # @env CLOUDFLARE_API_KEY!
 # @option -m --model=@cf/meta/llama-3-8b-instruct $CLOUDFLARE_MODEL
@@ -495,20 +453,6 @@ _build_body() {
 }'
 
             ;;
-        vertexai-claude|bedrock-claude)
-            echo '{
-    "anthropic_version": "vertex-2023-10-16",
-    "messages": [
-        {
-            "role": "user",
-            "content": "'"$*"'"
-        }
-    ],
-    "max_tokens": 4096,
-    "stream": '$stream'
-}'
-
-            ;;
         gemini|vertexai)
             echo '{
     "contents": [{
diff --git a/config.example.yaml b/config.example.yaml
index d908fa1..fa97583 100644
--- a/config.example.yaml
+++ b/config.example.yaml
@@ -157,6 +157,12 @@ clients:
     api_base: https://api.mistral.ai/v1
     api_key: xxx
 
+  # See https://docs.ai21.com/docs/quickstart
+  - type: openai-compatible
+    name: ai12
+    api_key: xxx
+    api_base: https://api.ai21.com/studio/v1
+
   # See https://docs.cohere.com/docs/the-cohere-platform
   - type: cohere
     api_key: xxx
diff --git a/models.yaml b/models.yaml
index 3d328ac..9668d06 100644
--- a/models.yaml
+++ b/models.yaml
@@ -108,7 +108,7 @@
   models:
     - name: claude-3-5-sonnet-20240620
       max_input_tokens: 200000
-      max_output_tokens: 4096
+      max_output_tokens: 8192
       require_max_tokens: true
       input_price: 3
       output_price: 15
@@ -181,6 +181,22 @@
       default_chunk_size: 2000
       max_batch_size: 3
 
+- platform: ai21
+  # docs:
+  #   - https://docs.ai21.com/reference/jamba-15-api-ref
+  #   - https://www.ai21.com/pricing
+  models:
+    - name: jamba-1.5-large
+      max_input_tokens: 256000
+      input_price: 2
+      output_price: 8
+      supports_function_calling: true
+    - name: jamba-1.5-mini
+      max_input_tokens: 256000
+      input_price: 0.2
+      output_price: 0.4
+      supports_function_calling: true
+
 - platform: cohere
   # docs:
   #   - https://docs.cohere.com/docs/command-r
@@ -328,7 +344,7 @@
       supports_function_calling: true
     - name: claude-3-5-sonnet@20240620
       max_input_tokens: 200000
-      max_output_tokens: 4096
+      max_output_tokens: 8192
       require_max_tokens: true
       input_price: 3
       output_price: 15
@@ -397,7 +413,7 @@
   models:
     - name: anthropic.claude-3-5-sonnet-20240620-v1:0
       max_input_tokens: 200000
-      max_output_tokens: 4096
+      max_output_tokens: 8192
       require_max_tokens: true
       input_price: 3
       output_price: 15
@@ -752,7 +768,6 @@
 - platform: github
   # docs:
   #   - https://github.com/marketplace/models
-  #   - https://deepinfra.com/pricing
   models:
     - name: gpt-4o
       max_input_tokens: 128000
@@ -1101,6 +1116,16 @@
       max_input_tokens: 32768
       input_price: 0.24
       output_price: 0.24
+    - name: ai21/jamba-1-5-large
+      max_input_tokens: 256000
+      input_price: 2
+      output_price: 8
+      supports_function_calling: true
+    - name: ai21/jamba-1-5-mini
+      max_input_tokens: 256000
+      input_price: 0.2
+      output_price: 0.4
+      supports_function_calling: true
     - name: cohere/command-r-plus
       max_input_tokens: 128000
       input_price: 3
diff --git a/src/client/mod.rs b/src/client/mod.rs
index 5a4a9ac..9d6dfb0 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -40,7 +40,8 @@ register_client!(
     (qianwen, "qianwen", QianwenConfig, QianwenClient),
 );
 
-pub const OPENAI_COMPATIBLE_PLATFORMS: [(&str, &str); 15] = [
+pub const OPENAI_COMPATIBLE_PLATFORMS: [(&str, &str); 16] = [
+    ("ai21", "https://api.ai21.com/studio/v1"),
     ("deepinfra", "https://api.deepinfra.com/v1/openai"),
     ("deepseek", "https://api.deepseek.com"),
     ("fireworks", "https://api.fireworks.ai/inference/v1"),
diff --git a/src/client/openai.rs b/src/client/openai.rs
index 3e1707f..ed577bc 100644
--- a/src/client/openai.rs
+++ b/src/client/openai.rs
@@ -206,6 +206,7 @@ pub fn openai_build_chat_completions_body(data: ChatCompletionsData, model: &Mod
                             },
                         })
                     }).collect();
+                    let text = if text.is_empty() { Value::Null } else { text.into() };
                     let mut messages = vec![
                         json!({ "role": MessageRole::Assistant, "content": text, "tool_calls": tool_calls })
                     ];