Update OpenAPI generator to output discriminator (#848)

oneOf should have discriminators so Stainless can generate better code

## Test Plan

Going to generate the SDK now and check.

docs/notebooks/Llama_Stack_Building_AI_Applications.ipynb

@@ -2088,7 +2088,7 @@
         "from llama_stack_client.lib.agents.event_logger import EventLogger\n",
         "from llama_stack_client.types.agent_create_params import AgentConfig\n",
         "from termcolor import cprint\n",
-        "from llama_stack_client.types.tool_runtime import DocumentParam as Document\n",
+        "from llama_stack_client.types import Document\n",
         "\n",
         "urls = [\"chat.rst\", \"llama3.rst\", \"datasets.rst\", \"lora_finetune.rst\"]\n",
         "documents = [\n",

docs/openapi_generator/strong_typing/classdef.py

@@ -125,6 +125,7 @@ class JsonSchemaAnyOf(JsonSchemaNode):
 @dataclass
 class JsonSchemaOneOf(JsonSchemaNode):
     oneOf: List["JsonSchemaAny"]
+    discriminator: Optional[str]
 
 
 JsonSchemaAny = Union[

docs/openapi_generator/strong_typing/schema.py

@@ -36,6 +36,7 @@ from typing import (
 )
 
 import jsonschema
+from typing_extensions import Annotated
 
 from . import docstring
 from .auxiliary import (
@@ -329,7 +330,6 @@ class JsonSchemaGenerator:
         if metadata is not None:
             # type is Annotated[T, ...]
             typ = typing.get_args(data_type)[0]
-
             schema = self._simple_type_to_schema(typ)
             if schema is not None:
                 # recognize well-known auxiliary types
@@ -446,12 +446,20 @@ class JsonSchemaGenerator:
                 ],
             }
         elif origin_type is Union:
-            return {
+            discriminator = None
+            if typing.get_origin(data_type) is Annotated:
+                discriminator = typing.get_args(data_type)[1].discriminator
+            ret = {
                 "oneOf": [
                     self.type_to_schema(union_type)
                     for union_type in typing.get_args(typ)
                 ]
             }
+            if discriminator:
+                ret["discriminator"] = {
+                    "propertyName": discriminator,
+                }
+            return ret
         elif origin_type is Literal:
             (literal_value,) = typing.get_args(typ)  # unpack value of literal type
             schema = self.type_to_schema(type(literal_value))

docs/resources/llama-stack-spec.html

@@ -3810,7 +3810,10 @@
                     {
                         "$ref": "#/components/schemas/TextContentItem"
                     }
-                ]
+                ],
+                "discriminator": {
+                    "propertyName": "type"
+                }
             },
             "Message": {
                 "oneOf": [
@@ -3826,7 +3829,10 @@
                     {
                         "$ref": "#/components/schemas/CompletionMessage"
                     }
-                ]
+                ],
+                "discriminator": {
+                    "propertyName": "role"
+                }
             },
             "SamplingParams": {
                 "type": "object",
@@ -3842,7 +3848,10 @@
                             {
                                 "$ref": "#/components/schemas/TopKSamplingStrategy"
                             }
-                        ]
+                        ],
+                        "discriminator": {
+                            "propertyName": "type"
+                        }
                     },
                     "max_tokens": {
                         "type": "integer",
@@ -4386,7 +4395,10 @@
                             "bnf"
                         ]
                     }
