From af1ba002275bbd8d1f746064540188c5a8e7f507 Mon Sep 17 00:00:00 2001
From: Botao Chen <markchen1015@meta.com>
Date: Sun, 12 Jan 2025 22:52:53 -0800
Subject: [PATCH] address_comment

---
 llama_stack/apis/common/type_system.py                    | 5 +++++
 .../providers/inline/post_training/common/validator.py    | 8 ++++++--
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/llama_stack/apis/common/type_system.py b/llama_stack/apis/common/type_system.py
index a653efef9..c5653f373 100644
--- a/llama_stack/apis/common/type_system.py
+++ b/llama_stack/apis/common/type_system.py
@@ -54,6 +54,11 @@ class AgentTurnInputType(BaseModel):
     type: Literal["agent_turn_input"] = "agent_turn_input"
 
 
+class DialogType(BaseModel):
+    # expects List[Message] for messages
+    type: Literal["dialog"] = "dialog"
+
+
 ParamType = register_schema(
     Annotated[
         Union[
diff --git a/llama_stack/providers/inline/post_training/common/validator.py b/llama_stack/providers/inline/post_training/common/validator.py
index e4456ea98..836e20c85 100644
--- a/llama_stack/providers/inline/post_training/common/validator.py
+++ b/llama_stack/providers/inline/post_training/common/validator.py
@@ -9,7 +9,11 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
-from llama_stack.apis.common.type_system import ChatCompletionInputType, StringType
+from llama_stack.apis.common.type_system import (
+    ChatCompletionInputType,
+    DialogType,
+    StringType,
+)
 from llama_stack.apis.datasets import Datasets
 from llama_stack.providers.utils.common.data_schema_validator import (
     ColumnName,
@@ -25,7 +29,7 @@ EXPECTED_DATASET_SCHEMA = {
     ],
     "dialog": [
         {
-            ColumnName.dialog.value: StringType(),
+            ColumnName.dialog.value: DialogType(),
         }
     ],
 }