test: Added a unittest

Author: asafgardin

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

@@ -2,6 +2,7 @@
 from pytest_mock import MockerFixture
 
 from ai21.ai21_http_client import AI21HTTPClient
+from ai21.clients.studio.resources.chat import ChatCompletions
 from ai21.clients.studio.resources.studio_answer import StudioAnswer
 from ai21.clients.studio.resources.studio_chat import StudioChat
 from ai21.clients.studio.resources.studio_completion import StudioCompletion
@@ -43,7 +44,13 @@
     SummarizeBySegmentResponse,
     SegmentSummary,
 )
+from ai21.models.chat import (
+    ChatMessage as ChatCompletionChatMessage,
+    ChatCompletionResponse,
+    ChatCompletionResponseChoice,
+)
 from ai21.models.responses.segmentation_response import Segment
+from ai21.models.usage_info import UsageInfo
 from ai21.utils.typing import to_lower_camel_case
 
 
@@ -110,6 +117,45 @@ def get_studio_chat():
     )
 
 
+def get_chat_completions():
+    _DUMMY_MODEL = "dummy-chat-model"
+    _DUMMY_MESSAGES = [
+        ChatCompletionChatMessage(content="Hello, I need help with a signup process.", role=RoleType.USER),
+        ChatCompletionChatMessage(
+            content="Hi Alice, I can help you with that. What seems to be the problem?",
+            role=RoleType.ASSISTANT,
+        ),
+    ]
+
+    return (
+        ChatCompletions,
+        {"model": _DUMMY_MODEL, "messages": _DUMMY_MESSAGES},
+        "chat/complete",
+        {
+            "model": _DUMMY_MODEL,
+            "messages": [message.to_dict() for message in _DUMMY_MESSAGES],
+        },
+        ChatCompletionResponse(
+            id="some-id",
+            choices=[
+                ChatCompletionResponseChoice(
+                    index=0,
+                    message=ChatCompletionChatMessage(
+                        content="Hello, I need help with a signup process.", role=RoleType.USER
+                    ),
+                    finish_reason="dummy_reason",
+                    logprobs=None,
+                )
+            ],
+            usage=UsageInfo(
+                prompt_tokens=10,
+                completion_tokens=20,
+                total_tokens=30,
+            ),
+        ),
+    )
+
+
 def get_studio_completion(**kwargs):
     _DUMMY_MODEL = "dummy-completion-model"
     _DUMMY_PROMPT = "dummy-prompt"

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

@@ -17,6 +17,7 @@
     get_studio_segmentation,
     get_studio_summarization,
     get_studio_summarize_by_segment,
+    get_chat_completions,
 )
 
 _BASE_URL = "https://test.api.ai21.com/studio/v1"
@@ -31,6 +32,7 @@ class TestStudioResources:
         ids=[
             "studio_answer",
             "studio_chat",
+            "chat_completions",
             "studio_completion",
             "studio_completion_with_extra_args",
             "studio_embed",
@@ -45,6 +47,7 @@ class TestStudioResources:
         argvalues=[
             (get_studio_answer()),
             (get_studio_chat()),
+            (get_chat_completions()),
             (get_studio_completion()),
             (get_studio_completion(temperature=0.5, max_tokens=50)),
             (get_studio_embed()),