feat: added json schema option for response_format (#217)

Author: adubovik

aidial_sdk/chat_completion/__init__.py

@@ -14,6 +14,10 @@
     MessageContentTextPart,
     Request,
     ResponseFormat,
+    ResponseFormatJsonObject,
+    ResponseFormatJsonSchema,
+    ResponseFormatJsonSchemaObject,
+    ResponseFormatText,
     Role,
 )
 from aidial_sdk.chat_completion.request import Stage as RequestStage

aidial_sdk/chat_completion/request.py

@@ -10,6 +10,7 @@
     ConstrainedFloat,
     ConstrainedInt,
     ConstrainedList,
+    Field,
     PositiveInt,
     StrictBool,
     StrictInt,
@@ -170,8 +171,35 @@ class ToolChoice(ExtraForbidModel):
     function: FunctionChoice
 
 
-class ResponseFormat(ExtraForbidModel):
-    type: Literal["text", "json_object"]
+class ResponseFormatText(ExtraForbidModel):
+    type: Literal["text"]
+
+
+class ResponseFormatJsonObject(ExtraForbidModel):
+    type: Literal["json_object"]
+
+
+class ResponseFormatJsonSchemaObject(ExtraForbidModel):
+    description: Optional[StrictStr] = None
+    name: StrictStr
+    schema_: Dict[str, Any] = Field(..., alias="schema")
+    strict: Optional[StrictBool] = False
+
+    def dict(self, *args, **kwargs):
+        kwargs["by_alias"] = True
+        return super().dict(*args, **kwargs)
+
+
+class ResponseFormatJsonSchema(ExtraForbidModel):
+    type: Literal["json_schema"]
+    json_schema: ResponseFormatJsonSchemaObject
+
+
+ResponseFormat = Union[
+    ResponseFormatText,
+    ResponseFormatJsonObject,
+    ResponseFormatJsonSchema,
+]
 
 
 class AzureChatCompletionRequest(ExtraForbidModel):

tests/test_serialization.py

@@ -1,6 +1,7 @@
 import json
 
-from aidial_sdk.chat_completion import Message, Role
+from aidial_sdk.chat_completion import Message, ResponseFormatJsonSchema, Role
+from aidial_sdk.chat_completion.request import ResponseFormatJsonSchemaObject
 
 
 def test_message_ser():
@@ -17,3 +18,53 @@ def test_message_deser():
     expected_obj = Message(role=Role.SYSTEM, content="test")
 
     assert actual_obj == expected_obj
+
+
+def test_response_format_serialization():
+    format_obj = ResponseFormatJsonSchema(
+        type="json_schema",
+        json_schema=ResponseFormatJsonSchemaObject(
+            description="desc",
+            name="name",
+            schema={"key": "value"},
+        ),
+    )
+
+    actual_dict = format_obj.dict()
+
+    expected_dict = {
+        "type": "json_schema",
+        "json_schema": {
+            "description": "desc",
+            "name": "name",
+            "schema": {"key": "value"},
+            "strict": False,
+        },
+    }
+
+    assert actual_dict == expected_dict
+
+
+def test_response_format_deserialization():
+    format_dict = {
+        "type": "json_schema",
+        "json_schema": {
+            "description": "desc",
+            "name": "name",
+            "schema": {"key": "value"},
+        },
+    }
+
+    actual_obj = ResponseFormatJsonSchema.parse_obj(format_dict)
+
+    expected_obj = ResponseFormatJsonSchema(
+        type="json_schema",
+        json_schema=ResponseFormatJsonSchemaObject(
+            description="desc",
+            name="name",
+            schema={"key": "value"},
+            strict=False,
+        ),
+    )
+
+    assert actual_obj == expected_obj