refactor: Removed old model names (#307)

Josephasafg · web-flow · commit 3e76e8be76b2 · 2025-08-12T11:58:36.000Z
* refactor: Removed old model names

* ci: Fix action permission
diff --git a/.github/workflows/pr-approval.yml b/.github/workflows/pr-approval.yml
@@ -13,7 +13,7 @@ jobs:
     if: github.event.review.state == 'APPROVED'
     runs-on: ubuntu-latest
     permissions:
-      contents: read
+      contents: write
       pull-requests: write
       statuses: write
 
diff --git a/README.md b/README.md
@@ -81,7 +81,7 @@ messages = [
 
 chat_completions = client.chat.completions.create(
     messages=messages,
-    model="jamba-mini-1.6-2025-03",
+    model="jamba-mini",
 )
 ```
 
@@ -111,7 +111,7 @@ client = AsyncAI21Client(
 async def main():
     response = await client.chat.completions.create(
         messages=messages,
-        model="jamba-mini-1.6-2025-03",
+        model="jamba-mini",
     )
 
     print(response)
@@ -235,7 +235,7 @@ client = AsyncAI21Client()
 async def main():
     response = await client.chat.completions.create(
         messages=messages,
-        model="jamba-mini-1.6-2025-03",
+        model="jamba-mini",
         stream=True,
     )
     async for chunk in response:
@@ -590,7 +590,7 @@ messages = [
 ]
 
 response = client.chat.completions.create(
-  model="jamba-mini-1.6-2025-03",
+  model="jamba-mini",
   messages=messages,
 )
 ```
diff --git a/ai21/clients/studio/resources/chat/base_chat_completions.py b/ai21/clients/studio/resources/chat/base_chat_completions.py
@@ -15,9 +15,9 @@
 
 
 _MODEL_DEPRECATION_WARNING = """
-The 'jamba-1.5-mini' and 'jamba-1.5-large' models are deprecated and will
+The 'jamba-mini-1.6' and 'jamba-large-1.6' models are deprecated and will
 be removed in a future version.
-Please use jamba-mini-1.6-2025-03 or jamba-large-1.6-2025-03 instead.
+Please use jamba-mini or jamba-large instead.
 """
 
 
@@ -28,7 +28,7 @@ def _check_model(self, model: Optional[str]) -> str:
         if not model:
             raise ValueError("model should be provided 'create' method call")
 
-        if model in ["jamba-1.5-mini", "jamba-1.5-large"]:
+        if model in ["jamba-mini-1.6", "jamba-large-1.6"]:
             warnings.warn(
                 _MODEL_DEPRECATION_WARNING,
                 DeprecationWarning,
diff --git a/examples/studio/batches/batches.py b/examples/studio/batches/batches.py
@@ -13,7 +13,7 @@
         "method": "POST",
         "url": "/v1/chat/completions",
         "body": {
-            "model": "jamba-1.5",
+            "model": "jamba-mini",
             "messages": [{"role": "user", "content": "What is your favorite color?"}],
         },
     },
@@ -22,7 +22,7 @@
         "method": "POST",
         "url": "/v1/chat/completions",
         "body": {
-            "model": "jamba-1.5",
+            "model": "jamba-mini",
             "messages": [{"role": "user", "content": "Tell me about your hobbies."}],
         },
     },
@@ -31,7 +31,7 @@
         "method": "POST",
         "url": "/v1/chat/completions",
         "body": {
-            "model": "jamba-1.5",
+            "model": "jamba-mini",
             "messages": [{"role": "user", "content": "Tell me about your favorite food."}],
         },
     },
