From e0859039b61e87d03068db18330b29d20e54aea5 Mon Sep 17 00:00:00 2001 From: Varun Chawla Date: Sat, 7 Feb 2026 21:50:28 -0800 Subject: [PATCH] Fix JSONSchema serialization inconsistency with alias field Fixes #310 The JSONSchema.model_serializer was not handling the case where model_dump() is called without by_alias=True. When by_alias=False, the serialized dict contains field names (e.g., "schema_definition") instead of aliases (e.g., "schema"), causing the serializer to incorrectly set the schema field to None. This commit updates the serializer to check both the alias and field name when retrieving values, ensuring consistent behavior regardless of the by_alias parameter. Changes: - Updated JSONSchema.serialize_model to try both alias and field name - Added test cases to verify serialization works with and without by_alias - Ensured both serialization methods produce identical results --- src/mistralai/models/jsonschema.py | 4 +- tests/test_jsonschema_serialization.py | 56 ++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) create mode 100644 tests/test_jsonschema_serialization.py diff --git a/src/mistralai/models/jsonschema.py b/src/mistralai/models/jsonschema.py index e2b6a45e..23e163db 100644 --- a/src/mistralai/models/jsonschema.py +++ b/src/mistralai/models/jsonschema.py @@ -36,8 +36,10 @@ def serialize_model(self, handler): for n, f in type(self).model_fields.items(): k = f.alias or n - val = serialized.get(k) + # Try alias first, then field name as fallback to handle both by_alias=True and by_alias=False + val = serialized.get(k) or serialized.get(n) serialized.pop(k, None) + serialized.pop(n, None) optional_nullable = k in optional_fields and k in nullable_fields is_set = ( diff --git a/tests/test_jsonschema_serialization.py b/tests/test_jsonschema_serialization.py new file mode 100644 index 00000000..47417fc7 --- /dev/null +++ b/tests/test_jsonschema_serialization.py @@ -0,0 +1,56 @@ +"""Test for JSONSchema serialization with and without by_alias parameter.""" +import json +from mistralai.extra.utils.response_format import response_format_from_pydantic_model +from pydantic import BaseModel, Field + + +class TestModel(BaseModel): + name: str = Field(description="A name field") + value: int = Field(description="A value field") + + +def test_jsonschema_serialization_without_by_alias(): + """Test that model_dump() works without by_alias=True.""" + response_format = response_format_from_pydantic_model(TestModel) + + # This should work without by_alias=True + serialized = response_format.model_dump(mode='json') + + # Verify that schema is not None + assert serialized['json_schema']['schema'] is not None, "Schema should not be None" + + # Verify that schema contains expected fields + schema = serialized['json_schema']['schema'] + assert 'type' in schema + assert 'properties' in schema + assert 'name' in schema['properties'] + assert 'value' in schema['properties'] + + +def test_jsonschema_serialization_with_by_alias(): + """Test that model_dump(by_alias=True) continues to work.""" + response_format = response_format_from_pydantic_model(TestModel) + + # This should work with by_alias=True + serialized = response_format.model_dump(mode='json', by_alias=True) + + # Verify that schema is not None + assert serialized['json_schema']['schema'] is not None, "Schema should not be None" + + # Verify that schema contains expected fields + schema = serialized['json_schema']['schema'] + assert 'type' in schema + assert 'properties' in schema + assert 'name' in schema['properties'] + assert 'value' in schema['properties'] + + +def test_jsonschema_serialization_consistency(): + """Test that both serialization methods produce the same result.""" + response_format = response_format_from_pydantic_model(TestModel) + + serialized_without_alias = response_format.model_dump(mode='json') + serialized_with_alias = response_format.model_dump(mode='json', by_alias=True) + + # Both should produce the same schema + assert serialized_without_alias['json_schema']['schema'] == serialized_with_alias['json_schema']['schema']