Skip to content

feat!: remove JsonSchema class, remove JsonSchemaType class, remove conversion between Schema and JsonSchema #2730

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat!: remove JsonSchema class, remove JsonSchemaType class, remove conversion between Schema and JsonSchema #2730

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
198 changes: 12 additions & 186 deletions google/genai/_automatic_function_calling_util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,6 @@
'_add_unevaluated_items_to_fixed_len_tuple_schema',
'_is_builtin_primitive_or_compound',
'_is_default_value_compatible',
'_parse_schema_from_parameter',
'_get_required_fields',
]

Expand Down Expand Up @@ -137,189 +136,16 @@ def _is_default_value_compatible(
return False


def _parse_schema_from_parameter( # type: ignore[return]
api_option: Literal['VERTEX_AI', 'GEMINI_API'],
param: inspect.Parameter,
func_name: str,
) -> types.Schema:
"""parse schema from parameter.

from the simplest case to the most complex case.
"""
schema = types.Schema()
default_value_error_msg = (
f'Default value {param.default} of parameter {param} of function'
f' {func_name} is not compatible with the parameter annotation'
f' {param.annotation}.'
)
if _is_builtin_primitive_or_compound(param.annotation):
if param.default is not inspect.Parameter.empty:
if not _is_default_value_compatible(param.default, param.annotation):
raise ValueError(default_value_error_msg)
schema.default = param.default
schema.type = _py_builtin_type_to_schema_type[param.annotation]
return schema
if (
isinstance(param.annotation, VersionedUnionType)
# only parse simple UnionType, example int | str | float | bool
# complex UnionType will be invoked in raise branch
and all(
(_is_builtin_primitive_or_compound(arg) or arg is type(None))
for arg in get_args(param.annotation)
)
):
schema.type = _py_builtin_type_to_schema_type[dict]
schema.any_of = []
unique_types = set()
for arg in get_args(param.annotation):
if arg.__name__ == 'NoneType': # Optional type
schema.nullable = True
continue
schema_in_any_of = _parse_schema_from_parameter(
api_option,
inspect.Parameter(
'item', inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=arg
),
func_name,
)
if (
schema_in_any_of.model_dump_json(exclude_none=True)
not in unique_types
):
schema.any_of.append(schema_in_any_of)
unique_types.add(schema_in_any_of.model_dump_json(exclude_none=True))
if len(schema.any_of) == 1: # param: list | None -> Array
schema.type = schema.any_of[0].type
schema.any_of = None
if (
param.default is not inspect.Parameter.empty
and param.default is not None
):
if not _is_default_value_compatible(param.default, param.annotation):
raise ValueError(default_value_error_msg)
schema.default = param.default
return schema
if isinstance(param.annotation, _GenericAlias) or isinstance(
param.annotation, builtin_types.GenericAlias
):
origin = get_origin(param.annotation)
args = get_args(param.annotation)
if origin is dict:
schema.type = _py_builtin_type_to_schema_type[dict]
if param.default is not inspect.Parameter.empty:
if not _is_default_value_compatible(param.default, param.annotation):
raise ValueError(default_value_error_msg)
schema.default = param.default
return schema
if origin is Literal:
if not all(isinstance(arg, str) for arg in args):
raise ValueError(
f'Literal type {param.annotation} must be a list of strings.'
)
schema.type = _py_builtin_type_to_schema_type[str]
schema.enum = list(args)
if param.default is not inspect.Parameter.empty:
if not _is_default_value_compatible(param.default, param.annotation):
raise ValueError(default_value_error_msg)
schema.default = param.default
return schema
if origin is list:
schema.type = _py_builtin_type_to_schema_type[list]
schema.items = _parse_schema_from_parameter(
api_option,
inspect.Parameter(
'item',
inspect.Parameter.POSITIONAL_OR_KEYWORD,
annotation=args[0],
),
func_name,
)
if param.default is not inspect.Parameter.empty:
if not _is_default_value_compatible(param.default, param.annotation):
raise ValueError(default_value_error_msg)
schema.default = param.default
return schema
if origin is Union:
schema.any_of = []
schema.type = _py_builtin_type_to_schema_type[dict]
unique_types = set()
for arg in args:
# The first check is for NoneType in Python 3.9, since the __name__
# attribute is not available in Python 3.9
if type(arg) is type(None) or (
hasattr(arg, '__name__') and arg.__name__ == 'NoneType'
): # Optional type
schema.nullable = True
continue
schema_in_any_of = _parse_schema_from_parameter(
api_option,
inspect.Parameter(
'item',
inspect.Parameter.POSITIONAL_OR_KEYWORD,
annotation=arg,
),
func_name,
)
if (
len(param.annotation.__args__) == 2
and type(None) in param.annotation.__args__
): # Optional type
for optional_arg in param.annotation.__args__:
if (
hasattr(optional_arg, '__origin__')
and optional_arg.__origin__ is list
):
# Optional type with list, for example Optional[list[str]]
schema.items = schema_in_any_of.items
if (
schema_in_any_of.model_dump_json(exclude_none=True)
not in unique_types
):
schema.any_of.append(schema_in_any_of)
unique_types.add(schema_in_any_of.model_dump_json(exclude_none=True))
if len(schema.any_of) == 1: # param: Union[List, None] -> Array
schema.type = schema.any_of[0].type
schema.any_of = None
if (
param.default is not None
and param.default is not inspect.Parameter.empty
):
if not _is_default_value_compatible(param.default, param.annotation):
raise ValueError(default_value_error_msg)
schema.default = param.default
return schema
# all other generic alias will be invoked in raise branch
if (
# for user defined class, we only support pydantic model
_extra_utils.is_annotation_pydantic_model(param.annotation)
):
if (
param.default is not inspect.Parameter.empty
and param.default is not None
):
schema.default = param.default
schema.type = _py_builtin_type_to_schema_type[dict]
schema.properties = {}
for field_name, field_info in param.annotation.model_fields.items():
schema.properties[field_name] = _parse_schema_from_parameter(
api_option,
inspect.Parameter(
field_name,
inspect.Parameter.POSITIONAL_OR_KEYWORD,
annotation=field_info.annotation,
),
func_name,
)
schema.required = _get_required_fields(schema)
return schema
_raise_for_unsupported_param(param, func_name, ValueError)


def _get_required_fields(schema: types.Schema) -> Optional[list[str]]:
if not schema.properties:
def _get_required_fields(json_schema: dict[str, Any]) -> Optional[list[str]]:
properties = json_schema.get('properties', {})
if not properties:
return None
return [
field_name
for field_name, field_schema in schema.properties.items()
if not field_schema.nullable and field_schema.default is None
]
required_fields = []
for field_name, field_schema in properties.items():
if not field_schema:
continue
if 'nullable' in field_schema and not field_schema['nullable']:
required_fields.append(field_name)
if 'default' not in field_schema:
required_fields.append(field_name)
return required_fields
42 changes: 1 addition & 41 deletions google/genai/_mcp_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,11 +49,7 @@ def mcp_to_gemini_tool(tool: McpTool) -> types.Tool:
function_declarations=[{
"name": tool.name,
"description": tool.description,
"parameters": types.Schema.from_json_schema(
json_schema=types.JSONSchema(
**_filter_to_supported_schema(tool.inputSchema)
)
),
"parameters_json_schema": tool.inputSchema,
}]
)

