phoenix-oss/llama-stack
1
0
Fork 0
forked from phoenix-oss/llama-stack-mirror
Code Pull requests Releases Packages 2 Activity Actions

fix: Support types.UnionType in schemas (#1721)

Browse source 
# What does this PR do?

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.

## Test Plan

Switch a schema type from Union to `|` and confirm the generator doesn't
crash with:

```
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/Users/ihrachys/src/llama-stack/docs/openapi_generator/generate.py", line 91, in <module>
    fire.Fire(main)
  File "/Users/ihrachys/.cache/uv/archive-v0/FBgkcwcN-PaJ0NAur__7J/lib/python3.11/site-packages/fire/core.py", line 135, in Fire
    component_trace = _Fire(component, args, parsed_flag_args, context, name)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ihrachys/.cache/uv/archive-v0/FBgkcwcN-PaJ0NAur__7J/lib/python3.11/site-packages/fire/core.py", line 468, in _Fire
    component, remaining_args = _CallAndUpdateTrace(
                                ^^^^^^^^^^^^^^^^^^^^
  File "/Users/ihrachys/.cache/uv/archive-v0/FBgkcwcN-PaJ0NAur__7J/lib/python3.11/site-packages/fire/core.py", line 684, in _CallAndUpdateTrace
    component = fn(*varargs, **kwargs)
                ^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ihrachys/src/llama-stack/docs/openapi_generator/generate.py", line 55, in main
    spec = Specification(
           ^^^^^^^^^^^^^^
  File "/Users/ihrachys/src/llama-stack/docs/openapi_generator/pyopenapi/utility.py", line 30, in __init__
    self.document = generator.generate()
                    ^^^^^^^^^^^^^^^^^^^^
  File "/Users/ihrachys/src/llama-stack/docs/openapi_generator/pyopenapi/generator.py", line 782, in generate
    operation = self._build_operation(op)
                ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ihrachys/src/llama-stack/docs/openapi_generator/pyopenapi/generator.py", line 648, in _build_operation
    "application/json": builder.build_media_type(
                        ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ihrachys/src/llama-stack/docs/openapi_generator/pyopenapi/generator.py", line 221, in build_media_type
    schema = self.schema_builder.classdef_to_ref(item_type)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ihrachys/src/llama-stack/docs/openapi_generator/pyopenapi/generator.py", line 135, in classdef_to_ref
    type_schema = self.classdef_to_schema(typ)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ihrachys/src/llama-stack/docs/openapi_generator/pyopenapi/generator.py", line 116, in classdef_to_schema
    type_schema, type_definitions = self.schema_generator.classdef_to_schema(typ)
                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ihrachys/src/llama-stack/llama_stack/strong_typing/schema.py", line 607, in classdef_to_schema
    types_defined[sub_name] = self._type_to_schema_with_lookup(sub_type)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ihrachys/src/llama-stack/llama_stack/strong_typing/schema.py", line 564, in _type_to_schema_with_lookup
    type_schema = self.type_to_schema(data_type, force_expand=True)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ihrachys/src/llama-stack/llama_stack/strong_typing/schema.py", line 320, in type_to_schema
    return self._type_to_schema(data_type, force_expand, json_schema_extra) | common_info
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ihrachys/src/llama-stack/llama_stack/strong_typing/schema.py", line 487, in _type_to_schema
    property_docstrings = get_class_property_docstrings(typ, self.options.property_description_fun)
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ihrachys/src/llama-stack/llama_stack/strong_typing/schema.py", line 94, in get_class_property_docstrings
    for base in inspect.getmro(data_type):
                ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/w2wykgpkzidnnr6cpw8wf94ghb0p8big-python3-3.11.11/lib/python3.11/inspect.py", line 731, in getmro
    return cls.__mro__
           ^^^^^^^^^^^
AttributeError: 'types.UnionType' object has no attribute '__mro__'. Did you mean: '__or__'?
```

Signed-off-by: Ihar Hrachyshka <ihar.hrachyshka@gmail.com>
This commit is contained in:
Ihar Hrachyshka 2025-03-20 12:54:02 -04:00 • committed by GitHub
parent 5403582582
commit 355134f51d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 2 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

3
llama_stack/strong_typing/schema.py
View file

@ -17,6 +17,7 @@ import enum
import functools import functools
import inspect import inspect
import json import json
import types
import typing import typing
import uuid import uuid
from copy import deepcopy from copy import deepcopy
@ -455,7 +456,7 @@ class JsonSchemaGenerator:
"maxItems": len(args), "maxItems": len(args),
"prefixItems": [self.type_to_schema(member_type) for member_type in 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 discriminator = None
if typing.get_origin(data_type) is Annotated: if typing.get_origin(data_type) is Annotated:
discriminator = typing.get_args(data_type)[1].discriminator discriminator = typing.get_args(data_type)[1].discriminator