feat: inherit all Enums from str to make JSON serialization possible (e…

…pam#132)

aidial_sdk/chat_completion/enums.py

@@ -1,14 +1,14 @@
 from enum import Enum
 
 
-class FinishReason(Enum):
+class FinishReason(str, Enum):
     STOP = "stop"
     LENGTH = "length"
     FUNCTION_CALL = "function_call"
     TOOL_CALLS = "tool_calls"
     CONTENT_FILTER = "content_filter"
 
 
-class Status(Enum):
+class Status(str, Enum):
     COMPLETED = "completed"
     FAILED = "failed"

aidial_sdk/chat_completion/request.py

@@ -48,7 +48,7 @@ class ToolCall(ExtraForbidModel):
     function: FunctionCall
 
 
-class Role(Enum):
+class Role(str, Enum):
     SYSTEM = "system"
     USER = "user"
     ASSISTANT = "assistant"

tests/test_serialization.py

@@ -0,0 +1,19 @@
+import json
+
+from aidial_sdk.chat_completion import Message, Role
+
+
+def test_message_ser():
+    msg_obj = Message(role=Role.SYSTEM, content="test")
+    actual_dict = msg_obj.dict(exclude_none=True)
+    expected_dict = {"role": "system", "content": "test"}
+
+    assert json.loads(json.dumps(actual_dict)) == expected_dict
+
+
+def test_message_deser():
+    msg_dict = {"role": "system", "content": "test"}
+    actual_obj = Message.parse_raw(json.dumps(msg_dict))
+    expected_obj = Message(role=Role.SYSTEM, content="test")
+
+    assert actual_obj == expected_obj