-                ]
+                ],
+                "discriminator": {
+                    "propertyName": "type"
+                }
             },
             "ChatCompletionRequest": {
                 "type": "object",
@@ -4515,7 +4527,10 @@
                     {
                         "$ref": "#/components/schemas/ToolCallDelta"
                     }
-                ]
+                ],
+                "discriminator": {
+                    "propertyName": "type"
+                }
             },
             "ImageDelta": {
                 "type": "object",
@@ -5019,7 +5034,10 @@
                             {
                                 "$ref": "#/components/schemas/AgentTurnResponseTurnCompletePayload"
                             }
-                        ]
+                        ],
+                        "discriminator": {
+                            "propertyName": "event_type"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -5062,7 +5080,10 @@
                             {
                                 "$ref": "#/components/schemas/MemoryRetrievalStep"
                             }
-                        ]
+                        ],
+                        "discriminator": {
+                            "propertyName": "step_type"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -5462,7 +5483,10 @@
                                 {
                                     "$ref": "#/components/schemas/MemoryRetrievalStep"
                                 }
-                            ]
+                            ],
+                            "discriminator": {
+                                "propertyName": "step_type"
+                            }
                         }
                     },
                     "output_message": {
@@ -5612,7 +5636,10 @@
                             {
                                 "$ref": "#/components/schemas/AgentCandidate"
                             }
-                        ]
+                        ],
+                        "discriminator": {
+                            "propertyName": "type"
+                        }
                     },
                     "scoring_params": {
                         "type": "object",
@@ -5627,7 +5654,10 @@
                                 {
                                     "$ref": "#/components/schemas/BasicScoringFnParams"
                                 }
-                            ]
+                            ],
+                            "discriminator": {
+                                "propertyName": "type"
+                            }
                         }
                     },
                     "num_examples": {
@@ -5677,7 +5707,10 @@
                             {
                                 "$ref": "#/components/schemas/AgentCandidate"
                             }
-                        ]
+                        ],
+                        "discriminator": {
+                            "propertyName": "type"
+                        }
                     },
                     "num_examples": {
                         "type": "integer"
@@ -5818,7 +5851,10 @@
                             {
                                 "$ref": "#/components/schemas/AppEvalTaskConfig"
                             }
-                        ]
+                        ],
+                        "discriminator": {
+                            "propertyName": "type"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -5981,7 +6017,10 @@
                             {
                                 "$ref": "#/components/schemas/MemoryRetrievalStep"
                             }
-                        ]
+                        ],
+                        "discriminator": {
+                            "propertyName": "step_type"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -6196,7 +6235,10 @@
                     {
                         "$ref": "#/components/schemas/AgentTurnInputType"
                     }
-                ]
+                ],
+                "discriminator": {
+                    "propertyName": "type"
+                }
             },
             "StringType": {
                 "type": "object",
@@ -6456,7 +6498,10 @@
                             {
                                 "$ref": "#/components/schemas/BasicScoringFnParams"
                             }
-                        ]
+                        ],
+                        "discriminator": {
+                            "propertyName": "type"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -7542,7 +7587,10 @@
                             {
                                 "$ref": "#/components/schemas/SpanEndPayload"
                             }
-                        ]
+                        ],
+                        "discriminator": {
+                            "propertyName": "type"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -7628,7 +7676,10 @@
                             {
                                 "$ref": "#/components/schemas/StructuredLogEvent"
                             }
-                        ]
+                        ],
+                        "discriminator": {
+                            "propertyName": "type"
+                        }
                     },
                     "ttl_seconds": {
                         "type": "integer"
@@ -7958,7 +8009,10 @@
                     {
                         "$ref": "#/components/schemas/LLMRAGQueryGeneratorConfig"
                     }
-                ]
+                ],
+                "discriminator": {
+                    "propertyName": "type"
+                }
             },
             "QueryRequest": {
                 "type": "object",
@@ -8350,7 +8404,10 @@
                             {
                                 "$ref": "#/components/schemas/BasicScoringFnParams"
                             }
-                        ]
+                        ],
+                        "discriminator": {
+                            "propertyName": "type"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -8483,7 +8540,10 @@
                             {
                                 "$ref": "#/components/schemas/AppEvalTaskConfig"
                             }
-                        ]
+                        ],
+                        "discriminator": {
+                            "propertyName": "type"
+                        }
                     }
                 },
                 "additionalProperties": false,
@@ -8632,7 +8692,10 @@
                                         {
                                             "$ref": "#/components/schemas/BasicScoringFnParams"
                                         }
-                                    ]
+                                    ],
+                                    "discriminator": {
+                                        "propertyName": "type"
+                                    }
                                 },
                                 {
                                     "type": "null"
@@ -8683,7 +8746,10 @@
                                         {
                                             "$ref": "#/components/schemas/BasicScoringFnParams"
                                         }
-                                    ]
+                                    ],
+                                    "discriminator": {
+                                        "propertyName": "type"
+                                    }
                                 },
                                 {
                                     "type": "null"
@@ -8860,7 +8926,10 @@
                             {
                                 "$ref": "#/components/schemas/QATFinetuningConfig"
                             }
-                        ]
+                        ],
+                        "discriminator": {
+                            "propertyName": "type"
+                        }
                     }
                 },
                 "additionalProperties": false,

docs/resources/llama-stack-spec.yaml

@@ -76,6 +76,8 @@ components:
       additionalProperties: false
       properties:
         step:
