outlines

1
import re
2

3
import pytest
4
from pydantic import BaseModel
5

6
from outlines import types
7
from outlines.fsm.types import python_types_to_regex
8

9

10
@pytest.mark.parametrize(
11
    "custom_type,test_string,should_match",
12
    [
13
        (types.phone_numbers.USPhoneNumber, "12", False),
14
        (types.phone_numbers.USPhoneNumber, "(123) 123-1234", True),
15
        (types.phone_numbers.USPhoneNumber, "123-123-1234", True),
16
        (types.zip_codes.USZipCode, "12", False),
17
        (types.zip_codes.USZipCode, "12345", True),
18
        (types.zip_codes.USZipCode, "12345-1234", True),
19
        (types.ISBN, "ISBN 0-1-2-3-4-5", False),
20
        (types.ISBN, "ISBN 978-0-596-52068-7", True),
21
        # (types.ISBN, "ISBN 978-0-596-52068-1", True), wrong check digit
22
        (types.ISBN, "ISBN-13: 978-0-596-52068-7", True),
23
        (types.ISBN, "978 0 596 52068 7", True),
24
        (types.ISBN, "9780596520687", True),
25
        (types.ISBN, "ISBN-10: 0-596-52068-9", True),
26
        (types.ISBN, "0-596-52068-9", True),
27
        (types.Email, "eitan@gmail.com", True),
28
        (types.Email, "99@yahoo.com", True),
29
        (types.Email, "eitan@.gmail.com", False),
30
        (types.Email, "myemail", False),
31
        (types.Email, "eitan@gmail", False),
32
        (types.Email, "eitan@my.custom.domain", True),
33
    ],
34
)
35
def test_type_regex(custom_type, test_string, should_match):
36
    class Model(BaseModel):
37
        attr: custom_type
38

39
    schema = Model.model_json_schema()
40
    assert schema["properties"]["attr"]["type"] == "string"
41
    regex_str = schema["properties"]["attr"]["pattern"]
42
    does_match = re.match(regex_str, test_string) is not None
43
    assert does_match is should_match
44

45
    regex_str, format_fn = python_types_to_regex(custom_type)
46
    assert isinstance(format_fn(1), str)
47
    does_match = re.match(regex_str, test_string) is not None
48
    assert does_match is should_match
49

50

51
def test_locale_not_implemented():
52
    with pytest.raises(NotImplementedError):
53
        types.locale("fr")
54

55

56
@pytest.mark.parametrize(
57
    "locale_str,base_types,locale_types",
58
    [
59
        (
60
            "us",
61
            ["ZipCode", "PhoneNumber"],
62
            [types.zip_codes.USZipCode, types.phone_numbers.USPhoneNumber],
63
        )
64
    ],
65
)
66
def test_locale(locale_str, base_types, locale_types):
67
    for base_type, locale_type in zip(base_types, locale_types):
68
        type = getattr(types.locale(locale_str), base_type)
69
        assert type == locale_type
70

71

72
@pytest.mark.parametrize(
73
    "custom_type,test_string,should_match",
74
    [
75
        (types.airports.IATA, "CDG", True),
76
        (types.airports.IATA, "XXX", False),
77
        (types.countries.Alpha2, "FR", True),
78
        (types.countries.Alpha2, "XX", False),
79
        (types.countries.Alpha3, "UKR", True),
80
        (types.countries.Alpha3, "XXX", False),
81
        (types.countries.Numeric, "004", True),
82
        (types.countries.Numeric, "900", False),
83
        (types.countries.Name, "Ukraine", True),
84
        (types.countries.Name, "Wonderland", False),
85
        (types.countries.Flag, "🇿🇼", True),
86
        (types.countries.Flag, "🤗", False),
87
    ],
88
)
89
def test_type_enum(custom_type, test_string, should_match):
90
    type_name = custom_type.__name__
91

92
    class Model(BaseModel):
93
        attr: custom_type
94

95
    schema = Model.model_json_schema()
96
    assert isinstance(schema["$defs"][type_name]["enum"], list)
97
    does_match = test_string in schema["$defs"][type_name]["enum"]
98
    assert does_match is should_match
99

100
    regex_str, format_fn = python_types_to_regex(custom_type)
101
    assert isinstance(format_fn(1), str)
102
    does_match = re.match(regex_str, test_string) is not None
103
    assert does_match is should_match
104