fix(responses): type aliasing not supported for pydantic code generation and discrimintated unions

This commit is contained in:
Emilio Garcia 2025-08-20 18:05:23 -04:00
parent 8fb17ba18e
commit 80b82c070c
5 changed files with 287 additions and 73 deletions

View file

@ -4,7 +4,7 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
from typing import Annotated, Any, Literal, Union
from typing import Annotated, Any, Literal
from pydantic import BaseModel, Field
from typing_extensions import TypedDict
@ -14,21 +14,20 @@ from llama_stack.apis.tools.openai_tool_choice import (
ToolChoiceCustom,
ToolChoiceFunction,
ToolChoiceMcp,
ToolChoiceOptions,
ToolChoiceTypes,
)
from llama_stack.apis.vector_io import SearchRankingOptions as FileSearchRankingOptions
from llama_stack.schema_utils import json_schema_type, register_schema
type OpenAIResponsesToolChoice = Annotated[
Union[
ToolChoiceTypes,
ToolChoiceAllowed,
ToolChoiceFunction,
ToolChoiceMcp,
ToolChoiceCustom
],
Field(discriminator="type"),
]
OpenAIResponsesToolChoice = (
ToolChoiceOptions
| ToolChoiceTypes # Multiple type values - can't use a discriminator here
| Annotated[
ToolChoiceAllowed | ToolChoiceFunction | ToolChoiceMcp | ToolChoiceCustom,
Field(discriminator="type"),
]
)
register_schema(OpenAIResponsesToolChoice, name="OpenAIResponsesToolChoice")

View file

@ -10,7 +10,7 @@ from pydantic import BaseModel
from llama_stack.schema_utils import json_schema_type, register_schema
type ToolChoiceOptions = Literal["none", "auto", "required"]
ToolChoiceOptions = Literal["none", "auto", "required"]
register_schema(ToolChoiceOptions, name="ToolChoiceOptions")
@ -24,7 +24,7 @@ class ToolChoiceTypes(BaseModel):
"image_generation",
"code_interpreter",
]
"""The type of hosted tool the model should to use.
"""The type of hosted tool the model should use.
Allowed values are:
@ -61,7 +61,7 @@ class ToolChoiceAllowed(BaseModel):
```
"""
type: Literal["allowed_tools"]
type: Literal["allowed_tools"] = "allowed_tools"
"""Allowed tool configuration type. Always `allowed_tools`."""
@ -70,7 +70,7 @@ class ToolChoiceFunction(BaseModel):
name: str
"""The name of the function to call."""
type: Literal["function"]
type: Literal["function"] = "function"
"""For function calling, the type is always `function`."""
@ -79,7 +79,7 @@ class ToolChoiceMcp(BaseModel):
server_label: str
"""The label of the MCP server to use."""
type: Literal["mcp"]
type: Literal["mcp"] = "mcp"
"""For MCP tools, the type is always `mcp`."""
name: str | None = None
@ -91,5 +91,5 @@ class ToolChoiceCustom(BaseModel):
name: str
"""The name of the custom tool to call."""
type: Literal["custom"]
type: Literal["custom"] = "custom"
"""For custom tool calling, the type is always `custom`."""