fix(types): quote non-string Dict keys so generated output is valid JSON

src/outlines/types/dsl.py

@@ -919,6 +919,17 @@ def _handle_dict(args: tuple, recursion_depth: int) -> Sequence:
         raise TypeError(f"Dict must have exactly two type arguments. Got {args}.")
     # Add dict support with key:value pairs
     key_type = _ensure_json_quoted(python_types_to_terms(args[0], recursion_depth + 1))
+    if isinstance(key_type, Regex) and key_type.pattern != types.string.pattern:
+        # JSON object keys must always be double-quoted strings, even when
+        # the Python key type is not `str` (e.g. `Dict[int, str]`,
+        # `Dict[date, str]`). `_ensure_json_quoted` only quotes `String`/
+        # `Alternatives` terms (from `Literal`/`Enum`); bare `Regex` terms
+        # like `types.integer`, `types.number`, `types.boolean`, `types.date`
+        # etc. are otherwise left unquoted, which produces a regex that
+        # matches invalid JSON (e.g. `{1:"a"}` instead of `{"1":"a"}`).
+        # `types.string` is excluded since its pattern is already
+        # self-quoted (`"[^"]*"`).
+        key_type = Sequence([String('"'), key_type, String('"')])
     value_type = _ensure_json_quoted(python_types_to_terms(args[1], recursion_depth + 1))
     return Sequence(
         [

tests/types/test_dsl.py

@@ -839,21 +839,50 @@ def test_dsl_handle_dict():
         incorrect_dict_type = dict[int, str, int]
         _handle_dict(get_args(incorrect_dict_type), recursion_depth=0)
 
-    # correct type
-    dict_type = dict[int, str]
+    # correct type with a str key: no extra quoting needed, types.string
+    # is already self-quoted
+    dict_type = dict[str, int]
     result = _handle_dict(get_args(dict_type), recursion_depth=0)
     assert isinstance(result, Sequence)
     assert len(result.terms) == 3
     assert result.terms[0] == String("{")
     assert isinstance(result.terms[1], Optional)
     assert isinstance(result.terms[1].term, Sequence)
     assert len(result.terms[1].term.terms) == 4
-    assert result.terms[1].term.terms[0] == types.integer
+    assert result.terms[1].term.terms[0] == types.string
     assert result.terms[1].term.terms[1] == String(":")
-    assert result.terms[1].term.terms[2] == types.string
-    assert result.terms[1].term.terms[3] == KleeneStar(Sequence([String(", "), types.integer, String(":"), types.string]))
+    assert result.terms[1].term.terms[2] == types.integer
+    assert result.terms[1].term.terms[3] == KleeneStar(Sequence([String(", "), types.string, String(":"), types.integer]))
     assert result.terms[2] == String("}")
 
+    # non-str key (e.g. int): JSON object keys must always be double-quoted
+    # strings, so the bare `types.integer` regex must be wrapped in quotes
+    # even though the Python key type is `int`.
+    dict_type = dict[int, str]
+    result = _handle_dict(get_args(dict_type), recursion_depth=0)
+    quoted_int_key = Sequence([String('"'), types.integer, String('"')])
+    assert result.terms[1].term.terms[0] == quoted_int_key
+    assert result.terms[1].term.terms[2] == types.string
+
+
+def test_dsl_handle_dict_non_string_key_produces_valid_json():
+    """Regression test: Dict[int, str] (and other non-str key types) must
+    only match strings that are valid JSON, i.e. the key must be quoted."""
+    import json
+
+    dict_type = dict[int, str]
+    result = _handle_dict(get_args(dict_type), recursion_depth=0)
+    pattern = to_regex(result)
+
+    # unquoted key: not valid JSON, must not match
+    assert _re.fullmatch(pattern, '{1:"a"}') is None
+    with pytest.raises(json.JSONDecodeError):
+        json.loads('{1:"a"}')
+
+    # quoted key: valid JSON, must match
+    assert _re.fullmatch(pattern, '{"1":"a"}') is not None
+    json.loads('{"1":"a"}')  # does not raise
+
 
 def test_ensure_json_quoted_string():
     """String terms are wrapped in double-quote delimiters."""