Ensure regex matches valid JSON for "const" and "enum" with booleans,…

… nulls, and strings Use JSON's serialization, rather than Python's, in this case. Fixes #971. Note that this still does not correctly handle arrays and objects, which are allowed by the JSON schema spec; however, those would be more complex to handle correctly.
dottxt-ai · Jun 18, 2024 · 7d8269f · 7d8269f
1 parent 63ac124
commit 7d8269f
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 13 deletions.
diff --git a/outlines/fsm/json_schema.py b/outlines/fsm/json_schema.py
@@ -269,19 +269,18 @@ def to_regex(
     elif "enum" in instance:
         choices = []
         for choice in instance["enum"]:
-            if type(choice) in [int, float, bool, None]:
-                choices.append(re.escape(str(choice)))
-            elif type(choice) == str:
-                choices.append(f'"{re.escape(choice)}"')
-
+            if type(choice) in [int, float, bool, type(None), str]:
+                choices.append(re.escape(json.dumps(choice)))
+            else:
+                raise TypeError(f"Unsupported data type in enum: {type(choice)}")
         return f"({'|'.join(choices)})"
 
     elif "const" in instance:
         const = instance["const"]
-        if type(const) in [int, float, bool, None]:
-            const = re.escape(str(const))
-        elif type(const) == str:
-            const = f'"{re.escape(const)}"'
+        if type(const) in [int, float, bool, type(None), str]:
+            const = re.escape(json.dumps(const))
+        else:
+            raise TypeError(f"Unsupported data type in const: {type(const)}")
         return const
 
     elif "$ref" in instance:

diff --git a/tests/fsm/test_json_schema.py b/tests/fsm/test_json_schema.py
@@ -177,36 +177,75 @@ def test_match_number(pattern, does_match):
             '"Marc"',
             [('"Marc"', True), ('"Jean"', False), ('"John"', False)],
         ),
-        # Make sure strings are escaped
+        # Make sure strings are escaped with regex escaping
         (
             {"title": "Foo", "const": ".*", "type": "string"},
             r'"\.\*"',
             [('".*"', True), (r'"\s*"', False), (r'"\.\*"', False)],
         ),
+        # Make sure strings are escaped with JSON escaping
+        (
+            {"title": "Foo", "const": '"', "type": "string"},
+            r'"\\""',
+            [('"\\""', True), ('"""', False)],
+        ),
         # Const integer
         (
             {"title": "Foo", "const": 0, "type": "integer"},
             "0",
             [("0", True), ("1", False), ("a", False)],
         ),
+        # Const float
+        (
+            {"title": "Foo", "const": 0.2, "type": "float"},
+            r"0\.2",
+            [("0.2", True), ("032", False)],
+        ),
+        # Const boolean
+        (
+            {"title": "Foo", "const": True, "type": "boolean"},
+            "true",
+            [("true", True), ("True", False)],
+        ),
+        # Const null
+        (
+            {"title": "Foo", "const": None, "type": "null"},
+            "null",
+            [("null", True), ("None", False), ("", False)],
+        ),
         # Enum string
         (
             {"title": "Foo", "enum": ["Marc", "Jean"], "type": "string"},
             '("Marc"|"Jean")',
             [('"Marc"', True), ('"Jean"', True), ('"John"', False)],
         ),
-        # Make sure strings are escaped
+        # Make sure strings are escaped with regex and JSON escaping
         (
             {"title": "Foo", "enum": [".*", r"\s*"], "type": "string"},
-            r'("\.\*"|"\\s\*")',
-            [('".*"', True), (r'"\s*"', True), (r'"\.\*"', False)],
+            r'("\.\*"|"\\\\s\*")',
+            [('".*"', True), (r'"\\s*"', True), (r'"\.\*"', False)],
         ),
         # Enum integer
         (
             {"title": "Foo", "enum": [0, 1], "type": "integer"},
             "(0|1)",
             [("0", True), ("1", True), ("a", False)],
         ),
+        # Enum mix of types
+        (
+            {"title": "Foo", "enum": [6, 5.3, "potato", True, None]},
+            r'(6|5\.3|"potato"|true|null)',
+            [
+                ("6", True),
+                ("5.3", True),
+                ('"potato"', True),
+                ("true", True),
+                ("null", True),
+                ("523", False),
+                ("True", False),
+                ("None", False),
+            ],
+        ),
         # integer
         (
             {