diff --git a/outlines/generate/json.py b/outlines/generate/json.py index ab05df1f4..f0968df0b 100644 --- a/outlines/generate/json.py +++ b/outlines/generate/json.py @@ -50,7 +50,7 @@ def json( schema = pyjson.dumps(schema_object.model_json_schema()) regex_str = build_regex_from_schema(schema, whitespace_pattern) generator = regex(model, regex_str, sampler) - generator.format_sequence = lambda x: schema_object.parse_raw(x) + generator.format_sequence = lambda x: schema_object.model_validate_json(x) elif isinstance(schema_object, type(Enum)): schema = pyjson.dumps(get_schema_from_enum(schema_object)) regex_str = build_regex_from_schema(schema, whitespace_pattern) @@ -95,7 +95,7 @@ def json_openai( schema = schema_object.model_json_schema() schema["additionalProperties"] = False schema = pyjson.dumps(schema) - format_sequence = lambda x: schema_object.parse_raw(x) + format_sequence = lambda x: schema_object.model_validate_json(x) elif isinstance(schema_object, str): schema = schema_object format_sequence = lambda x: pyjson.loads(x) diff --git a/tests/models/test_openai.py b/tests/models/test_openai.py index f4b97f36a..3907a903a 100644 --- a/tests/models/test_openai.py +++ b/tests/models/test_openai.py @@ -99,7 +99,7 @@ class Person(BaseModel): # assert success for valid response with patched_openai(completion=completion) as model: generator = generate.json(model, Person) - assert generator("fastest person") == Person.parse_raw(completion) + assert generator("fastest person") == Person.model_validate_json(completion) # assert fail for non-json response with patched_openai(completion="usain bolt") as model: