From cd38b2594e96a752da0213e8e3540bead789c0a3 Mon Sep 17 00:00:00 2001 From: Ihar Hrachyshka Date: Thu, 20 Mar 2025 10:40:26 -0400 Subject: [PATCH] fix: Support types.UnionType in schemas Since Python 3.10, unions can be expressed as `type1 | type2`. Sadly, while this is functionally equivalent to `Union[type1, type2]`, the type of the expression is different (`types.UnionType`, not `typing.Union`). We should handle both in schemas. Signed-off-by: Ihar Hrachyshka --- llama_stack/strong_typing/schema.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/llama_stack/strong_typing/schema.py b/llama_stack/strong_typing/schema.py index dfc51ea78..de69c9b82 100644 --- a/llama_stack/strong_typing/schema.py +++ b/llama_stack/strong_typing/schema.py @@ -17,6 +17,7 @@ import enum import functools import inspect import json +import types import typing import uuid from copy import deepcopy @@ -455,7 +456,7 @@ class JsonSchemaGenerator: "maxItems": len(args), "prefixItems": [self.type_to_schema(member_type) for member_type in args], } - elif origin_type is Union: + elif origin_type in (Union, types.UnionType): discriminator = None if typing.get_origin(data_type) is Annotated: discriminator = typing.get_args(data_type)[1].discriminator