+          discriminator:
+            propertyName: step_type
           oneOf:
           - $ref: '#/components/schemas/InferenceStep'
           - $ref: '#/components/schemas/ToolExecutionStep'
@@ -119,6 +121,8 @@ components:
       additionalProperties: false
       properties:
         payload:
+          discriminator:
+            propertyName: event_type
           oneOf:
           - $ref: '#/components/schemas/AgentTurnResponseStepStartPayload'
           - $ref: '#/components/schemas/AgentTurnResponseStepProgressPayload'
@@ -137,6 +141,8 @@ components:
           default: step_complete
           type: string
         step_details:
+          discriminator:
+            propertyName: step_type
           oneOf:
           - $ref: '#/components/schemas/InferenceStep'
           - $ref: '#/components/schemas/ToolExecutionStep'
@@ -258,6 +264,8 @@ components:
       additionalProperties: false
       properties:
         eval_candidate:
+          discriminator:
+            propertyName: type
           oneOf:
           - $ref: '#/components/schemas/ModelCandidate'
           - $ref: '#/components/schemas/AgentCandidate'
@@ -265,6 +273,8 @@ components:
           type: integer
         scoring_params:
           additionalProperties:
+            discriminator:
+              propertyName: type
             oneOf:
             - $ref: '#/components/schemas/LLMAsJudgeScoringFnParams'
             - $ref: '#/components/schemas/RegexParserScoringFnParams'
@@ -402,6 +412,8 @@ components:
       additionalProperties: false
       properties:
         eval_candidate:
+          discriminator:
+            propertyName: type
           oneOf:
           - $ref: '#/components/schemas/ModelCandidate'
           - $ref: '#/components/schemas/AgentCandidate'
@@ -619,6 +631,8 @@ components:
       title: streamed completion response.
       type: object
     ContentDelta:
+      discriminator:
+        propertyName: type
       oneOf:
       - $ref: '#/components/schemas/TextDelta'
       - $ref: '#/components/schemas/ImageDelta'
@@ -897,6 +911,8 @@ components:
             type: string
           type: array
         task_config:
+          discriminator:
+            propertyName: type
           oneOf:
           - $ref: '#/components/schemas/BenchmarkEvalTaskConfig'
           - $ref: '#/components/schemas/AppEvalTaskConfig'
@@ -1038,6 +1054,8 @@ components:
           $ref: '#/components/schemas/InterleavedContentItem'
         type: array
     InterleavedContentItem:
+      discriminator:
+        propertyName: type
       oneOf:
       - $ref: '#/components/schemas/ImageContentItem'
       - $ref: '#/components/schemas/TextContentItem'
@@ -1244,6 +1262,8 @@ components:
       additionalProperties: false
       properties:
         event:
+          discriminator:
+            propertyName: type
           oneOf:
           - $ref: '#/components/schemas/UnstructuredLogEvent'
           - $ref: '#/components/schemas/MetricEvent'
@@ -1325,6 +1345,8 @@ components:
       - inserted_context
       type: object
     Message:
+      discriminator:
+        propertyName: role
       oneOf:
       - $ref: '#/components/schemas/UserMessage'
       - $ref: '#/components/schemas/SystemMessage'
@@ -1495,6 +1517,8 @@ components:
       - total_count
       type: object
     ParamType:
+      discriminator:
+        propertyName: type
       oneOf:
       - $ref: '#/components/schemas/StringType'
       - $ref: '#/components/schemas/NumberType'
@@ -1805,6 +1829,8 @@ components:
       - max_chunks
       type: object
     RAGQueryGeneratorConfig:
+      discriminator:
+        propertyName: type
       oneOf:
       - $ref: '#/components/schemas/DefaultRAGQueryGeneratorConfig'
       - $ref: '#/components/schemas/LLMRAGQueryGeneratorConfig'
@@ -1922,6 +1948,8 @@ components:
         description:
           type: string
         params:
+          discriminator:
+            propertyName: type
           oneOf:
           - $ref: '#/components/schemas/LLMAsJudgeScoringFnParams'
           - $ref: '#/components/schemas/RegexParserScoringFnParams'
@@ -2002,6 +2030,8 @@ components:
       - embedding_model
       type: object
     ResponseFormat:
+      discriminator:
+        propertyName: type
       oneOf:
       - additionalProperties: false
         properties:
@@ -2063,6 +2093,8 @@ components:
       additionalProperties: false
       properties:
         task_config:
+          discriminator:
+            propertyName: type
           oneOf:
           - $ref: '#/components/schemas/BenchmarkEvalTaskConfig'
           - $ref: '#/components/schemas/AppEvalTaskConfig'
