diff --git a/src/outlines/types/dsl.py b/src/outlines/types/dsl.py index 19504c392..69fc7b97f 100644 --- a/src/outlines/types/dsl.py +++ b/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( [ diff --git a/tests/types/test_dsl.py b/tests/types/test_dsl.py index 2bb813c9f..74153470c 100644 --- a/tests/types/test_dsl.py +++ b/tests/types/test_dsl.py @@ -839,8 +839,9 @@ 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 @@ -848,12 +849,31 @@ def test_dsl_handle_dict(): 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."""