Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 25 additions & 3 deletions libs/genai/langchain_google_genai/chat_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -2220,14 +2220,26 @@ def _prepare_params(
gen_config = {**gen_config, **generation_config}

response_mime_type = kwargs.get("response_mime_type", self.response_mime_type)
if response_mime_type is not None:
gen_config["response_mime_type"] = response_mime_type

response_schema = kwargs.get("response_schema", self.response_schema)

# In case passed in as a direct kwarg
response_json_schema = kwargs.get("response_json_schema")

# Handle response_format from ProviderStrategy (OpenAI-style format)
# Format: {'type': 'json_schema', 'json_schema': {'schema': {...}}}
# Only extract from response_format if response_json_schema is not already set
if response_json_schema is None:
response_format = kwargs.get("response_format")
if response_format is not None and isinstance(response_format, dict):
if response_format.get("type") == "json_schema":
json_schema_obj = response_format.get("json_schema", {})
if (
isinstance(json_schema_obj, dict)
and "schema" in json_schema_obj
):
# Extract the actual schema from the nested structure
response_json_schema = json_schema_obj["schema"]

# Handle both response_schema and response_json_schema
# (Regardless, we use `response_json_schema` in the request)
schema_to_use = (
Expand All @@ -2236,6 +2248,16 @@ def _prepare_params(
else response_schema
)

# Automatically set response_mime_type to "application/json" when
# response_schema is provided but response_mime_type is not set. This
# enables seamless support for structured output strategies like
# ProviderStrategy.
if schema_to_use is not None and response_mime_type is None:
response_mime_type = "application/json"

if response_mime_type is not None:
gen_config["response_mime_type"] = response_mime_type

if schema_to_use is not None:
if response_mime_type != "application/json":
param_name = (
Expand Down
106 changes: 106 additions & 0 deletions libs/genai/tests/unit_tests/test_chat_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -2957,6 +2957,112 @@ def test_response_schema_mime_type_validation() -> None:
assert llm_with_json_schema is not None


def _convert_proto_to_dict(obj: Any) -> Any:
"""Recursively convert proto objects to dicts for comparison."""
if isinstance(obj, dict):
return {k: _convert_proto_to_dict(v) for k, v in obj.items()}
if isinstance(obj, (list, tuple)):
return [_convert_proto_to_dict(item) for item in obj]
if hasattr(obj, "__class__") and "proto" in str(type(obj)):
# Try to convert proto object to dict
try:
if hasattr(obj, "__iter__") and not isinstance(obj, str):
converted = dict(obj)
# Recursively convert nested proto objects
return {k: _convert_proto_to_dict(v) for k, v in converted.items()}
return obj
except (TypeError, ValueError):
return obj
return obj


def test_response_format_provider_strategy() -> None:
"""Test that `response_format` from ProviderStrategy is correctly handled."""
llm = ChatGoogleGenerativeAI(
model=MODEL_NAME, google_api_key=SecretStr(FAKE_API_KEY)
)

schema_dict: dict[str, Any] = {
"type": "object",
"properties": {
"sentiment": {
"type": "string",
"enum": ["positive", "negative", "neutral"],
},
"confidence": {"type": "number", "minimum": 0, "maximum": 1},
},
"required": ["sentiment", "confidence"],
"additionalProperties": False,
}

# Test response_format with ProviderStrategy format (OpenAI-style)
response_format: dict[str, Any] = {
"type": "json_schema",
"json_schema": {
"name": "response_format_test",
"schema": schema_dict,
},
}

gen_config = llm._prepare_params(stop=None, response_format=response_format)

# response_json_schema may be converted to proto object, so convert to
# dict for comparison
schema = _convert_proto_to_dict(gen_config.response_json_schema)
assert schema == schema_dict
assert gen_config.response_mime_type == "application/json"

# Test that response_json_schema takes precedence over response_format
different_schema: dict[str, Any] = {
"type": "object",
"properties": {"age": {"type": "integer"}},
"required": ["age"],
}

gen_config_2 = llm._prepare_params(
stop=None,
response_format=response_format,
response_json_schema=different_schema,
)

# response_json_schema may be converted to proto object, so convert to
# dict for comparison
schema_2 = _convert_proto_to_dict(gen_config_2.response_json_schema)
assert schema_2 == different_schema
assert gen_config_2.response_mime_type == "application/json"

old_schema: dict[str, Any] = {
"type": "object",
"properties": {"old_field": {"type": "string"}},
"required": ["old_field"],
}

gen_config_3 = llm._prepare_params(
stop=None,
response_schema=old_schema,
response_format=response_format,
)

# response_json_schema may be converted to proto object, so convert to
# dict for comparison
schema_3 = _convert_proto_to_dict(gen_config_3.response_json_schema)
assert schema_3 == schema_dict
assert gen_config_3.response_mime_type == "application/json"

invalid_response_format: dict[str, Any] = {"type": "invalid_type"}
gen_config_4 = llm._prepare_params(
stop=None,
response_format=invalid_response_format,
response_schema=schema_dict,
)
# Should fall back to response_schema
# response_json_schema may be converted to proto object, so convert to
# dict for comparison
schema_4 = _convert_proto_to_dict(gen_config_4.response_json_schema)
assert schema_4 == schema_dict
assert gen_config_4.response_mime_type == "application/json"


def test_is_new_gemini_model() -> None:
assert _is_gemini_3_or_later("gemini-3.0-pro") is True
assert _is_gemini_3_or_later("gemini-2.5-pro") is False
Expand Down