@@ -2130,6 +2162,8 @@ components:
           default: 1.0
           type: number
         strategy:
+          discriminator:
+            propertyName: type
           oneOf:
           - $ref: '#/components/schemas/GreedySamplingStrategy'
           - $ref: '#/components/schemas/TopPSamplingStrategy'
@@ -2167,7 +2201,9 @@ components:
         scoring_functions:
           additionalProperties:
             oneOf:
-            - oneOf:
+            - discriminator:
+                propertyName: type
+              oneOf:
               - $ref: '#/components/schemas/LLMAsJudgeScoringFnParams'
               - $ref: '#/components/schemas/RegexParserScoringFnParams'
               - $ref: '#/components/schemas/BasicScoringFnParams'
@@ -2208,7 +2244,9 @@ components:
         scoring_functions:
           additionalProperties:
             oneOf:
-            - oneOf:
+            - discriminator:
+                propertyName: type
+              oneOf:
               - $ref: '#/components/schemas/LLMAsJudgeScoringFnParams'
               - $ref: '#/components/schemas/RegexParserScoringFnParams'
               - $ref: '#/components/schemas/BasicScoringFnParams'
@@ -2246,6 +2284,8 @@ components:
             - type: object
           type: object
         params:
+          discriminator:
+            propertyName: type
           oneOf:
           - $ref: '#/components/schemas/LLMAsJudgeScoringFnParams'
           - $ref: '#/components/schemas/RegexParserScoringFnParams'
@@ -2503,6 +2543,8 @@ components:
             - type: object
           type: object
         payload:
+          discriminator:
+            propertyName: type
           oneOf:
           - $ref: '#/components/schemas/SpanStartPayload'
           - $ref: '#/components/schemas/SpanEndPayload'
@@ -2528,6 +2570,8 @@ components:
       additionalProperties: false
       properties:
         algorithm_config:
+          discriminator:
+            propertyName: type
           oneOf:
           - $ref: '#/components/schemas/LoraFinetuningConfig'
           - $ref: '#/components/schemas/QATFinetuningConfig'
@@ -3115,6 +3159,8 @@ components:
           type: string
         steps:
           items:
+            discriminator:
+              propertyName: step_type
             oneOf:
             - $ref: '#/components/schemas/InferenceStep'
             - $ref: '#/components/schemas/ToolExecutionStep'

llama_stack/distribution/store/registry.py

@@ -35,7 +35,7 @@ class DistributionRegistry(Protocol):
 
 
 REGISTER_PREFIX = "distributions:registry"
-KEY_VERSION = "v6"
+KEY_VERSION = "v7"
 KEY_FORMAT = f"{REGISTER_PREFIX}:{KEY_VERSION}::" + "{type}:{identifier}"
 
 

tests/client-sdk/tool_runtime/test_rag_tool.py

@@ -8,7 +8,7 @@ import random
 
 import pytest
 
-from llama_stack_client.types.tool_runtime import DocumentParam
+from llama_stack_client.types import Document
 
 
 @pytest.fixture(scope="function")
@@ -38,22 +38,22 @@ def single_entry_vector_db_registry(llama_stack_client, empty_vector_db_registry
 @pytest.fixture(scope="session")
 def sample_documents():
     return [
-        DocumentParam(
+        Document(
             document_id="test-doc-1",
             content="Python is a high-level programming language.",
             metadata={"category": "programming", "difficulty": "beginner"},
         ),
-        DocumentParam(
+        Document(
             document_id="test-doc-2",
             content="Machine learning is a subset of artificial intelligence.",
             metadata={"category": "AI", "difficulty": "advanced"},
         ),
-        DocumentParam(
+        Document(
             document_id="test-doc-3",
             content="Data structures are fundamental to computer science.",
             metadata={"category": "computer science", "difficulty": "intermediate"},
         ),
-        DocumentParam(
+        Document(
             document_id="test-doc-4",
             content="Neural networks are inspired by biological neural networks.",
             metadata={"category": "AI", "difficulty": "advanced"},
@@ -148,7 +148,7 @@ def test_vector_db_insert_from_url_and_query(
         "llama3.rst",
     ]
     documents = [
-        DocumentParam(
+        Document(
             document_id=f"num-{i}",
             content=f"https://raw.githubusercontent.com/pytorch/torchtune/main/docs/source/tutorials/{url}",
             mime_type="text/plain",