fix: revert completion message to dataclass

ai21/clients/studio/resources/chat/chat_completions.py

@@ -50,7 +50,7 @@ def _create_body(
         return remove_not_given(
             {
                 "model": model,
-                "messages": messages,
+                "messages": [message.to_dict() for message in messages],
                 "temperature": temperature,
                 "maxTokens": max_tokens,
                 "topP": top_p,

ai21/models/chat/chat_message.py

@@ -1,7 +1,11 @@
 from __future__ import annotations
-from typing_extensions import TypedDict, Literal
 
+from dataclasses import dataclass
 
-class ChatMessage(TypedDict):
-    role: str | Literal["user", "assistant", "system"]
+from ai21.models.ai21_base_model_mixin import AI21BaseModelMixin
+
+
+@dataclass
+class ChatMessage(AI21BaseModelMixin):
+    role: str
     content: str

tests/integration_tests/clients/studio/test_chat_completions.py

@@ -18,18 +18,18 @@
 # TODO: When the api is officially released, update the test to assert the actual response
 @pytest.mark.skip(reason="API is not officially released")
 def test_chat_completion():
-    num_results = 5
     messages = _MESSAGES
 
     client = AI21Client()
     response = client.chat.completions.create(
         model=_MODEL,
         messages=messages,
-        num_results=num_results,
         max_tokens=64,
         temperature=0.7,
         stop=["\n"],
         top_p=0.3,
     )
 
     assert isinstance(response, ChatCompletionResponse)
+    assert response.choices[0].message.content
+    assert response.choices[0].message.role

tests/unittests/clients/studio/resources/conftest.py

@@ -126,14 +126,15 @@ def get_chat_completions():
             role="assistant",
         ),
     ]
+    _EXPECTED_SERIALIZED_MESSAGES = [message.to_dict() for message in _DUMMY_MESSAGES]
 
     return (
         ChatCompletions,
         {"model": _DUMMY_MODEL, "messages": _DUMMY_MESSAGES},
         "chat/completions",
         {
             "model": _DUMMY_MODEL,
-            "messages": _DUMMY_MESSAGES,
+            "messages": _EXPECTED_SERIALIZED_MESSAGES,
         },
         ChatCompletionResponse(
             id="some-id",