From c28e172937ea7dad195d6d90c17940640e2413bc Mon Sep 17 00:00:00 2001 From: Priyanshu Sharma Date: Fri, 30 May 2025 22:48:12 +0530 Subject: [PATCH 1/2] feat: Implement native JSON mode with schema for Gemini Enables direct use of Pydantic models for response_schema in Gemini API calls when json_mode=True. Includes updates to warnings for strict mode and relevant unit tests. --- mirascope/core/base/tool.py | 6 +- mirascope/core/google/_call.py | 2 +- mirascope/core/google/_utils/_setup_call.py | 15 +- tests/core/google/_utils/test_setup_call.py | 177 ++++++++++++++++++++ 4 files changed, 193 insertions(+), 7 deletions(-) diff --git a/mirascope/core/base/tool.py b/mirascope/core/base/tool.py index 17e023043..e39118849 100644 --- a/mirascope/core/base/tool.py +++ b/mirascope/core/base/tool.py @@ -205,7 +205,11 @@ def warn_for_unsupported_configurations(cls) -> None: UserWarning, ) - if "strict" in cls.model_config and cls.__provider__ not in ["openai", "azure"]: + if "strict" in cls.model_config and cls.__provider__ not in [ + "openai", + "azure", + "google", + ]: warnings.warn( f"{cls.__provider__} does not support strict structured outputs, but " "you have configured `strict=True` in your `ResponseModelConfigDict`. " diff --git a/mirascope/core/google/_call.py b/mirascope/core/google/_call.py index c34acc6c1..261a44378 100644 --- a/mirascope/core/google/_call.py +++ b/mirascope/core/google/_call.py @@ -56,7 +56,7 @@ def recommend_book(genre: str) -> str: output_parser (Callable[[GoogleCallResponse | ResponseModelT], Any]): A function for parsing the call response whose value will be returned in place of the original call response. - json_modem (bool): Whether to use JSON Mode. + json_mode (bool): Whether to use JSON Mode. client (object): An optional custom client to use in place of the default client. call_params (GoogleCallParams): The `GoogleCallParams` call parameters to use in the API call. diff --git a/mirascope/core/google/_utils/_setup_call.py b/mirascope/core/google/_utils/_setup_call.py index 4ea97fdad..754d72be2 100644 --- a/mirascope/core/google/_utils/_setup_call.py +++ b/mirascope/core/google/_utils/_setup_call.py @@ -147,11 +147,16 @@ def setup_call( if json_mode: with _generate_content_config_context(call_kwargs) as config: - if not tools: - config.response_mime_type = "application/json" - messages[-1]["parts"].append( # pyright: ignore [reportTypedDictNotRequiredAccess, reportOptionalMemberAccess, reportArgumentType] - PartDict(text=_utils.json_mode_content(response_model)) - ) # pyright: ignore [reportTypedDictNotRequiredAccess, reportOptionalMemberAccess, reportArgumentType] + config.response_mime_type = "application/json" + if response_model: + config.response_schema = response_model + + elif not tools: + messages[-1][ + "parts" + ].append( # pyright: ignore [reportTypedDictNotRequiredAccess, reportOptionalMemberAccess, reportArgumentType] + PartDict(text=_utils.json_mode_content(None)) + ) # pyright: ignore [reportTypedDictNotRequiredAccess, reportOptionalMemberAccess, reportArgumentType] elif response_model: assert tool_types, "At least one tool must be provided for extraction." with _generate_content_config_context(call_kwargs) as config: diff --git a/tests/core/google/_utils/test_setup_call.py b/tests/core/google/_utils/test_setup_call.py index f364422cc..4805b93ee 100644 --- a/tests/core/google/_utils/test_setup_call.py +++ b/tests/core/google/_utils/test_setup_call.py @@ -384,3 +384,180 @@ def test_setup_call_with_call_params( call_params, # Call params should be passed directly convert_common_call_params, ) + + +class MyTestModel(BaseModel): + param: str + optional_param: int | None = None + + +# Helper to structure mock returns for _utils.setup_call +def _get_mock_base_setup_call_value(messages=None, call_kwargs_override=None): + base_messages = ( + messages + if messages is not None + else [{"role": "user", "parts": [{"text": "test content"}]}] + ) + base_call_kwargs = call_kwargs_override if call_kwargs_override is not None else {} + return [ + MagicMock(), + base_messages, + MagicMock(), + base_call_kwargs, + ] # prompt_template, messages, tool_types, call_kwargs + + +@patch( + "mirascope.core.google._utils._setup_call.convert_message_params", + new_callable=MagicMock, +) +@patch("mirascope.core.google._utils._setup_call._utils", new_callable=MagicMock) +@patch("mirascope.core.google._utils._setup_call.Client", new_callable=MagicMock) +def test_setup_call_json_mode_with_response_model( + mock_client_class: MagicMock, + mock_utils: MagicMock, + mock_convert_message_params: MagicMock, +) -> None: + """Tests `setup_call` with `json_mode=True` and a `response_model`. + It should set `response_schema` in the config and not add generic JSON content. + """ + mock_client = MagicMock() + mock_client_class.return_value = mock_client + mock_utils.setup_call.return_value = _get_mock_base_setup_call_value() + mock_convert_message_params.side_effect = lambda x, y: x # Passthrough + + _, _, messages_out, _, call_kwargs = setup_call( + model="google-1.5-flash", + client=None, + fn=MagicMock(), + fn_args={}, + dynamic_config=None, + tools=None, + json_mode=True, + call_params={}, + response_model=MyTestModel, + stream=False, + ) + + assert "config" in call_kwargs + config = call_kwargs["config"] + assert isinstance(config, GenerateContentConfig) + assert config.response_mime_type == "application/json" + assert config.response_schema == MyTestModel + + # Ensure json_mode_content was NOT added to messages + # Assuming last message is user message and we check its parts + last_message_parts = messages_out[-1]["parts"] + for part in last_message_parts: + if isinstance(part, dict) and "text" in part: + assert mock_utils.json_mode_content.call_count == 0 + + +@patch( + "mirascope.core.google._utils._setup_call.convert_message_params", + new_callable=MagicMock, +) +@patch("mirascope.core.google._utils._setup_call._utils", new_callable=MagicMock) +@patch("mirascope.core.google._utils._setup_call.Client", new_callable=MagicMock) +def test_setup_call_json_mode_no_response_model_no_tools( + mock_client_class: MagicMock, + mock_utils: MagicMock, + mock_convert_message_params: MagicMock, +) -> None: + """Tests `setup_call` with `json_mode=True`, no `response_model`, and no `tools`.""" + mock_client = MagicMock() + mock_client_class.return_value = mock_client + # Simulate initial messages from base setup call + initial_messages = [{"role": "user", "parts": [{"text": "initial prompt"}]}] + mock_utils.setup_call.return_value = _get_mock_base_setup_call_value( + messages=initial_messages + ) + mock_convert_message_params.side_effect = lambda x, y: x # Passthrough + mock_utils.json_mode_content.return_value = "JSON_MODE_PROMPT_TEXT" + + _, _, messages_out, _, call_kwargs = setup_call( + model="google-1.5-flash", + client=None, + fn=MagicMock(), + fn_args={}, + dynamic_config=None, + tools=None, # Explicitly no tools + json_mode=True, + call_params={}, + response_model=None, # Explicitly no response model + stream=False, + ) + + assert "config" in call_kwargs + config = call_kwargs["config"] + assert isinstance(config, GenerateContentConfig) + assert config.response_mime_type == "application/json" + assert getattr(config, "response_schema", None) is None # Schema should not be set + + # Ensure json_mode_content WAS added to the last message + assert len(messages_out) > 0 + last_message_parts = messages_out[-1]["parts"] + assert any( + part.get("text") == "JSON_MODE_PROMPT_TEXT" + for part in last_message_parts + if isinstance(part, dict) + ) + mock_utils.json_mode_content.assert_called_once_with(None) + + +@patch( + "mirascope.core.google._utils._setup_call.convert_message_params", + new_callable=MagicMock, +) +@patch("mirascope.core.google._utils._setup_call._utils", new_callable=MagicMock) +@patch("mirascope.core.google._utils._setup_call.Client", new_callable=MagicMock) +def test_setup_call_json_mode_with_response_model_sets_schema( + mock_client_class: MagicMock, + mock_utils: MagicMock, + mock_convert_message_params: MagicMock, + mock_base_setup_call: MagicMock, +) -> None: + """Tests `setup_call` with `json_mode=True` and a `response_model`. + + It should set `response_schema` in the config and not add generic JSON content by calling + `_utils.json_mode_content`. + """ + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + mock_utils.setup_call = mock_base_setup_call + + initial_messages = [{"role": "user", "parts": [{"text": "initial prompt"}]}] + mock_base_setup_call.return_value[1] = initial_messages + mock_base_setup_call.return_value[3] = {} + + mock_convert_message_params.side_effect = lambda x, y: x + + class TestResponseModel(BaseModel): + foo: str + bar: int | None = None + + _, _, messages_out, _, call_kwargs_out = setup_call( + model="google-1.5-flash", + client=None, + fn=MagicMock(), + fn_args={}, + dynamic_config=None, + tools=None, + json_mode=True, + call_params={}, + response_model=TestResponseModel, # Use the concrete model class + stream=False, + ) + + assert "config" in call_kwargs_out + config = call_kwargs_out["config"] + assert isinstance(config, GenerateContentConfig) + assert config.response_mime_type == "application/json" + assert config.response_schema == TestResponseModel + + mock_utils.json_mode_content.assert_not_called() + + assert len(messages_out) == len(initial_messages) + if initial_messages and messages_out: + assert messages_out[-1]["parts"] == initial_messages[-1]["parts"] From 5d9086233d74f0d9328aa0d001d28e34d0c5afaa Mon Sep 17 00:00:00 2001 From: Priyanshu Sharma Date: Fri, 30 May 2025 23:51:19 +0530 Subject: [PATCH 2/2] feat: Force response_schema usage when strict=True is set in response_model config --- mirascope/core/google/_utils/_setup_call.py | 3 +- tests/core/google/_utils/test_setup_call.py | 81 +++++++++++++++++++++ 2 files changed, 82 insertions(+), 2 deletions(-) diff --git a/mirascope/core/google/_utils/_setup_call.py b/mirascope/core/google/_utils/_setup_call.py index 754d72be2..4ca78bc42 100644 --- a/mirascope/core/google/_utils/_setup_call.py +++ b/mirascope/core/google/_utils/_setup_call.py @@ -145,12 +145,11 @@ def setup_call( Part.model_validate(part) for part in message_parts ] - if json_mode: + if json_mode or (response_model and getattr(response_model, "model_config", {}).get("strict", False)): with _generate_content_config_context(call_kwargs) as config: config.response_mime_type = "application/json" if response_model: config.response_schema = response_model - elif not tools: messages[-1][ "parts" diff --git a/tests/core/google/_utils/test_setup_call.py b/tests/core/google/_utils/test_setup_call.py index 4805b93ee..e0784085e 100644 --- a/tests/core/google/_utils/test_setup_call.py +++ b/tests/core/google/_utils/test_setup_call.py @@ -561,3 +561,84 @@ class TestResponseModel(BaseModel): assert len(messages_out) == len(initial_messages) if initial_messages and messages_out: assert messages_out[-1]["parts"] == initial_messages[-1]["parts"] + + +@patch( + "mirascope.core.google._utils._setup_call.convert_message_params", + new_callable=MagicMock, +) +@patch("mirascope.core.google._utils._setup_call._utils", new_callable=MagicMock) +@patch("mirascope.core.google._utils._setup_call.Client", new_callable=MagicMock) +def test_setup_call_strict_response_model_forces_schema( + mock_client_class: MagicMock, + mock_utils: MagicMock, + mock_convert_message_params: MagicMock, +) -> None: + """Test that setup_call uses response_schema when strict=True is set in response_model's config, + even when json_mode is False. + """ + mock_client = MagicMock() + mock_client_class.return_value = mock_client + mock_utils.setup_call.return_value = _get_mock_base_setup_call_value() + mock_convert_message_params.side_effect = lambda x, y: x + + class StrictResponseModel(BaseModel): + name: str + age: int + + model_config = {"strict": True} + + # Test with json_mode=False but strict=True + _, _, messages_out, _, call_kwargs = setup_call( + model="google-1.5-flash", + client=None, + fn=MagicMock(), + fn_args={}, + dynamic_config=None, + tools=None, + json_mode=False, # json_mode is False + call_params={}, + response_model=StrictResponseModel, + stream=False, + ) + + # Should still use response_schema because strict=True + assert "config" in call_kwargs + config = call_kwargs["config"] + assert isinstance(config, GenerateContentConfig) + assert config.response_mime_type == "application/json" + assert config.response_schema == StrictResponseModel + + # Should not use tools path + assert config.tool_config is None + + # Should not append json_mode_content + mock_utils.json_mode_content.assert_not_called() + + # Test with both json_mode=True and strict=True + mock_utils.json_mode_content.reset_mock() + _, _, messages_out, _, call_kwargs = setup_call( + model="google-1.5-flash", + client=None, + fn=MagicMock(), + fn_args={}, + dynamic_config=None, + tools=None, + json_mode=True, # json_mode is True + call_params={}, + response_model=StrictResponseModel, + stream=False, + ) + + # Should still use response_schema + assert "config" in call_kwargs + config = call_kwargs["config"] + assert isinstance(config, GenerateContentConfig) + assert config.response_mime_type == "application/json" + assert config.response_schema == StrictResponseModel + + # Should not use tools path + assert config.tool_config is None + + # Should not append json_mode_content + mock_utils.json_mode_content.assert_not_called()