Escape special characters in JSON structure

dottxt-ai · Sep 10, 2023 · 23299e0 · 23299e0
1 parent ce0fad4
commit 23299e0
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 8 deletions.
diff --git a/outlines/text/json_schema.py b/outlines/text/json_schema.py
@@ -1,5 +1,6 @@
 import itertools
 import json
+import re
 from typing import Dict
 
 STRING = r'".*"'
@@ -192,7 +193,7 @@ def match_step_to_regex(step):
     Parameters
     ----------
     step:
-        A string that represents the schema's structure, or a dictionnary
+        A string that represents the schema's structure, or a dictionary
         that represents a field in the schema.
 
     Returns
@@ -203,13 +204,13 @@ def match_step_to_regex(step):
     """
     match step:
         case str() as step:
-            return step
+            return re.escape(step)
 
         case {"enum": choices, "type": "string"}:
-            choices = [f'"{choice}"' for choice in choices]
+            choices = [f'"{re.escape(choice)}"' for choice in choices]
             return f"({'|'.join(choices)})"
         case {"enum": choices}:
-            choices = [str(choice) for choice in choices]
+            choices = [re.escape(str(choice)) for choice in choices]
             return f"({'|'.join(choices)})"
 
         case {"type": "array", "items": items}:

diff --git a/tests/text/test_json_schema.py b/tests/text/test_json_schema.py
@@ -299,6 +299,11 @@ def test_match_number(pattern, does_match):
             '("Marc"|"Jean")',
             [('"Marc"', True), ('"Jean"', True), ('"John"', False)],
         ),
+        (
+            {"title": "Foo", "enum": [".*", r"\s*"], "type": "string"},
+            r'("\.\*"|"\\s\*")',
+            [('".*"', True), (r'"\s*"', True), (r'"\.\*"', False)],
+        ),
         (
             {"title": "Foo", "enum": [0, 1], "type": "integer"},
             "(0|1)",
@@ -310,7 +315,7 @@ def test_match_number(pattern, does_match):
                 "type": "object",
                 "properties": {"count": {"title": "Count", "type": "integer"}},
             },
-            '{\n  "count": ' + INTEGER + "\n}",
+            '\\{\\\n\\ \\ "count":\\ ' + INTEGER + "\\\n\\}",
             [('{\n  "count": 100\n}', True)],
         ),
         (
@@ -339,16 +344,20 @@ def test_match_number(pattern, does_match):
                     }
                 },
             },
-            '{\n  "fuzz": {\n    "spam": ' + INTEGER + "\n  }\n}",
+            '\\{\\\n\\ \\ "fuzz":\\ \\{\\\n\\ \\ \\ \\ "spam":\\ '
+            + INTEGER
+            + "\\\n\\ \\ \\}\\\n\\}",
             [('{\n  "fuzz": {\n    "spam": 100\n  }\n}', True)],
         ),
     ],
 )
 def test_match(step, regex, examples):
-    assert match_step_to_regex(step) == regex
+    test_regex = match_step_to_regex(step)
+
+    assert test_regex == regex
 
     for string, does_match in examples:
-        match = re.fullmatch(regex, string)
+        match = re.fullmatch(test_regex, string)
         if does_match:
             assert match[0] == string
             assert match.span() == (0, len(string))