diff --git a/examples/studio/chat/async_chat_completions.py b/examples/studio/chat/async_chat_completions.py
@@ -18,7 +18,7 @@
 async def main():
     response = await client.chat.completions.create(
         messages=messages,
-        model="jamba-mini-1.6-2025-03",
+        model="jamba-mini",
         max_tokens=100,
         temperature=0.7,
         top_p=1.0,
diff --git a/examples/studio/chat/async_stream_chat_completions.py b/examples/studio/chat/async_stream_chat_completions.py
@@ -18,7 +18,7 @@
 async def main():
     response = await client.chat.completions.create(
         messages=messages,
-        model="jamba-mini-1.6-2025-03",
+        model="jamba-mini",
         max_tokens=100,
         stream=True,
     )
diff --git a/examples/studio/chat/chat_completions.py b/examples/studio/chat/chat_completions.py
@@ -14,7 +14,7 @@
 
 response = client.chat.completions.create(
     messages=messages,
-    model="jamba-mini-1.6-2025-03",
+    model="jamba-mini",
     max_tokens=100,
     temperature=0.7,
     top_p=1.0,
diff --git a/examples/studio/chat/chat_documents.py b/examples/studio/chat/chat_documents.py
@@ -42,7 +42,7 @@
 
 response = client.chat.completions.create(
     messages=messages,
-    model="jamba-mini-1.6-2025-03",
+    model="jamba-mini",
     documents=documents,
 )
 
diff --git a/examples/studio/chat/chat_function_calling.py b/examples/studio/chat/chat_function_calling.py
@@ -46,7 +46,7 @@ def get_order_delivery_date(order_id: str) -> str:
 
 client = AI21Client()
 
-response = client.chat.completions.create(messages=messages, model="jamba-large-1.6-2025-03", tools=tools)
+response = client.chat.completions.create(messages=messages, model="jamba-large", tools=tools)
 
 """ AI models can be error-prone, it's crucial to ensure that the tool calls align with the expectations.
 The below code snippet demonstrates how to handle tool calls in the response and invoke the tool function
@@ -79,5 +79,5 @@ def get_order_delivery_date(order_id: str) -> str:
     tool_message = ToolMessage(role="tool", tool_call_id=tool_calls[0].id, content=delivery_date)
     messages.append(tool_message)
 
-    response = client.chat.completions.create(messages=messages, model="jamba-large-1.6-2025-03", tools=tools)
+    response = client.chat.completions.create(messages=messages, model="jamba-large", tools=tools)
     print(response.choices[0].message.content)
diff --git a/examples/studio/chat/chat_function_calling_multiple_tools.py b/examples/studio/chat/chat_function_calling_multiple_tools.py
@@ -75,7 +75,7 @@ def get_sunset_hour(place: str, date: str) -> str:
 
 client = AI21Client()
 
-response = client.chat.completions.create(messages=messages, model="jamba-large-1.6-2025-03", tools=tools)
+response = client.chat.completions.create(messages=messages, model="jamba-large", tools=tools)
 
 """ AI models can be error-prone, it's crucial to ensure that the tool calls align with the expectations.
 The below code snippet demonstrates how to handle tool calls in the response and invoke the tool function
@@ -123,5 +123,5 @@ def get_sunset_hour(place: str, date: str) -> str:
         tool_message = ToolMessage(role="tool", tool_call_id=tool_id_called, content=str(result))
         messages.append(tool_message)
 
-    response = client.chat.completions.create(messages=messages, model="jamba-large-1.6-2025-03", tools=tools)
+    response = client.chat.completions.create(messages=messages, model="jamba-large", tools=tools)
     print(response.choices[0].message.content)
diff --git a/examples/studio/chat/chat_response_format.py b/examples/studio/chat/chat_response_format.py
@@ -39,7 +39,7 @@ class ZooTicketsOrder(BaseModel):
 
 response = client.chat.completions.create(
     messages=messages,
-    model="jamba-large-1.6-2025-03",
+    model="jamba-large",
     max_tokens=800,
     temperature=0,
     response_format=ResponseFormat(type="json_object"),
diff --git a/examples/studio/chat/stream_chat_completions.py b/examples/studio/chat/stream_chat_completions.py
@@ -14,7 +14,7 @@
 
 response = client.chat.completions.create(
     messages=messages,
-    model="jamba-large-1.6-2025-03",
+    model="jamba-large",
     max_tokens=100,
     stream=True,
 )
diff --git a/tests/integration_tests/clients/studio/test_chat_completions.py b/tests/integration_tests/clients/studio/test_chat_completions.py
@@ -16,7 +16,7 @@
 )
 
 
-_MODEL = "jamba-mini-1.6-2025-03"
+_MODEL = "jamba-mini"
 _MESSAGES = [
     ChatMessage(
         content="Hello, I need help studying for the coming test, can you teach me about the US constitution? ",
diff --git a/tests/unittests/clients/studio/resources/chat/test_chat_completions.py b/tests/unittests/clients/studio/resources/chat/test_chat_completions.py
@@ -23,7 +23,7 @@
 def test_chat_create__when_bad_import_to_chat_message__raise_error():
     with pytest.raises(ValueError) as e:
         AI21Client(api_key=_FAKE_API_KEY).chat.completions.create(
-            model="jamba-1.5",
+            model="jamba-mini",
             messages=[ChatMessage(role=RoleType.USER, text="Hello")],
             system="System Test",
         )
@@ -38,7 +38,7 @@ def test_chat_create__when_bad_import_to_chat_message__raise_error():
 async def test_async_chat_create__when_bad_import_to_chat_message__raise_error():
     with pytest.raises(ValueError) as e:
         await AsyncAI21Client(api_key=_FAKE_API_KEY).chat.completions.create(
-            model="jamba-1.5",
+            model="jamba-mini",
             messages=[ChatMessage(role=RoleType.USER, text="Hello")],
             system="System Test",
         )
@@ -72,7 +72,7 @@ def test_chat_completion_basic_happy_flow(mocker: MockerFixture) -> None:
     mocked_client.send.return_value = httpx.Response(status_code=200, json=_FAKE_BASIC_HAPPY_FLOW_RESPONSE_JSON)
     client = AI21Client(api_key=_FAKE_API_KEY, http_client=mocked_client)
     response: ChatCompletionResponse = client.chat.completions.create(
-        model="jamba-mini-1.6-2025-03", messages=[UserMessage(role="user", content="Hello")]
+        model="jamba-mini", messages=[UserMessage(role="user", content="Hello")]
     )
     assert response.choices[0].message.content == _FAKE_BASIC_HAPPY_FLOW_EXPECTED_CONTENT
 
@@ -85,7 +85,7 @@ async def test_async_chat_completion_basic_happy_flow(mocker: MockerFixture) ->
     )
     async_client = AsyncAI21Client(api_key=_FAKE_API_KEY, http_client=mocked_client)
     response: ChatCompletionResponse = await async_client.chat.completions.create(
-        model="jamba-mini-1.6-2025-03", messages=[UserMessage(role="user", content="Hello")]
+        model="jamba-mini", messages=[UserMessage(role="user", content="Hello")]
     )
     assert response.choices[0].message.content == _FAKE_BASIC_HAPPY_FLOW_EXPECTED_CONTENT
 
@@ -149,7 +149,7 @@ def test_chat_completion_with_tool_calls_happy_flow(mocker: MockerFixture) -> No
     mocked_client.send.return_value = httpx.Response(status_code=200, json=_FAKE_TOOL_CALLS_HAPPY_FLOW_RESPONSE_JSON)
     client = AI21Client(api_key=_FAKE_API_KEY, http_client=mocked_client)
     response = client.chat.completions.create(
-        model="jamba-mini-1.6-2025-03", messages=_FAKE_TOOL_CALL_TEST_MESSAGES, tools=_FAKE_TOOL_CALL_TEST_TOOLS
+        model="jamba-mini", messages=_FAKE_TOOL_CALL_TEST_MESSAGES, tools=_FAKE_TOOL_CALL_TEST_TOOLS
     )
     assert response.choices[0].message.tool_calls[0].function.name == _FAKE_TOOL_CALL_TEST_EXPECTED_FUNCTION_NAME
     assert (
@@ -165,7 +165,7 @@ async def test_async_chat_completion_with_tool_calls_happy_flow(mocker: MockerFi
     )
     async_client = AsyncAI21Client(api_key=_FAKE_API_KEY, http_client=mocked_client)
     response = await async_client.chat.completions.create(
-        model="jamba-mini-1.6-2025-03", messages=_FAKE_TOOL_CALL_TEST_MESSAGES, tools=_FAKE_TOOL_CALL_TEST_TOOLS
+        model="jamba-mini", messages=_FAKE_TOOL_CALL_TEST_MESSAGES, tools=_FAKE_TOOL_CALL_TEST_TOOLS
     )
     assert response.choices[0].message.tool_calls[0].function.name == _FAKE_TOOL_CALL_TEST_EXPECTED_FUNCTION_NAME
     assert (
@@ -231,7 +231,7 @@ def test_chat_completion_with_documents_happy_flow(mocker: MockerFixture) -> Non
     mocked_client.send.return_value = httpx.Response(status_code=200, json=_FAKE_DOCUMENTS_TEST_RESPONSE_JSON)
     client = AI21Client(api_key=_FAKE_API_KEY, http_client=mocked_client)
     response = client.chat.completions.create(
-        model="jamba-mini-1.6-2025-03", messages=_FAKE_DOCUMENTS_TEST_MESSAGES, documents=_FAKE_DOCUMENTS_TEST_DOCUMENTS
+        model="jamba-mini", messages=_FAKE_DOCUMENTS_TEST_MESSAGES, documents=_FAKE_DOCUMENTS_TEST_DOCUMENTS
     )
     assert response.choices[0].message.content == _FAKE_DOCUMENTS_TEST_EXPECTED_CONTENT
 
@@ -244,7 +244,7 @@ async def test_async_chat_completion_with_documents_happy_flow(mocker: MockerFix
     )
     async_client = AsyncAI21Client(api_key=_FAKE_API_KEY, http_client=mocked_client)
     response = await async_client.chat.completions.create(
-        model="jamba-mini-1.6-2025-03", messages=_FAKE_DOCUMENTS_TEST_MESSAGES, documents=_FAKE_DOCUMENTS_TEST_DOCUMENTS
+        model="jamba-mini", messages=_FAKE_DOCUMENTS_TEST_MESSAGES, documents=_FAKE_DOCUMENTS_TEST_DOCUMENTS
     )
     assert response.choices[0].message.content == _FAKE_DOCUMENTS_TEST_EXPECTED_CONTENT
 
@@ -301,7 +301,7 @@ def test_chat_completion_response_format_json_happy_flow(mocker: MockerFixture)
     )
     client = AI21Client(api_key=_FAKE_API_KEY, http_client=mocked_client)
     response = client.chat.completions.create(
-        model="jamba-mini-1.6-2025-03",
+        model="jamba-mini",
         messages=_FAKE_RESPONSE_FORMAT_JSON_TEST_MESSAGES,
         response_format=ResponseFormat(type="json_object"),
     )
@@ -316,7 +316,7 @@ async def test_async_chat_completion_response_format_json_happy_flow(mocker: Moc
     )
     async_client = AsyncAI21Client(api_key=_FAKE_API_KEY, http_client=mocked_client)
     response = await async_client.chat.completions.create(
-        model="jamba-mini-1.6-2025-03",
+        model="jamba-mini",
         messages=_FAKE_RESPONSE_FORMAT_JSON_TEST_MESSAGES,
         response_format=ResponseFormat(type="json_object"),
     )

Original file line number	Diff line number	Diff line change
`@@ -81,7 +81,7 @@ messages = [`
`81`	`81`
`82`	`82`	`chat_completions = client.chat.completions.create(`
`83`	`83`	`messages=messages,`
`84`		`- model="jamba-mini-1.6-2025-03",`
	`84`	`+ model="jamba-mini",`
`85`	`85`	`)`
`86`	`86`	```
`87`	`87`
`@@ -111,7 +111,7 @@ client = AsyncAI21Client(`
`111`	`111`	`async def main():`
`112`	`112`	`response = await client.chat.completions.create(`
`113`	`113`	`messages=messages,`
`114`		`- model="jamba-mini-1.6-2025-03",`
	`114`	`+ model="jamba-mini",`
`115`	`115`	`)`
`116`	`116`
`117`	`117`	`print(response)`
`@@ -235,7 +235,7 @@ client = AsyncAI21Client()`
`235`	`235`	`async def main():`
`236`	`236`	`response = await client.chat.completions.create(`
`237`	`237`	`messages=messages,`
`238`		`- model="jamba-mini-1.6-2025-03",`
	`238`	`+ model="jamba-mini",`
`239`	`239`	`stream=True,`
`240`	`240`	`)`
`241`	`241`	`async for chunk in response:`
`@@ -590,7 +590,7 @@ messages = [`
`590`	`590`	`]`
`591`	`591`
`592`	`592`	`response = client.chat.completions.create(`
`593`		`- model="jamba-mini-1.6-2025-03",`
	`593`	`+ model="jamba-mini",`
`594`	`594`	`messages=messages,`
`595`	`595`	`)`
`596`	`596`	```
Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,7 @@`
`18`	`18`	`async def main():`
`19`	`19`	`response = await client.chat.completions.create(`
`20`	`20`	`messages=messages,`
`21`		`- model="jamba-mini-1.6-2025-03",`
	`21`	`+ model="jamba-mini",`
`22`	`22`	`max_tokens=100,`
`23`	`23`	`stream=True,`
`24`	`24`	`)`
Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,7 @@`
`42`	`42`
`43`	`43`	`response = client.chat.completions.create(`
`44`	`44`	`messages=messages,`
`45`		`- model="jamba-mini-1.6-2025-03",`
	`45`	`+ model="jamba-mini",`
`46`	`46`	`documents=documents,`
`47`	`47`	`)`
`48`	`48`