Expand Down Expand Up @@ -127,42 +123,6 @@ def set_mcp_usage_header(headers: dict[str, str]) -> None:
).lstrip()


def _filter_to_supported_schema(
schema: _common.StringDict,
) -> _common.StringDict:
"""Filters the schema to only include fields that are supported by JSONSchema."""
supported_fields: set[str] = set(types.JSONSchema.model_fields.keys())

supported_fields.update([
"additionalProperties", "anyOf", "oneOf", "$defs", "$ref"
])

schema_field_names = (
"items",
"additionalProperties",
"additional_properties",
)
list_schema_field_names = ("anyOf", "any_of", "oneOf", "one_of")
dict_schema_field_names = ("properties", "defs", "$defs")

filtered_schema: dict[str, Any] = {}
for field_name, field_value in schema.items():
if field_name in schema_field_names:
filtered_schema[field_name] = _filter_to_supported_schema(field_value)
elif field_name in list_schema_field_names:
filtered_schema[field_name] = [
_filter_to_supported_schema(value) for value in field_value
]
elif field_name in dict_schema_field_names:
filtered_schema[field_name] = {
key: _filter_to_supported_schema(value)
for key, value in field_value.items()
}
elif field_name in supported_fields:
filtered_schema[field_name] = field_value

return filtered_schema

@contextlib.asynccontextmanager
async def _connect_agent_platform_mcp(api_client: Any, toolset_name: str) -> typing.AsyncIterator[Any]:
"""Internal helper to manage the Agent Platform MCP lifecycle per request."""
Expand Down
6 changes: 6 additions & 0 deletions google/genai/tests/afc/test_generate_content_stream_afc_thoughts.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,9 @@ def get_current_weather(location: str) -> str:
pytest_plugins = ('pytest_asyncio',)


@pytest.mark.skip(
'AFC is in progress of refactoring, this case will be updated by Yvonne'
)
def test_generate_content_stream_with_function_and_thought_summaries(client):
"""Test when function tools are provided and thought summaries are enabled.

Expand All @@ -54,6 +57,9 @@ def test_generate_content_stream_with_function_and_thought_summaries(client):
assert chunk is not None


@pytest.mark.skip(
'AFC is in progress of refactoring, this case will be updated by Yvonne'
)
@pytest.mark.asyncio
async def test_generate_content_stream_with_function_and_thought_summaries_async(
client,
Expand Down
Loading
Loading