From 4f3fd22dd4fb59f4f965915e9d8614f2a7e8e5be Mon Sep 17 00:00:00 2001
From: Xi Yan <xiyan@meta.com>
Date: Tue, 14 Jan 2025 18:08:30 -0800
Subject: [PATCH] add schema type register

---
 docs/resources/llama-stack-spec.html     | 100 +++++++++++++++--------
 docs/resources/llama-stack-spec.yaml     |  71 ++++++++++------
 llama_stack/apis/common/content_types.py |   2 +
 3 files changed, 114 insertions(+), 59 deletions(-)

diff --git a/docs/resources/llama-stack-spec.html b/docs/resources/llama-stack-spec.html
index 5ed8701a4..4f6ec51dc 100644
--- a/docs/resources/llama-stack-spec.html
+++ b/docs/resources/llama-stack-spec.html
@@ -4228,47 +4228,53 @@
             "ContentDelta": {
                 "oneOf": [
                     {
-                        "type": "object",
-                        "properties": {
-                            "type": {
-                                "type": "string",
-                                "const": "text",
-                                "default": "text"
-                            },
-                            "text": {
-                                "type": "string"
-                            }
-                        },
-                        "additionalProperties": false,
-                        "required": [
-                            "type",
-                            "text"
-                        ]
+                        "$ref": "#/components/schemas/TextDelta"
                     },
                     {
-                        "type": "object",
-                        "properties": {
-                            "type": {
-                                "type": "string",
-                                "const": "image",
-                                "default": "image"
-                            },
-                            "data": {
-                                "type": "string",
-                                "contentEncoding": "base64"
-                            }
-                        },
-                        "additionalProperties": false,
-                        "required": [
-                            "type",
-                            "data"
-                        ]
+                        "$ref": "#/components/schemas/ImageDelta"
                     },
                     {
                         "$ref": "#/components/schemas/ToolCallDelta"
                     }
                 ]
             },
+            "ImageDelta": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "image",
+                        "default": "image"
+                    },
+                    "data": {
+                        "type": "string",
+                        "contentEncoding": "base64"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type",
+                    "data"
+                ]
+            },
+            "TextDelta": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "text",
+                        "default": "text"
+                    },
+                    "text": {
+                        "type": "string"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type",
+                    "text"
+                ]
+            },
             "TokenLogProbs": {
                 "type": "object",
                 "properties": {
@@ -7261,6 +7267,9 @@
                     "shuffle": {
                         "type": "boolean"
                     },
+                    "data_format": {
+                        "$ref": "#/components/schemas/DatasetFormat"
+                    },
                     "validation_dataset_id": {
                         "type": "string"
                     },
@@ -7277,7 +7286,15 @@
                 "required": [
                     "dataset_id",
                     "batch_size",
-                    "shuffle"
+                    "shuffle",
+                    "data_format"
+                ]
+            },
+            "DatasetFormat": {
+                "type": "string",
+                "enum": [
+                    "instruct",
+                    "dialog"
                 ]
             },
             "EfficiencyConfig": {
@@ -8827,6 +8844,10 @@
             "name": "Dataset",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/Dataset\" />"
         },
+        {
+            "name": "DatasetFormat",
+            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/DatasetFormat\" />"
+        },
         {
             "name": "DatasetIO"
         },
@@ -8895,6 +8916,10 @@
             "name": "ImageContentItem",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/ImageContentItem\" />"
         },
+        {
+            "name": "ImageDelta",
+            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/ImageDelta\" />"
+        },
         {
             "name": "Inference"
         },
@@ -9243,6 +9268,10 @@
             "name": "TextContentItem",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/TextContentItem\" />"
         },
