diff --git a/src/mistralai/extra/tests/test_utils.py b/src/mistralai/extra/tests/test_utils.py index 41fa53e3..bb8ed96a 100644 --- a/src/mistralai/extra/tests/test_utils.py +++ b/src/mistralai/extra/tests/test_utils.py @@ -3,7 +3,7 @@ response_format_from_pydantic_model, rec_strict_json_schema, ) -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, ValidationError, conlist from ...models import ResponseFormat, JSONSchema from ...types.basemodel import Unset @@ -26,6 +26,11 @@ class MathDemonstration(BaseModel): final_answer: str +class ConlistModel(BaseModel): + """Model with constrained list to test numeric values in schema""" + words: conlist(str, min_length=3, max_length=5) # type: ignore + + mathdemo_schema = { "$defs": { "Explanation": { @@ -149,13 +154,26 @@ def test_response_format_from_pydantic_model(self): ) def test_rec_strict_json_schema(self): - invalid_schema = mathdemo_schema | {"wrong_value": 1} + # Test that numeric values (like constraints) are now handled correctly + schema_with_numeric = mathdemo_schema | {"numeric_value": 42, "float_value": 3.14} self.assertEqual( rec_strict_json_schema(mathdemo_schema), mathdemo_strict_schema ) - with self.assertRaises(ValueError): - rec_strict_json_schema(invalid_schema) + # Numeric values should now be accepted as valid schema values + result = rec_strict_json_schema(schema_with_numeric) + self.assertEqual(result["numeric_value"], 42) + self.assertEqual(result["float_value"], 3.14) + + def test_conlist_schema(self): + """Test that Pydantic models with conlist work correctly (issue #279)""" + # This should not raise a ValueError + response_format = response_format_from_pydantic_model(ConlistModel) + + # Verify the response format was created successfully + self.assertIsNotNone(response_format) + self.assertEqual(response_format.json_schema.name, "ConlistModel") + self.assertTrue(response_format.json_schema.strict) if __name__ == "__main__": diff --git a/src/mistralai/extra/utils/_pydantic_helper.py b/src/mistralai/extra/utils/_pydantic_helper.py index f042c394..74aa8a5a 100644 --- a/src/mistralai/extra/utils/_pydantic_helper.py +++ b/src/mistralai/extra/utils/_pydantic_helper.py @@ -6,7 +6,7 @@ def rec_strict_json_schema(schema_node: Any) -> Any: Recursively set the additionalProperties property to False for all objects in the JSON Schema. This makes the JSON Schema strict (i.e. no additional properties are allowed). """ - if isinstance(schema_node, (str, bool)) or schema_node is None: + if isinstance(schema_node, (str, bool, int, float)) or schema_node is None: return schema_node if isinstance(schema_node, dict): if "type" in schema_node and schema_node["type"] == "object":