Support PEP 604 unions (int | str) as output types

src/outlines/types/utils.py

@@ -4,6 +4,7 @@
 import datetime
 import inspect
 import sys
+import types
 import warnings
 from enum import Enum, EnumMeta
 from typing import (
@@ -113,7 +114,7 @@ def is_typing_tuple(value: Any) -> bool:
 
 
 def is_union(value: Any) -> bool:
-    return get_origin(value) is Union
+    return get_origin(value) is Union or isinstance(value, types.UnionType)
 
 
 def is_literal(value: Any) -> bool:

tests/types/test_dsl.py

@@ -638,6 +638,10 @@ class SomeEnum(Enum):
     # convert to terms are tested in distinct tests below
     assert python_types_to_terms(Literal["a", "b"]) == _handle_literal(("a", "b"))
     assert python_types_to_terms(Union[int, str]) == _handle_union((int, str), recursion_depth=0)
+    # PEP 604 unions must dispatch the same way as typing.Union/Optional
+    assert python_types_to_terms(int | str) == python_types_to_terms(Union[int, str])
+    assert python_types_to_terms(str | None) == python_types_to_terms(PyOptional[str])
+    assert python_types_to_terms(list[int | str]) == python_types_to_terms(list[Union[int, str]])
     assert python_types_to_terms(list[int]) == _handle_list((int,), recursion_depth=0)
     assert python_types_to_terms(tuple[int, str]) == _handle_tuple((int, str), recursion_depth=0)
     assert python_types_to_terms(dict[int, str]) == _handle_dict((int, str), recursion_depth=0)

tests/types/test_types_utils.py

@@ -274,6 +274,8 @@ def test_is_typing_tuple():
 def test_is_union():
     assert is_union(Union[int, str])
     assert is_union(Optional[int])
+    assert is_union(int | str)
+    assert is_union(str | None)
     assert not is_union(list)
     assert not is_union(["a", "b"])
     assert not is_union(Literal[int, str])