+        {
+            "name": "TextDelta",
+            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/TextDelta\" />"
+        },
         {
             "name": "TokenLogProbs",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/TokenLogProbs\" />"
@@ -9444,6 +9473,7 @@
                 "DPOAlignmentConfig",
                 "DataConfig",
                 "Dataset",
+                "DatasetFormat",
                 "DeleteAgentsRequest",
                 "DeleteAgentsSessionRequest",
                 "EfficiencyConfig",
@@ -9458,6 +9488,7 @@
                 "GraphMemoryBankParams",
                 "HealthInfo",
                 "ImageContentItem",
+                "ImageDelta",
                 "InferenceStep",
                 "InsertDocumentsRequest",
                 "InterleavedContent",
@@ -9536,6 +9567,7 @@
                 "SyntheticDataGenerationResponse",
                 "SystemMessage",
                 "TextContentItem",
+                "TextDelta",
                 "TokenLogProbs",
                 "Tool",
                 "ToolCall",
diff --git a/docs/resources/llama-stack-spec.yaml b/docs/resources/llama-stack-spec.yaml
index 2a573959f..5c74713e6 100644
--- a/docs/resources/llama-stack-spec.yaml
+++ b/docs/resources/llama-stack-spec.yaml
@@ -570,31 +570,8 @@ components:
       type: object
     ContentDelta:
       oneOf:
-      - additionalProperties: false
-        properties:
-          text:
-            type: string
-          type:
-            const: text
-            default: text
-            type: string
-        required:
-        - type
-        - text
-        type: object
-      - additionalProperties: false
-        properties:
-          data:
-            contentEncoding: base64
-            type: string
-          type:
-            const: image
-            default: image
-            type: string
-        required:
-        - type
-        - data
-        type: object
+      - $ref: '#/components/schemas/TextDelta'
+      - $ref: '#/components/schemas/ImageDelta'
       - $ref: '#/components/schemas/ToolCallDelta'
     CreateAgentRequest:
       additionalProperties: false
@@ -680,6 +657,8 @@ components:
       properties:
         batch_size:
           type: integer
+        data_format:
+          $ref: '#/components/schemas/DatasetFormat'
         dataset_id:
           type: string
         packed:
@@ -696,6 +675,7 @@ components:
       - dataset_id
       - batch_size
       - shuffle
+      - data_format
       type: object
     Dataset:
       additionalProperties: false
@@ -735,6 +715,11 @@ components:
       - url
       - metadata
       type: object
+    DatasetFormat:
+      enum:
+      - instruct
+      - dialog
+      type: string
     DeleteAgentsRequest:
       additionalProperties: false
       properties:
@@ -960,6 +945,20 @@ components:
       required:
       - type
       type: object
+    ImageDelta:
+      additionalProperties: false
+      properties:
+        data:
+          contentEncoding: base64
+          type: string
+        type:
+          const: image
+          default: image
+          type: string
+      required:
+      - type
+      - data
+      type: object
     InferenceStep:
       additionalProperties: false
       properties:
@@ -2588,6 +2587,19 @@ components:
       - type
       - text
       type: object
+    TextDelta:
+      additionalProperties: false
+      properties:
+        text:
+          type: string
+        type:
+          const: text
+          default: text
+          type: string
+      required:
+      - type
+      - text
+      type: object
     TokenLogProbs:
       additionalProperties: false
       properties:
@@ -5548,6 +5560,8 @@ tags:
   name: DataConfig
 - description: <SchemaDefinition schemaRef="#/components/schemas/Dataset" />
   name: Dataset
+- description: <SchemaDefinition schemaRef="#/components/schemas/DatasetFormat" />
+  name: DatasetFormat
 - name: DatasetIO
 - name: Datasets
 - description: <SchemaDefinition schemaRef="#/components/schemas/DeleteAgentsRequest"
@@ -5592,6 +5606,8 @@ tags:
 - description: <SchemaDefinition schemaRef="#/components/schemas/ImageContentItem"
     />
   name: ImageContentItem
+- description: <SchemaDefinition schemaRef="#/components/schemas/ImageDelta" />
+  name: ImageDelta
 - name: Inference
 - description: <SchemaDefinition schemaRef="#/components/schemas/InferenceStep" />
   name: InferenceStep
@@ -5824,6 +5840,8 @@ tags:
 - description: <SchemaDefinition schemaRef="#/components/schemas/TextContentItem"
     />
   name: TextContentItem
+- description: <SchemaDefinition schemaRef="#/components/schemas/TextDelta" />
+  name: TextDelta
 - description: <SchemaDefinition schemaRef="#/components/schemas/TokenLogProbs" />
   name: TokenLogProbs
 - description: <SchemaDefinition schemaRef="#/components/schemas/Tool" />
@@ -5978,6 +5996,7 @@ x-tagGroups:
   - DPOAlignmentConfig
   - DataConfig
   - Dataset
+  - DatasetFormat
   - DeleteAgentsRequest
   - DeleteAgentsSessionRequest
   - EfficiencyConfig
@@ -5992,6 +6011,7 @@ x-tagGroups:
   - GraphMemoryBankParams
   - HealthInfo
   - ImageContentItem
+  - ImageDelta
   - InferenceStep
   - InsertDocumentsRequest
   - InterleavedContent
@@ -6070,6 +6090,7 @@ x-tagGroups:
   - SyntheticDataGenerationResponse
   - SystemMessage
   - TextContentItem
+  - TextDelta
   - TokenLogProbs
   - Tool
   - ToolCall
diff --git a/llama_stack/apis/common/content_types.py b/llama_stack/apis/common/content_types.py
index 3b61fa243..b845d09dd 100644
--- a/llama_stack/apis/common/content_types.py
+++ b/llama_stack/apis/common/content_types.py
@@ -64,11 +64,13 @@ InterleavedContent = register_schema(
 )
 
 
+@json_schema_type
 class TextDelta(BaseModel):
     type: Literal["text"] = "text"
     text: str
 
 
+@json_schema_type
 class ImageDelta(BaseModel):
     type: Literal["image"] = "image"
     data: bytes