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

src/outlines/types/dsl.py

@@ -838,19 +838,24 @@ def _handle_literal(args: tuple) -> Alternatives:
     return Alternatives([python_types_to_terms(arg) for arg in args])
 
 
-def _ensure_json_quoted(term: Term) -> Term:
-    """Wrap bare ``String`` terms in double quotes for JSON container contexts.
-
-    When string literal values (from ``Literal`` or ``Enum``) appear inside
-    container types (``List``, ``Tuple``, ``Dict``), they must be JSON-quoted
-    so the generated regex matches valid JSON.  ``Regex``-based terms (e.g.
-    ``types.string``) already include their own quotes and are left untouched.
+def _ensure_json_quoted(term: Term, quote_regex: bool = False) -> Term:
+    """Wrap ``String`` terms in double quotes for JSON container contexts.
+
+    String literals (from ``Literal``/``Enum``) inside containers must be
+    JSON-quoted so the generated regex matches valid JSON. With
+    ``quote_regex``, bare ``Regex`` terms are quoted too (needed for ``Dict``
+    keys, which JSON always requires to be strings); ``types.string`` is never
+    re-quoted since its pattern already includes the quotes. Both cases recurse
+    into ``Alternatives``, so ``Regex`` members nested there (e.g.
+    ``Dict[Literal[1, 2], str]``) are quoted as well.
     """
     if isinstance(term, String):
         return String(f'"{term.value}"')
     if isinstance(term, Alternatives):
-        quoted = [_ensure_json_quoted(t) for t in term.terms]
+        quoted = [_ensure_json_quoted(t, quote_regex) for t in term.terms]
         return Alternatives(quoted)
+    if quote_regex and isinstance(term, Regex) and term.pattern != types.string.pattern:
+        return Sequence([String('"'), term, String('"')])
     return term
 
 
@@ -917,8 +922,11 @@ def _handle_tuple(args: tuple, recursion_depth: int) -> Union[Sequence, String]:
 def _handle_dict(args: tuple, recursion_depth: int) -> Sequence:
     if args is None or len(args) != 2:
         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))
+    # JSON object keys must always be quoted strings, so quote the key term
+    # even when the Python key type isn't `str` (e.g. `Dict[int, str]`).
+    key_type = _ensure_json_quoted(
+        python_types_to_terms(args[0], recursion_depth + 1), quote_regex=True
+    )
     value_type = _ensure_json_quoted(python_types_to_terms(args[1], recursion_depth + 1))
     return Sequence(
         [

tests/types/test_dsl.py

@@ -839,21 +839,41 @@ 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
+
+    # non-str key nested in an Alternatives (Literal ints go through
+    # _handle_literal): each member must be quoted, not just a top-level Regex
+    dict_type = dict[Literal[1, 2, 3], str]
+    result = _handle_dict(get_args(dict_type), recursion_depth=0)
+    key_term = result.terms[1].term.terms[0]
+    assert isinstance(key_term, Alternatives)
+    assert key_term.terms == [
+        Sequence([String('"'), Regex(str(i)), String('"')]) for i in (1, 2, 3)
+    ]
+
 
 def test_ensure_json_quoted_string():
     """String terms are wrapped in double-quote delimiters."""