diff --git a/docs/resources/llama-stack-spec.html b/docs/resources/llama-stack-spec.html
index 38cabdd3e..459a53888 100644
--- a/docs/resources/llama-stack-spec.html
+++ b/docs/resources/llama-stack-spec.html
@@ -6067,6 +6067,76 @@
                     "step"
                 ]
             },
+            "AgentTurnInputType": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "agent_turn_input",
+                        "default": "agent_turn_input"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type"
+                ]
+            },
+            "ArrayType": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "array",
+                        "default": "array"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type"
+                ]
+            },
+            "BooleanType": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "boolean",
+                        "default": "boolean"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type"
+                ]
+            },
+            "ChatCompletionInputType": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "chat_completion_input",
+                        "default": "chat_completion_input"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type"
+                ]
+            },
+            "CompletionInputType": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "completion_input",
+                        "default": "completion_input"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type"
+                ]
+            },
             "Dataset": {
                 "type": "object",
                 "properties": {
@@ -6130,150 +6200,110 @@
                     "metadata"
                 ]
             },
+            "JsonType": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "json",
+                        "default": "json"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type"
+                ]
+            },
+            "NumberType": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "number",
+                        "default": "number"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type"
+                ]
+            },
+            "ObjectType": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "object",
+                        "default": "object"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type"
+                ]
+            },
             "ParamType": {
                 "oneOf": [
                     {
-                        "type": "object",
-                        "properties": {
-                            "type": {
-                                "type": "string",
-                                "const": "string",
-                                "default": "string"
-                            }
-                        },
-                        "additionalProperties": false,
-                        "required": [
-                            "type"
-                        ]
+                        "$ref": "#/components/schemas/StringType"
                     },
                     {
-                        "type": "object",
-                        "properties": {
-                            "type": {
-                                "type": "string",
-                                "const": "number",
-                                "default": "number"
-                            }
-                        },
-                        "additionalProperties": false,
-                        "required": [
-                            "type"
-                        ]
+                        "$ref": "#/components/schemas/NumberType"
                     },
                     {
-                        "type": "object",
-                        "properties": {
-                            "type": {
-                                "type": "string",
-                                "const": "boolean",
-                                "default": "boolean"
-                            }
-                        },
-                        "additionalProperties": false,
-                        "required": [
-                            "type"
-                        ]
+                        "$ref": "#/components/schemas/BooleanType"
                     },
                     {
-                        "type": "object",
-                        "properties": {
-                            "type": {
-                                "type": "string",
-                                "const": "array",
-                                "default": "array"
-                            }
-                        },
-                        "additionalProperties": false,
-                        "required": [
-                            "type"
-                        ]
+                        "$ref": "#/components/schemas/ArrayType"
                     },
                     {
-                        "type": "object",
-                        "properties": {
-                            "type": {
-                                "type": "string",
-                                "const": "object",
-                                "default": "object"
-                            }
-                        },
-                        "additionalProperties": false,
-                        "required": [
-                            "type"
-                        ]
+                        "$ref": "#/components/schemas/ObjectType"
                     },
                     {
-                        "type": "object",
-                        "properties": {
-                            "type": {
-                                "type": "string",
-                                "const": "json",
-                                "default": "json"
-                            }
-                        },
-                        "additionalProperties": false,
-                        "required": [
-                            "type"
-                        ]
+                        "$ref": "#/components/schemas/JsonType"
                     },
                     {
-                        "type": "object",
-                        "properties": {
-                            "type": {
-                                "type": "string",
-                                "const": "union",
-                                "default": "union"
-                            }
-                        },
-                        "additionalProperties": false,
-                        "required": [
-                            "type"
-                        ]
+                        "$ref": "#/components/schemas/UnionType"
                     },
                     {
-                        "type": "object",
-                        "properties": {
-                            "type": {
-                                "type": "string",
-                                "const": "chat_completion_input",
-                                "default": "chat_completion_input"
-                            }
-                        },
-                        "additionalProperties": false,
-                        "required": [
-                            "type"
-                        ]
+                        "$ref": "#/components/schemas/ChatCompletionInputType"
                     },
                     {
-                        "type": "object",
-                        "properties": {
-                            "type": {
-                                "type": "string",
-                                "const": "completion_input",
-                                "default": "completion_input"
-                            }
-                        },
-                        "additionalProperties": false,
-                        "required": [
-                            "type"
-                        ]
+                        "$ref": "#/components/schemas/CompletionInputType"
                     },
                     {
-                        "type": "object",
-                        "properties": {
-                            "type": {
-                                "type": "string",
-                                "const": "agent_turn_input",
-                                "default": "agent_turn_input"
-                            }
-                        },
-                        "additionalProperties": false,
-                        "required": [
-                            "type"
-                        ]
+                        "$ref": "#/components/schemas/AgentTurnInputType"
                     }
                 ]
             },
+            "StringType": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "string",
+                        "default": "string"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type"
+                ]
+            },
+            "UnionType": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "union",
+                        "default": "union"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type"
+                ]
+            },
             "EvalTask": {
                 "type": "object",
                 "properties": {
@@ -8922,6 +8952,10 @@
             "name": "AgentTool",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/AgentTool\" />"
         },
+        {
+            "name": "AgentTurnInputType",
+            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/AgentTurnInputType\" />"
+        },
         {
             "name": "AgentTurnResponseEvent",
             "description": "Streamed agent execution response.\n\n<SchemaDefinition schemaRef=\"#/components/schemas/AgentTurnResponseEvent\" />"
@@ -8965,6 +8999,10 @@
             "name": "AppendRowsRequest",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/AppendRowsRequest\" />"
         },
+        {
+            "name": "ArrayType",
+            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/ArrayType\" />"
+        },
         {
             "name": "BasicScoringFnParams",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/BasicScoringFnParams\" />"
@@ -8992,6 +9030,10 @@
             "name": "BenchmarkEvalTaskConfig",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/BenchmarkEvalTaskConfig\" />"
         },
+        {
+            "name": "BooleanType",
+            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/BooleanType\" />"
+        },
         {
             "name": "BuiltinTool",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/BuiltinTool\" />"
@@ -9000,6 +9042,10 @@
             "name": "CancelTrainingJobRequest",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/CancelTrainingJobRequest\" />"
         },
+        {
+            "name": "ChatCompletionInputType",
+            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/ChatCompletionInputType\" />"
+        },
         {
             "name": "ChatCompletionRequest",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/ChatCompletionRequest\" />"
@@ -9024,6 +9070,10 @@
             "name": "Checkpoint",
             "description": "Checkpoint created during training runs\n\n<SchemaDefinition schemaRef=\"#/components/schemas/Checkpoint\" />"
         },
+        {
+            "name": "CompletionInputType",
+            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/CompletionInputType\" />"
+        },
         {
             "name": "CompletionMessage",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/CompletionMessage\" />"
@@ -9166,6 +9216,10 @@
             "name": "JobStatus",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/JobStatus\" />"
         },
+        {
+            "name": "JsonType",
+            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/JsonType\" />"
+        },
         {
             "name": "KeyValueMemoryBank",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/KeyValueMemoryBank\" />"
@@ -9283,6 +9337,14 @@
         {
             "name": "Models"
         },
+        {
+            "name": "NumberType",
+            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/NumberType\" />"
+        },
+        {
+            "name": "ObjectType",
+            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/ObjectType\" />"
+        },
         {
             "name": "OptimizerConfig",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/OptimizerConfig\" />"
@@ -9490,6 +9552,10 @@
             "name": "StopReason",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/StopReason\" />"
         },
+        {
+            "name": "StringType",
+            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/StringType\" />"
+        },
         {
             "name": "StructuredLogEvent",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/StructuredLogEvent\" />"
@@ -9622,6 +9688,10 @@
             "name": "URL",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/URL\" />"
         },
+        {
+            "name": "UnionType",
+            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/UnionType\" />"
+        },
         {
             "name": "UnstructuredLogEvent",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/UnstructuredLogEvent\" />"
@@ -9682,6 +9752,7 @@
                 "AgentSessionCreateResponse",
                 "AgentStepResponse",
                 "AgentTool",
+                "AgentTurnInputType",
                 "AgentTurnResponseEvent",
                 "AgentTurnResponseStepCompletePayload",
                 "AgentTurnResponseStepProgressPayload",
@@ -9692,20 +9763,24 @@
                 "AggregationFunctionType",
                 "AppEvalTaskConfig",
                 "AppendRowsRequest",
+                "ArrayType",
                 "BasicScoringFnParams",
                 "BatchChatCompletionRequest",
                 "BatchChatCompletionResponse",
                 "BatchCompletionRequest",
                 "BatchCompletionResponse",
                 "BenchmarkEvalTaskConfig",
+                "BooleanType",
                 "BuiltinTool",
                 "CancelTrainingJobRequest",
+                "ChatCompletionInputType",
                 "ChatCompletionRequest",
                 "ChatCompletionResponse",
                 "ChatCompletionResponseEvent",
                 "ChatCompletionResponseEventType",
                 "ChatCompletionResponseStreamChunk",
                 "Checkpoint",
+                "CompletionInputType",
                 "CompletionMessage",
                 "CompletionRequest",
                 "CompletionResponse",
@@ -9737,6 +9812,7 @@
                 "InvokeToolRequest",
                 "Job",
                 "JobStatus",
+                "JsonType",
                 "KeyValueMemoryBank",
                 "KeyValueMemoryBankParams",
                 "KeywordMemoryBank",
@@ -9764,6 +9840,8 @@
                 "Model",
                 "ModelCandidate",
                 "ModelType",
+                "NumberType",
+                "ObjectType",
                 "OptimizerConfig",
                 "OptimizerType",
                 "PaginatedRowsResult",
@@ -9812,6 +9890,7 @@
                 "SpanStatus",
                 "SpanWithStatus",
                 "StopReason",
+                "StringType",
                 "StructuredLogEvent",
                 "SupervisedFineTuneRequest",
                 "SyntheticDataGenerateRequest",
@@ -9842,6 +9921,7 @@
                 "TrainingConfig",
                 "Turn",
                 "URL",
+                "UnionType",
                 "UnstructuredLogEvent",
                 "UserMessage",
                 "VectorMemoryBank",
diff --git a/docs/resources/llama-stack-spec.yaml b/docs/resources/llama-stack-spec.yaml
index 75bc25e94..9aeac6db3 100644
--- a/docs/resources/llama-stack-spec.yaml
+++ b/docs/resources/llama-stack-spec.yaml
@@ -105,6 +105,16 @@ components:
         - name
         - args
         type: object
+    AgentTurnInputType:
+      additionalProperties: false
+      properties:
+        type:
+          const: agent_turn_input
+          default: agent_turn_input
+          type: string
+      required:
+      - type
+      type: object
     AgentTurnResponseEvent:
       additionalProperties: false
       properties:
@@ -290,6 +300,16 @@ components:
       - dataset_id
       - rows
       type: object
+    ArrayType:
+      additionalProperties: false
+      properties:
+        type:
+          const: array
+          default: array
+          type: string
+      required:
+      - type
+      type: object
     BasicScoringFnParams:
       additionalProperties: false
       properties:
@@ -395,6 +415,16 @@ components:
       - type
       - eval_candidate
       type: object
+    BooleanType:
+      additionalProperties: false
+      properties:
+        type:
+          const: boolean
+          default: boolean
+          type: string
+      required:
+      - type
+      type: object
     BuiltinTool:
       enum:
       - brave_search
@@ -410,6 +440,16 @@ components:
       required:
       - job_uuid
       type: object
+    ChatCompletionInputType:
+      additionalProperties: false
+      properties:
+        type:
+          const: chat_completion_input
+          default: chat_completion_input
+          type: string
+      required:
+      - type
+      type: object
     ChatCompletionRequest:
       additionalProperties: false
       properties:
@@ -492,6 +532,16 @@ components:
       type: object
     Checkpoint:
       description: Checkpoint created during training runs
+    CompletionInputType:
+      additionalProperties: false
+      properties:
+        type:
+          const: completion_input
+          default: completion_input
+          type: string
+      required:
+      - type
+      type: object
     CompletionMessage:
       additionalProperties: false
       properties:
@@ -1007,6 +1057,16 @@ components:
       - failed
       - scheduled
       type: string
+    JsonType:
+      additionalProperties: false
+      properties:
+        type:
+          const: json
+          default: json
+          type: string
+      required:
+      - type
+      type: object
     KeyValueMemoryBank:
       additionalProperties: false
       properties:
@@ -1440,6 +1500,26 @@ components:
       - llm
       - embedding
       type: string
+    NumberType:
+      additionalProperties: false
+      properties:
+        type:
+          const: number
+          default: number
+          type: string
+      required:
+      - type
+      type: object
+    ObjectType:
+      additionalProperties: false
+      properties:
+        type:
+          const: object
+          default: object
+          type: string
+      required:
+      - type
+      type: object
     OptimizerConfig:
       additionalProperties: false
       properties:
@@ -1488,96 +1568,16 @@ components:
       type: object
     ParamType:
       oneOf:
-      - additionalProperties: false
-        properties:
-          type:
-            const: string
-            default: string
-            type: string
-        required:
-        - type
-        type: object
-      - additionalProperties: false
-        properties:
-          type:
-            const: number
-            default: number
-            type: string
-        required:
-        - type
-        type: object
-      - additionalProperties: false
-        properties:
-          type:
-            const: boolean
-            default: boolean
-            type: string
-        required:
-        - type
-        type: object
-      - additionalProperties: false
-        properties:
-          type:
-            const: array
-            default: array
-            type: string
-        required:
-        - type
-        type: object
-      - additionalProperties: false
-        properties:
-          type:
-            const: object
-            default: object
-            type: string
-        required:
-        - type
-        type: object
-      - additionalProperties: false
-        properties:
-          type:
-            const: json
-            default: json
-            type: string
-        required:
-        - type
-        type: object
-      - additionalProperties: false
-        properties:
-          type:
-            const: union
-            default: union
-            type: string
-        required:
-        - type
-        type: object
-      - additionalProperties: false
-        properties:
-          type:
-            const: chat_completion_input
-            default: chat_completion_input
-            type: string
-        required:
-        - type
-        type: object
-      - additionalProperties: false
-        properties:
-          type:
-            const: completion_input
-            default: completion_input
-            type: string
-        required:
-        - type
-        type: object
-      - additionalProperties: false
-        properties:
-          type:
-            const: agent_turn_input
-            default: agent_turn_input
-            type: string
-        required:
-        - type
-        type: object
+      - $ref: '#/components/schemas/StringType'
+      - $ref: '#/components/schemas/NumberType'
+      - $ref: '#/components/schemas/BooleanType'
+      - $ref: '#/components/schemas/ArrayType'
+      - $ref: '#/components/schemas/ObjectType'
+      - $ref: '#/components/schemas/JsonType'
+      - $ref: '#/components/schemas/UnionType'
+      - $ref: '#/components/schemas/ChatCompletionInputType'
+      - $ref: '#/components/schemas/CompletionInputType'
+      - $ref: '#/components/schemas/AgentTurnInputType'
     PostTrainingJob:
       additionalProperties: false
       properties:
@@ -2479,6 +2479,16 @@ components:
       - end_of_message
       - out_of_tokens
       type: string
+    StringType:
+      additionalProperties: false
+      properties:
+        type:
+          const: string
+          default: string
+          type: string
+      required:
+      - type
+      type: object
     StructuredLogEvent:
       additionalProperties: false
       properties:
@@ -3131,6 +3141,16 @@ components:
       required:
       - uri
       type: object
+    UnionType:
+      additionalProperties: false
+      properties:
+        type:
+          const: union
+          default: union
+          type: string
+      required:
+      - type
+      type: object
     UnstructuredLogEvent:
       additionalProperties: false
       properties:
@@ -5588,6 +5608,9 @@ tags:
   name: AgentStepResponse
 - description: <SchemaDefinition schemaRef="#/components/schemas/AgentTool" />
   name: AgentTool
+- description: <SchemaDefinition schemaRef="#/components/schemas/AgentTurnInputType"
+    />
+  name: AgentTurnInputType
 - description: 'Streamed agent execution response.
 
 
@@ -5624,6 +5647,8 @@ tags:
 - description: <SchemaDefinition schemaRef="#/components/schemas/AppendRowsRequest"
     />
   name: AppendRowsRequest
+- description: <SchemaDefinition schemaRef="#/components/schemas/ArrayType" />
+  name: ArrayType
 - description: <SchemaDefinition schemaRef="#/components/schemas/BasicScoringFnParams"
     />
   name: BasicScoringFnParams
@@ -5643,11 +5668,16 @@ tags:
 - description: <SchemaDefinition schemaRef="#/components/schemas/BenchmarkEvalTaskConfig"
     />
   name: BenchmarkEvalTaskConfig
+- description: <SchemaDefinition schemaRef="#/components/schemas/BooleanType" />
+  name: BooleanType
 - description: <SchemaDefinition schemaRef="#/components/schemas/BuiltinTool" />
   name: BuiltinTool
 - description: <SchemaDefinition schemaRef="#/components/schemas/CancelTrainingJobRequest"
     />
   name: CancelTrainingJobRequest
+- description: <SchemaDefinition schemaRef="#/components/schemas/ChatCompletionInputType"
+    />
+  name: ChatCompletionInputType
 - description: <SchemaDefinition schemaRef="#/components/schemas/ChatCompletionRequest"
     />
   name: ChatCompletionRequest
@@ -5676,6 +5706,9 @@ tags:
 
     <SchemaDefinition schemaRef="#/components/schemas/Checkpoint" />'
   name: Checkpoint
+- description: <SchemaDefinition schemaRef="#/components/schemas/CompletionInputType"
+    />
+  name: CompletionInputType
 - description: <SchemaDefinition schemaRef="#/components/schemas/CompletionMessage"
     />
   name: CompletionMessage
@@ -5770,6 +5803,8 @@ tags:
   name: Job
 - description: <SchemaDefinition schemaRef="#/components/schemas/JobStatus" />
   name: JobStatus
+- description: <SchemaDefinition schemaRef="#/components/schemas/JsonType" />
+  name: JsonType
 - description: <SchemaDefinition schemaRef="#/components/schemas/KeyValueMemoryBank"
     />
   name: KeyValueMemoryBank
@@ -5847,6 +5882,10 @@ tags:
 - description: <SchemaDefinition schemaRef="#/components/schemas/ModelType" />
   name: ModelType
 - name: Models
+- description: <SchemaDefinition schemaRef="#/components/schemas/NumberType" />
+  name: NumberType
+- description: <SchemaDefinition schemaRef="#/components/schemas/ObjectType" />
+  name: ObjectType
 - description: <SchemaDefinition schemaRef="#/components/schemas/OptimizerConfig"
     />
   name: OptimizerConfig
@@ -5988,6 +6027,8 @@ tags:
   name: SpanWithStatus
 - description: <SchemaDefinition schemaRef="#/components/schemas/StopReason" />
   name: StopReason
+- description: <SchemaDefinition schemaRef="#/components/schemas/StringType" />
+  name: StringType
 - description: <SchemaDefinition schemaRef="#/components/schemas/StructuredLogEvent"
     />
   name: StructuredLogEvent
@@ -6080,6 +6121,8 @@ tags:
   name: Turn
 - description: <SchemaDefinition schemaRef="#/components/schemas/URL" />
   name: URL
+- description: <SchemaDefinition schemaRef="#/components/schemas/UnionType" />
+  name: UnionType
 - description: <SchemaDefinition schemaRef="#/components/schemas/UnstructuredLogEvent"
     />
   name: UnstructuredLogEvent
@@ -6126,6 +6169,7 @@ x-tagGroups:
   - AgentSessionCreateResponse
   - AgentStepResponse
   - AgentTool
+  - AgentTurnInputType
   - AgentTurnResponseEvent
   - AgentTurnResponseStepCompletePayload
   - AgentTurnResponseStepProgressPayload
@@ -6136,20 +6180,24 @@ x-tagGroups:
   - AggregationFunctionType
   - AppEvalTaskConfig
   - AppendRowsRequest
+  - ArrayType
   - BasicScoringFnParams
   - BatchChatCompletionRequest
   - BatchChatCompletionResponse
   - BatchCompletionRequest
   - BatchCompletionResponse
   - BenchmarkEvalTaskConfig
+  - BooleanType
   - BuiltinTool
   - CancelTrainingJobRequest
+  - ChatCompletionInputType
   - ChatCompletionRequest
   - ChatCompletionResponse
   - ChatCompletionResponseEvent
   - ChatCompletionResponseEventType
   - ChatCompletionResponseStreamChunk
   - Checkpoint
+  - CompletionInputType
   - CompletionMessage
   - CompletionRequest
   - CompletionResponse
@@ -6181,6 +6229,7 @@ x-tagGroups:
   - InvokeToolRequest
   - Job
   - JobStatus
+  - JsonType
   - KeyValueMemoryBank
   - KeyValueMemoryBankParams
   - KeywordMemoryBank
@@ -6208,6 +6257,8 @@ x-tagGroups:
   - Model
   - ModelCandidate
   - ModelType
+  - NumberType
+  - ObjectType
   - OptimizerConfig
   - OptimizerType
   - PaginatedRowsResult
@@ -6256,6 +6307,7 @@ x-tagGroups:
   - SpanStatus
   - SpanWithStatus
   - StopReason
+  - StringType
   - StructuredLogEvent
   - SupervisedFineTuneRequest
   - SyntheticDataGenerateRequest
@@ -6286,6 +6338,7 @@ x-tagGroups:
   - TrainingConfig
   - Turn
   - URL
+  - UnionType
   - UnstructuredLogEvent
   - UserMessage
   - VectorMemoryBank
diff --git a/llama_stack/apis/common/type_system.py b/llama_stack/apis/common/type_system.py
index e76cfde13..fa9c5e92e 100644
--- a/llama_stack/apis/common/type_system.py
+++ b/llama_stack/apis/common/type_system.py
@@ -6,54 +6,65 @@
 
 from typing import Literal, Union
 
-from llama_models.schema_utils import register_schema
+from llama_models.schema_utils import json_schema_type, register_schema
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated
 
 
+@json_schema_type
 class StringType(BaseModel):
     type: Literal["string"] = "string"
 
 
+@json_schema_type
 class NumberType(BaseModel):
     type: Literal["number"] = "number"
 
 
+@json_schema_type
 class BooleanType(BaseModel):
     type: Literal["boolean"] = "boolean"
 
 
+@json_schema_type
 class ArrayType(BaseModel):
     type: Literal["array"] = "array"
 
 
+@json_schema_type
 class ObjectType(BaseModel):
     type: Literal["object"] = "object"
 
 
+@json_schema_type
 class JsonType(BaseModel):
     type: Literal["json"] = "json"
 
 
+@json_schema_type
 class UnionType(BaseModel):
     type: Literal["union"] = "union"
 
 
+@json_schema_type
 class ChatCompletionInputType(BaseModel):
     # expects List[Message] for messages
     type: Literal["chat_completion_input"] = "chat_completion_input"
 
 
+@json_schema_type
 class CompletionInputType(BaseModel):
     # expects InterleavedTextMedia for content
     type: Literal["completion_input"] = "completion_input"
 
 
+@json_schema_type
 class AgentTurnInputType(BaseModel):
     # expects List[Message] for messages (may also include attachments?)
     type: Literal["agent_turn_input"] = "agent_turn_input"
 
 
+@json_schema_type
 class DialogType(BaseModel):
     # expects List[Message] for messages
     # this type semantically contains the output label whereas ChatCompletionInputType does not