diff --git a/google/genai/_automatic_function_calling_util.py b/google/genai/_automatic_function_calling_util.py index ec7f9a702..b8b66ac70 100644 --- a/google/genai/_automatic_function_calling_util.py +++ b/google/genai/_automatic_function_calling_util.py @@ -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', ] @@ -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 and field_name not in required_fields: + required_fields.append(field_name) + return required_fields diff --git a/google/genai/_extra_utils.py b/google/genai/_extra_utils.py index 20ad6a1fe..37abcff4d 100644 --- a/google/genai/_extra_utils.py +++ b/google/genai/_extra_utils.py @@ -53,8 +53,8 @@ def _create_generate_content_config_model( - config: types.GenerateContentConfigOrDict, -) -> types.GenerateContentConfig: + config: Union[types.GenerateContentConfigOrDict, types.ChatConfig], +) -> Union[types.GenerateContentConfig, types.ChatConfig]: if isinstance(config, dict): return types.GenerateContentConfig(**config) else: @@ -114,7 +114,9 @@ def format_destination( def find_afc_incompatible_tool_indexes( - config: Optional[types.GenerateContentConfigOrDict] = None, + config: Optional[ + Union[types.GenerateContentConfigOrDict, types.ChatConfig] + ] = None, is_agent_platform: bool = False, ) -> list[int]: """Checks if the config contains any AFC incompatible tools.""" @@ -408,82 +410,85 @@ async def get_function_response_parts_async( return func_response_parts -def should_disable_afc( - config: Optional[types.GenerateContentConfigOrDict] = None, + +def should_enable_afc( + config: Optional[types.ChatConfig] = None, ) -> bool: """Returns whether automatic function calling is enabled.""" if not config: return False - config_model = _create_generate_content_config_model(config) + else: + config_model = config + # If max_remote_calls is less or equal to 0, warn and disable AFC. if ( config_model - and config_model.automatic_function_calling - and config_model.automatic_function_calling.maximum_remote_calls + and config_model.automatic_function_calling_config + and config_model.automatic_function_calling_config.maximum_remote_calls is not None - and int(config_model.automatic_function_calling.maximum_remote_calls) <= 0 + and int(config_model.automatic_function_calling_config.maximum_remote_calls) <= 0 ): logger.warning( 'max_remote_calls in automatic_function_calling_config' - f' {config_model.automatic_function_calling.maximum_remote_calls} is' + f' {config_model.automatic_function_calling_config.maximum_remote_calls} is' ' less than or equal to 0. Disabling automatic function calling.' - ' Please set max_remote_calls to a positive integer.' ) - return True + return False - # Default to enable AFC if not specified. + # Default to disable AFC if not specified. if ( - not config_model.automatic_function_calling - or config_model.automatic_function_calling.disable is None + not config_model.automatic_function_calling_config + or config_model.automatic_function_calling_config.enable is None ): return False if ( - config_model.automatic_function_calling.disable - and config_model.automatic_function_calling.maximum_remote_calls + not config_model.automatic_function_calling_config.enable + and config_model.automatic_function_calling_config.maximum_remote_calls is not None # exclude the case where max_remote_calls is set to 10 by default. and 'maximum_remote_calls' - in config_model.automatic_function_calling.model_fields_set - and int(config_model.automatic_function_calling.maximum_remote_calls) > 0 + in config_model.automatic_function_calling_config.model_fields_set + and int(config_model.automatic_function_calling_config.maximum_remote_calls) > 0 ): logger.warning( - '`automatic_function_calling.disable` is set to `True`. And' + '`automatic_function_calling.enable` is set to `False`. And' ' `automatic_function_calling.maximum_remote_calls` is a' ' positive number' - f' {config_model.automatic_function_calling.maximum_remote_calls}.' + f' {config_model.automatic_function_calling_config.maximum_remote_calls}.' ' Disabling automatic function calling. If you want to enable' ' automatic function calling, please set' - ' `automatic_function_calling.disable` to `False` or leave it unset,' - ' and set `automatic_function_calling.maximum_remote_calls` to a' - ' positive integer or leave' + ' `automatic_function_calling.enable` to `True` and set' + ' `automatic_function_calling.maximum_remote_calls` to a positive' + ' integer or leave' ' `automatic_function_calling.maximum_remote_calls` unset.' ) + return False - return config_model.automatic_function_calling.disable + return config_model.automatic_function_calling_config.enable def get_max_remote_calls_afc( - config: Optional[types.GenerateContentConfigOrDict] = None, + config: Optional[types.ChatConfig] = None, ) -> int: if not config: return _DEFAULT_MAX_REMOTE_CALLS_AFC """Returns the remaining remote calls for automatic function calling.""" - if should_disable_afc(config): + if not should_enable_afc(config): raise ValueError( 'automatic function calling is not enabled, but SDK is trying to get' ' max remote calls.' ) config_model = _create_generate_content_config_model(config) if ( - not config_model.automatic_function_calling - or config_model.automatic_function_calling.maximum_remote_calls is None + not config_model.automatic_function_calling_config # type: ignore[attr-defined] + or config_model.automatic_function_calling_config.maximum_remote_calls is None # type: ignore[attr-defined] ): return _DEFAULT_MAX_REMOTE_CALLS_AFC - return int(config_model.automatic_function_calling.maximum_remote_calls) + return int(config_model.automatic_function_calling_config.maximum_remote_calls) # type: ignore[attr-defined] -def raise_error_for_afc_incompatible_config(config: Optional[types.GenerateContentConfig] +def raise_error_for_afc_incompatible_config(config: Optional[types.ChatConfig] ) -> None: """Raises an error if the config is not compatible with AFC.""" if ( @@ -492,36 +497,41 @@ def raise_error_for_afc_incompatible_config(config: Optional[types.GenerateConte or not config.tool_config.function_calling_config ): return - afc_config = config.automatic_function_calling - disable_afc_config = afc_config.disable if afc_config else False + afc_config = config.automatic_function_calling_config + enable_afc = afc_config.enable if afc_config else False stream_function_call = ( config.tool_config.function_calling_config.stream_function_call_arguments ) - if stream_function_call and not disable_afc_config: + if stream_function_call and enable_afc: raise ValueError( - 'Running in streaming mode with stream_function_call_arguments' - ' enabled, this feature is not compatible with automatic function' - ' calling (AFC). Please set config.automatic_function_calling.disable' - ' to True to disable AFC or leave config.tool_config.' - ' function_calling_config.stream_function_call_arguments to be empty' - ' or set to False to disable streaming function call arguments.' + 'Running in streaming mode with stream_function_call_arguments enabled,' + ' this feature is not compatible with automatic function calling (AFC).' + ' Please set ChatConfig.automatic_function_calling_config.enable to' + ' False to disable AFC or leave config.tool_config.' + ' function_calling_config.stream_function_call_arguments to be empty or' + ' set to False to disable streaming function call arguments.' ) + def should_append_afc_history( config: Optional[types.GenerateContentConfigOrDict] = None, ) -> bool: if not config: return True config_model = _create_generate_content_config_model(config) - if not config_model.automatic_function_calling: + if not config_model.automatic_function_calling_config: # type: ignore[attr-defined] return True - return not config_model.automatic_function_calling.ignore_call_history + return not config_model.automatic_function_calling_config.ignore_call_history # type: ignore[attr-defined] def parse_config_for_mcp_usage( - config: Optional[types.GenerateContentConfigOrDict] = None, -) -> Optional[types.GenerateContentConfig]: + config: Optional[ + Union[ + types.GenerateContentConfigOrDict, types.ChatConfig + ] + ] = None, +) -> Optional[Union[types.GenerateContentConfig, types.ChatConfig]]: """Returns a parsed config with an appended MCP header if MCP tools or sessions are used.""" if not config: return None @@ -684,28 +694,33 @@ def prepare_resumable_upload( def has_agent_platform_mcp_servers( config: Optional[types.GenerateContentConfigOrDict] = None, ) -> bool: - """Checks whether the configuration contains any MCP server requests.""" - if not config: - return False - config_model = _create_generate_content_config_model(config) - if not config_model.tools: - return False - - for tool in config_model.tools: - if getattr(tool, 'mcp_servers', None): - return True + """Checks whether the configuration contains any MCP server requests.""" + if not config: return False + config_model = _create_generate_content_config_model(config) + if not config_model.tools: + return False + + for tool in config_model.tools: + if getattr(tool, 'mcp_servers', None): + return True + return False def get_usage_header( - config: Optional[types.GenerateContentConfigOrDict] = None, + config: Optional[ + Union[types.GenerateContentConfigOrDict, types.ChatConfig] + ] = None, usage: str = 'afc', -) -> types.GenerateContentConfig: +) -> Union[types.GenerateContentConfig, types.ChatConfig]: """Sets the afc version label.""" usage_header = f'google-genai-sdk/{public_version.__version__}+{usage}' if not config: - config_model = types.GenerateContentConfig() + config_model = ( + types.ChatConfig() if usage == 'chat' else types.GenerateContentConfig() + ) elif isinstance(config, dict): + # ChatConfig doesn't support dict type config_model = types.GenerateContentConfig(**config) else: config_model = config diff --git a/google/genai/_mcp_utils.py b/google/genai/_mcp_utils.py index 74b24b363..6b6d35b39 100644 --- a/google/genai/_mcp_utils.py +++ b/google/genai/_mcp_utils.py @@ -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, }] ) @@ -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.""" diff --git a/google/genai/chats.py b/google/genai/chats.py index 8782f47ce..1b3d69138 100644 --- a/google/genai/chats.py +++ b/google/genai/chats.py @@ -14,15 +14,26 @@ # from collections.abc import Iterator +import contextlib +import logging +import pydantic +from pydantic import BaseModel import sys -from typing import AsyncIterator, Awaitable, Optional, Union, get_args - +from typing import Any, AsyncIterator, get_args, Optional, Union +from . import errors from . import _extra_utils from . import _transformers as t from . import types -from .models import AsyncModels, Models -from .types import Content, ContentOrDict, GenerateContentConfigOrDict, GenerateContentResponse, Part, PartUnionDict +from .models import AsyncModels +from .models import Models +from .types import ChatConfig +from .types import Content +from .types import ContentOrDict +from .types import GenerateContentConfig +from .types import GenerateContentResponse +from .types import Part +from .types import PartUnionDict if sys.version_info >= (3, 10): @@ -30,6 +41,8 @@ else: from typing_extensions import TypeGuard +logger = logging.getLogger("google_genai.chats") + def _validate_content(content: Content) -> bool: if not content.parts: @@ -104,6 +117,21 @@ def _extract_curated_history( return curated_history +def _extract_generate_content_config( + chat_config: Optional[ChatConfig] = None, +) -> Optional[GenerateContentConfig]: + """Slices a ChatConfig instance to a GenerateContentConfig instance. + """ + if not chat_config: + return None + + generate_content_config_kwargs = { + field: getattr(chat_config, field) + for field in GenerateContentConfig.model_fields + } + return GenerateContentConfig(**generate_content_config_kwargs) + + class _BaseChat: """Base chat session.""" @@ -111,11 +139,11 @@ def __init__( self, *, model: str, - config: Optional[GenerateContentConfigOrDict] = None, + config: Optional[ChatConfig] = None, history: list[ContentOrDict], ): self._model = model - self._config = _extra_utils.get_usage_header(config, usage="chat") + self._config: ChatConfig = _extra_utils.get_usage_header(config, usage="chat") # type: ignore[assignment] content_models = [] for content in history: if not isinstance(content, Content): @@ -134,7 +162,6 @@ def record_history( self, user_input: Content, model_output: list[Content], - automatic_function_calling_history: list[Content], is_valid: bool, ) -> None: """Records the chat history. @@ -145,20 +172,10 @@ def record_history( user_input: The user's input content. model_output: A list of `Content` from the model's response. This can be an empty list if the model produced no output. - automatic_function_calling_history: A list of `Content` representing the - history of automatic function calls, including the user input as the - first entry. is_valid: A boolean flag indicating whether the current model output is considered valid. """ - input_contents = ( - # Because the AFC input contains the entire curated chat history in - # addition to the new user input, we need to truncate the AFC history - # to deduplicate the existing chat history. - automatic_function_calling_history[len(self._curated_history) :] - if automatic_function_calling_history - else [user_input] - ) + input_contents = [user_input] # Appends an empty content when model returns empty response, so that the # history is always alternating between user and model. output_contents = ( @@ -212,7 +229,7 @@ def __init__( *, modules: Models, model: str, - config: Optional[GenerateContentConfigOrDict] = None, + config: Optional[ChatConfig] = None, history: list[ContentOrDict], ): self._modules = modules @@ -225,7 +242,7 @@ def __init__( def send_message( self, message: Union[list[PartUnionDict], PartUnionDict], - config: Optional[GenerateContentConfigOrDict] = None, + config: Optional[ChatConfig] = None, ) -> GenerateContentResponse: """Sends the conversation history with the additional message and returns the model's response. @@ -244,44 +261,149 @@ def send_message( chat = client.chats.create(model='gemini-2.0-flash') response = chat.send_message('tell me a story') """ + method_config = config if config else self._config + generate_content_config = _extract_generate_content_config(method_config) + parsed_config = _extra_utils.parse_config_for_mcp_usage( + generate_content_config + ) + if ( + parsed_config + and parsed_config.tools + and _extra_utils._mcp_utils.has_mcp_session_usage(parsed_config.tools) # type: ignore[attr-defined] + ): + raise errors.UnsupportedFunctionError( + "MCP sessions are not supported in synchronous methods." + ) if not _is_part_type(message): raise ValueError( f"Message must be a valid part type: {types.PartUnion} or" f" {types.PartUnionDict}, got {type(message)}" ) - input_content = t.t_content(message) - method_config = config if config else self._config - method_config = _extra_utils.get_usage_header( - method_config, usage="chat" + + incompatible_tools_indexes = ( + _extra_utils.find_afc_incompatible_tool_indexes(generate_content_config) + ) + user_input = t.t_content(message) + contents_to_model = self._curated_history + [user_input] # type: ignore[arg-type] + if not _extra_utils.should_enable_afc(method_config): + response = self._modules.generate_content( + model=self._model, + contents=contents_to_model, # type: ignore[arg-type] + config=generate_content_config, + ) + model_output = ( + [response.candidates[0].content] + if response.candidates and response.candidates[0].content + else [] + ) + self.record_history( + user_input=user_input, + model_output=model_output, + is_valid=_validate_response(response), + ) + return response + if incompatible_tools_indexes: + original_tools_length = 0 + if method_config.tools: + original_tools_length = len(method_config.tools) + if len(incompatible_tools_indexes) != original_tools_length: + indices_str = ", ".join(map(str, incompatible_tools_indexes)) + logger.warning( + "Tools at indices [%s] are not compatible with automatic" + " function calling (AFC). AFC is disabled. If AFC is" + " intended, please include python callables in the tool" + " list, and do not include function declaration and MCP" + " server in the tool list.", + indices_str, + ) + + response = self._modules.generate_content( + model=self._model, + contents=contents_to_model, # type: ignore[arg-type] + config=generate_content_config, + ) + model_output = ( + [response.candidates[0].content] + if response.candidates and response.candidates[0].content + else [] + ) + self.record_history( + user_input=user_input, + model_output=model_output, + is_valid=_validate_response(response), + ) + return response + # AFC handling + remaining_remote_calls_afc = _extra_utils.get_max_remote_calls_afc( + method_config ) - response = self._modules.generate_content( - model=self._model, - contents=self._curated_history + [input_content], # type: ignore[arg-type] - config=method_config, + logger.info( + f"AFC is enabled with max remote calls: {remaining_remote_calls_afc}." ) + response = types.GenerateContentResponse() + function_map = _extra_utils.get_function_map(generate_content_config) + i = 0 + while remaining_remote_calls_afc > 0: + if function_map: + generate_content_config = _extra_utils.get_usage_header( + generate_content_config, usage="afc" + ) + i += 1 + response = self._modules.generate_content( + model=self._model, + contents=contents_to_model, # type: ignore[arg-type] + config=generate_content_config, + ) + if ( + not function_map + or not response + or not response.candidates + or not response.candidates[0].content + or not response.candidates[0].content.parts + ): + break + + func_response_parts = _extra_utils.get_function_response_parts( + response, function_map + ) + if not func_response_parts: + break + logger.info(f"AFC remote call {i} is done.") + remaining_remote_calls_afc -= 1 + if remaining_remote_calls_afc == 0: + logger.info("Reached max remote calls for automatic function calling.") + func_call_content = response.candidates[0].content + func_response_content = types.Content( + role="user", parts=func_response_parts + ) + contents_to_model.append(func_call_content) + contents_to_model.append(func_response_content) + model_output = [func_call_content] + self.record_history( + user_input=user_input, + model_output=model_output, + is_valid=_validate_response(response), + ) + user_input = func_response_content + model_output = ( [response.candidates[0].content] if response.candidates and response.candidates[0].content else [] ) - automatic_function_calling_history = ( - response.automatic_function_calling_history - if response.automatic_function_calling_history - else [] - ) self.record_history( - user_input=input_content, + user_input=user_input, model_output=model_output, - automatic_function_calling_history=automatic_function_calling_history, is_valid=_validate_response(response), ) return response + def send_message_stream( self, message: Union[list[PartUnionDict], PartUnionDict], - config: Optional[GenerateContentConfigOrDict] = None, + config: Optional[ChatConfig] = None, ) -> Iterator[GenerateContentResponse]: """Sends the conversation history with the additional message and yields the model's response in chunks. @@ -302,45 +424,155 @@ def send_message_stream( print(chunk.text) """ + method_config = config if config else self._config + generate_content_config = _extract_generate_content_config(method_config) + parsed_config = _extra_utils.parse_config_for_mcp_usage( + generate_content_config + ) + if ( + parsed_config + and parsed_config.tools + and _extra_utils._mcp_utils.has_mcp_session_usage(parsed_config.tools) # type: ignore[attr-defined] + ): + raise errors.UnsupportedFunctionError( + "MCP sessions are not supported in synchronous methods." + ) if not _is_part_type(message): raise ValueError( f"Message must be a valid part type: {types.PartUnion} or" f" {types.PartUnionDict}, got {type(message)}" ) - input_content = t.t_content(message) - output_contents = [] + incompatible_tools_indexes = ( + _extra_utils.find_afc_incompatible_tool_indexes(generate_content_config) + ) + user_input = t.t_content(message) + contents_to_model = self._curated_history + [user_input] # type: ignore[arg-type] + model_output = [] finish_reason = None is_valid = True - chunk = None - method_config = config if config else self._config - method_config = _extra_utils.get_usage_header( - method_config, usage="chat" + enable_afc = _extra_utils.should_enable_afc(method_config) + if enable_afc and incompatible_tools_indexes: + original_tools_length = 0 + if method_config.tools: + original_tools_length = len(method_config.tools) + if len(incompatible_tools_indexes) != original_tools_length: + indices_str = ", ".join(map(str, incompatible_tools_indexes)) + logger.warning( + "Tools at indices [%s] are not compatible with automatic" + " function calling (AFC). AFC is disabled. If AFC is" + " intended, please include python callables in the tool" + " list, and do not include function declaration and MCP" + " server in the tool list.", + indices_str, + ) + enable_afc = False + + if not enable_afc: + if isinstance(self._modules, Models): + for chunk in self._modules.generate_content_stream( + model=self._model, + contents=contents_to_model, # type: ignore[arg-type] + config=generate_content_config, + ): + if not _validate_response(chunk): + is_valid = False + if chunk.candidates and chunk.candidates[0].content: + model_output.append(chunk.candidates[0].content) + if chunk.candidates and chunk.candidates[0].finish_reason: + finish_reason = chunk.candidates[0].finish_reason + yield chunk + self.record_history( + user_input=user_input, + model_output=model_output, + is_valid=is_valid + and model_output is not None + and finish_reason is not None, + ) + return + + # AFC handling + _extra_utils.raise_error_for_afc_incompatible_config(method_config) + remaining_remote_calls_afc = _extra_utils.get_max_remote_calls_afc( + method_config + ) + logger.info( + f"AFC is enabled with max remote calls: {remaining_remote_calls_afc}." ) + function_map = _extra_utils.get_function_map(generate_content_config) + i = 0 if isinstance(self._modules, Models): - for chunk in self._modules.generate_content_stream( - model=self._model, - contents=self._curated_history + [input_content], # type: ignore[arg-type] - config=method_config, - ): - if not _validate_response(chunk): - is_valid = False - if chunk.candidates and chunk.candidates[0].content: - output_contents.append(chunk.candidates[0].content) - if chunk.candidates and chunk.candidates[0].finish_reason: - finish_reason = chunk.candidates[0].finish_reason - yield chunk - automatic_function_calling_history = ( - chunk.automatic_function_calling_history - if chunk is not None and chunk.automatic_function_calling_history - else [] - ) + while remaining_remote_calls_afc > 0: + i += 1 + if function_map: + generate_content_config = _extra_utils.get_usage_header( + generate_content_config, usage="afc" + ) + response_stream = self._modules.generate_content_stream( + model=self._model, + contents=contents_to_model, # type: ignore[arg-type] + config=generate_content_config, + ) + + model_output = [] + finish_reason = None + is_valid = True + func_response_parts = [] + chunk = None # type: ignore[assignment] + + for chunk in response_stream: + if not _validate_response(chunk): + is_valid = False + + if ( + function_map + and chunk.candidates + and chunk.candidates[0].content + and chunk.candidates[0].content.parts + ): + chunk_func_response_parts = ( + _extra_utils.get_function_response_parts(chunk, function_map) + ) + if chunk_func_response_parts: + func_response_parts.extend(chunk_func_response_parts) + + if chunk.candidates and chunk.candidates[0].content: + model_output.append(chunk.candidates[0].content) + if chunk.candidates and chunk.candidates[0].finish_reason: + finish_reason = chunk.candidates[0].finish_reason + yield chunk + + if not function_map or not func_response_parts: + break + + logger.info(f"AFC remote call {i} is done.") + remaining_remote_calls_afc -= 1 + if remaining_remote_calls_afc == 0: + logger.info( + "Reached max remote calls for automatic function calling." + ) + + if chunk and chunk.candidates and chunk.candidates[0].content: + func_response_content = types.Content( + role="user", parts=func_response_parts + ) + contents_to_model.extend(model_output) + contents_to_model.append(func_response_content) + + self.record_history( + user_input=user_input, + model_output=model_output, + is_valid=is_valid, + ) + user_input = func_response_content + self.record_history( - user_input=input_content, - model_output=output_contents, - automatic_function_calling_history=automatic_function_calling_history, - is_valid=is_valid - and output_contents is not None - and finish_reason is not None, + user_input=user_input, + model_output=model_output, + is_valid=bool( + is_valid + and model_output is not None + and finish_reason is not None + ), ) @@ -354,7 +586,7 @@ def create( self, *, model: str, - config: Optional[GenerateContentConfigOrDict] = None, + config: Optional[ChatConfig] = None, history: Optional[list[ContentOrDict]] = None, ) -> Chat: """Creates a new chat session. @@ -383,7 +615,7 @@ def __init__( *, modules: AsyncModels, model: str, - config: Optional[GenerateContentConfigOrDict] = None, + config: Optional[ChatConfig] = None, history: list[ContentOrDict], ): self._modules = modules @@ -396,7 +628,7 @@ def __init__( async def send_message( self, message: Union[list[PartUnionDict], PartUnionDict], - config: Optional[GenerateContentConfigOrDict] = None, + config: Optional[ChatConfig] = None, ) -> GenerateContentResponse: """Sends the conversation history with the additional message and returns model's response. @@ -415,43 +647,227 @@ async def send_message( chat = client.aio.chats.create(model='gemini-2.0-flash') response = await chat.send_message('tell me a story') """ + method_config = config if config else self._config if not _is_part_type(message): raise ValueError( f"Message must be a valid part type: {types.PartUnion} or" f" {types.PartUnionDict}, got {type(message)}" ) - input_content = t.t_content(message) - method_config = config if config else self._config - method_config = _extra_utils.get_usage_header( - method_config, usage="chat" - ) - response = await self._modules.generate_content( - model=self._model, - contents=self._curated_history + [input_content], # type: ignore[arg-type] - config=method_config, - ) - model_output = ( - [response.candidates[0].content] - if response.candidates and response.candidates[0].content - else [] - ) - automatic_function_calling_history = ( - response.automatic_function_calling_history - if response.automatic_function_calling_history - else [] - ) - self.record_history( - user_input=input_content, - model_output=model_output, - automatic_function_calling_history=automatic_function_calling_history, - is_valid=_validate_response(response), + + user_input = t.t_content(message) + contents_to_model = self._curated_history + [user_input] # type: ignore[arg-type] + + generate_content_config = _extract_generate_content_config(method_config) + if not _extra_utils.should_enable_afc(method_config): + response = await self._modules.generate_content( + model=self._model, + contents=contents_to_model, # type: ignore[arg-type] + config=generate_content_config, + ) + model_output = ( + [response.candidates[0].content] + if response.candidates and response.candidates[0].content + else [] + ) + self.record_history( + user_input=user_input, + model_output=model_output, + is_valid=_validate_response(response), + ) + return response + + incompatible_tools_indexes = ( + _extra_utils.find_afc_incompatible_tool_indexes( + generate_content_config, + is_agent_platform=getattr( + self._modules._api_client, "vertexai", False + ), + ) ) - return response + + if not method_config: + parsed_config = None + else: + parsed_config = method_config.model_copy(deep=True) + + if incompatible_tools_indexes: + original_tools_length = 0 + if method_config.tools: + original_tools_length = len(method_config.tools) + + if len(incompatible_tools_indexes) != original_tools_length: + indices_str = ", ".join(map(str, incompatible_tools_indexes)) + logger.warning( + "Tools at indices [%s] are not compatible with automatic" + " function calling (AFC). AFC is disabled. If AFC is" + " intended, please include python callables in the tool" + " list, and do not include function declaration and MCP" + " server in the tool list.", + indices_str, + ) + + response = await self._modules.generate_content( + model=self._model, + contents=contents_to_model, # type: ignore[arg-type] + config=generate_content_config, + ) + model_output = ( + [response.candidates[0].content] + if response.candidates and response.candidates[0].content + else [] + ) + self.record_history( + user_input=user_input, + model_output=model_output, + is_valid=_validate_response(response), + ) + return response + + # AFC handling + async with contextlib.AsyncExitStack() as stack: + # Intercept Agent Platform MCP servers and open connections + if ( + self._modules._api_client.vertexai + and _extra_utils.has_agent_platform_mcp_servers( + generate_content_config + ) + and generate_content_config is not None + ): + new_tools: list[Any] = [] + if generate_content_config.tools: + for tool in generate_content_config.tools: + if isinstance(tool, types.Tool) and tool.mcp_servers: + # Only keep the tool if it has fields besides mcp_servers + if ( + tool.function_declarations + or tool.google_search + or tool.retrieval + or tool.google_search_retrieval + or tool.code_execution + ): + tool_copy = tool.model_copy(update={'mcp_servers': None}) + new_tools.append(tool_copy) + + for server in tool.mcp_servers: + if ( + getattr(server, 'streamable_http_transport', None) + is not None + ): + raise ValueError( + "The 'streamable_http_transport' parameter is only" + ' supported in Gemini Developer API mode, not in Gemini' + ' Enterprise Agent Platform mode.' + ) + + # Open the stream and tie its lifespan to the AsyncExitStack + if server.name is not None: + session = await stack.enter_async_context( + _extra_utils._mcp_utils._connect_agent_platform_mcp( # type: ignore[attr-defined] + self._modules._api_client, server.name + ) + ) + new_tools.append(session) + else: + raise ValueError( + "Agent Platform MCP servers require a 'name' field." + ) + else: + new_tools.append(tool) + generate_content_config.tools = new_tools + + # Convert active sessions to tools and adapters + final_generate_content_config, mcp_to_genai_tool_adapters = ( + await _extra_utils.parse_config_for_mcp_sessions( + generate_content_config, + is_agent_platform=getattr( + self._modules._api_client, "vertexai", False + ), + ) + ) + remaining_remote_calls_afc = _extra_utils.get_max_remote_calls_afc( + method_config + ) + + logger.info( + f"AFC is enabled with max remote calls: {remaining_remote_calls_afc}." + ) + + response = types.GenerateContentResponse() + function_map = _extra_utils.get_function_map( + final_generate_content_config, + mcp_to_genai_tool_adapters, + is_caller_method_async=True, + ) + + i = 0 + while remaining_remote_calls_afc > 0: + if function_map: + final_generate_content_config = _extra_utils.get_usage_header( + final_generate_content_config, usage="afc" + ) + i += 1 + response = await self._modules.generate_content( + model=self._model, + contents=contents_to_model, # type: ignore[arg-type] + config=final_generate_content_config, + ) + if ( + not function_map + or not response + or not response.candidates + or not response.candidates[0].content + or not response.candidates[0].content.parts + ): + break + + func_response_parts = ( + await _extra_utils.get_function_response_parts_async( + response, function_map + ) + ) + if not func_response_parts: + break + + logger.info(f"AFC remote call {i} is done.") + remaining_remote_calls_afc -= 1 + if remaining_remote_calls_afc == 0: + logger.info( + "Reached max remote calls for automatic function calling." + ) + + func_call_content = response.candidates[0].content + func_response_content = types.Content( + role="user", parts=func_response_parts + ) + + contents_to_model.append(func_call_content) + contents_to_model.append(func_response_content) + + model_output = [func_call_content] + self.record_history( + user_input=user_input, + model_output=model_output, + is_valid=_validate_response(response), + ) + user_input = func_response_content + + model_output = ( + [response.candidates[0].content] + if response.candidates and response.candidates[0].content + else [] + ) + self.record_history( + user_input=user_input, + model_output=model_output, + is_valid=_validate_response(response), + ) + return response + async def send_message_stream( self, message: Union[list[PartUnionDict], PartUnionDict], - config: Optional[GenerateContentConfigOrDict] = None, + config: Optional[ChatConfig] = None, ) -> AsyncIterator[GenerateContentResponse]: """Sends the conversation history with the additional message and yields the model's response in chunks. @@ -471,7 +887,6 @@ async def send_message_stream( async for chunk in await chat.send_message_stream('tell me a story'): print(chunk.text) """ - if not _is_part_type(message): raise ValueError( f"Message must be a valid part type: {types.PartUnion} or" @@ -479,40 +894,223 @@ async def send_message_stream( ) input_content = t.t_content(message) - method_config = config if config else self._config - method_config = _extra_utils.get_usage_header( - method_config, usage="chat" - ) - async def async_generator(): # type: ignore[no-untyped-def] - output_contents = [] - finish_reason = None - is_valid = True - chunk = None - async for chunk in await self._modules.generate_content_stream( # type: ignore[attr-defined] - model=self._model, - contents=self._curated_history + [input_content], # type: ignore[arg-type] - config=method_config, - ): - if not _validate_response(chunk): + method_config = config if config else self._config + generate_content_config = _extract_generate_content_config(method_config) + parsed_generate_content_config = _extra_utils.parse_config_for_mcp_usage( + generate_content_config + ) + enable_afc = _extra_utils.should_enable_afc(method_config) + incompatible_tools_indexes = ( + _extra_utils.find_afc_incompatible_tool_indexes( + parsed_generate_content_config, + is_agent_platform=getattr( + self._modules._api_client, "vertexai", False + ), + ) + ) + user_input = input_content + contents_to_model = self._curated_history + [user_input] # type: ignore[arg-type] + if enable_afc and incompatible_tools_indexes: + original_tools_length = 0 + if method_config.tools: + original_tools_length = len(method_config.tools) + if len(incompatible_tools_indexes) != original_tools_length: + indices_str = ", ".join(map(str, incompatible_tools_indexes)) + logger.warning( + "Tools at indices [%s] are not compatible with automatic" + " function calling (AFC). AFC is disabled. If AFC is" + " intended, please include python callables in the tool" + " list, and do not include function declaration and MCP" + " server in the tool list.", + indices_str, + ) + enable_afc = False + + if not enable_afc: + output_contents = [] + finish_reason = None + is_valid = True + async for chunk in await self._modules.generate_content_stream( # type: ignore[attr-defined] + model=self._model, + contents=contents_to_model, # type: ignore[arg-type] + config=parsed_generate_content_config, + ): + if not _validate_response(chunk): + is_valid = False + if chunk.candidates and chunk.candidates[0].content: + output_contents.append(chunk.candidates[0].content) + if chunk.candidates and chunk.candidates[0].finish_reason: + finish_reason = chunk.candidates[0].finish_reason + yield chunk + + if not output_contents or finish_reason is None: is_valid = False - if chunk.candidates and chunk.candidates[0].content: - output_contents.append(chunk.candidates[0].content) - if chunk.candidates and chunk.candidates[0].finish_reason: - finish_reason = chunk.candidates[0].finish_reason - yield chunk - if not output_contents or finish_reason is None: - is_valid = False + self.record_history( + user_input=user_input, + model_output=output_contents, + is_valid=is_valid, + ) + return - self.record_history( - user_input=input_content, - model_output=output_contents, - automatic_function_calling_history=chunk.automatic_function_calling_history - if chunk is not None and chunk.automatic_function_calling_history - else [], - is_valid=is_valid, + # AFC handling + _extra_utils.raise_error_for_afc_incompatible_config( + method_config ) + async with contextlib.AsyncExitStack() as stack: + # Intercept Agent Platform MCP servers and open connections + if ( + self._modules._api_client.vertexai + and _extra_utils.has_agent_platform_mcp_servers( + parsed_generate_content_config, + ) + and parsed_generate_content_config is not None + ): + new_tools: list[Any] = [] + if parsed_generate_content_config.tools: + for tool in parsed_generate_content_config.tools: + if isinstance(tool, types.Tool) and tool.mcp_servers: + # Only keep the tool if it has fields besides mcp_servers + if ( + tool.function_declarations + or tool.google_search + or tool.retrieval + or tool.google_search_retrieval + or tool.code_execution + ): + tool_copy = tool.model_copy(update={'mcp_servers': None}) + new_tools.append(tool_copy) + + for server in tool.mcp_servers: + if ( + getattr(server, 'streamable_http_transport', None) + is not None + ): + raise ValueError( + "The 'streamable_http_transport' parameter is only" + ' supported in Gemini Developer API mode, not in Gemini' + ' Enterprise Agent Platform mode.' + ) + + # Open the stream and tie its lifespan to the AsyncExitStack + if server.name is not None: + session = await stack.enter_async_context( + _extra_utils._mcp_utils._connect_agent_platform_mcp( # type: ignore[attr-defined] + self._modules._api_client, server.name + ) + ) + new_tools.append(session) + else: + raise ValueError( + "Agent Platform MCP servers require a 'name' field." + ) + else: + new_tools.append(tool) + parsed_generate_content_config.tools = new_tools + + # Convert active sessions to tools and adapters + final_parsed_generate_content_config, mcp_to_genai_tool_adapters = ( + await _extra_utils.parse_config_for_mcp_sessions( + parsed_generate_content_config, + is_agent_platform=getattr( + self._modules._api_client, "vertexai", False + ), + ) + ) + + remaining_remote_calls_afc = _extra_utils.get_max_remote_calls_afc( + method_config + ) + + logger.info( + "AFC is enabled with max remote calls:" + f" {remaining_remote_calls_afc}." + ) + + function_map = _extra_utils.get_function_map( + final_parsed_generate_content_config, + mcp_to_genai_tool_adapters, + is_caller_method_async=True, + ) + + i = 0 + while remaining_remote_calls_afc > 0: + i += 1 + if function_map: + final_parsed_generate_content_config = ( + _extra_utils.get_usage_header( + final_parsed_generate_content_config, usage="afc" + ) + ) + response_stream = await self._modules.generate_content_stream( + model=self._model, + contents=contents_to_model, # type: ignore[arg-type] + config=final_parsed_generate_content_config, + ) + + model_output: list[types.Content] = [] + finish_reason = None + is_valid = True + func_response_parts: list[types.Part] = [] + + async for chunk in response_stream: + if not _validate_response(chunk): + is_valid = False + + if ( + function_map + and chunk.candidates + and chunk.candidates[0].content + and chunk.candidates[0].content.parts + ): + chunk_func_response_parts = ( + await _extra_utils.get_function_response_parts_async( + chunk, function_map + ) + ) + if chunk_func_response_parts: + func_response_parts.extend(chunk_func_response_parts) + + if chunk.candidates and chunk.candidates[0].content: + model_output.append(chunk.candidates[0].content) + if chunk.candidates and chunk.candidates[0].finish_reason: + finish_reason = chunk.candidates[0].finish_reason + yield chunk + + if not function_map or not func_response_parts: + break + + logger.info(f"AFC remote call {i} is done.") + remaining_remote_calls_afc -= 1 + if remaining_remote_calls_afc == 0: + logger.info( + "Reached max remote calls for automatic function calling." + ) + + func_response_content = types.Content( + role="user", parts=func_response_parts + ) + + contents_to_model.extend(model_output) + contents_to_model.append(func_response_content) + + self.record_history( + user_input=user_input, + model_output=model_output, + is_valid=is_valid, + ) + user_input = func_response_content + + self.record_history( + user_input=user_input, + model_output=model_output, + is_valid=bool( + is_valid + and model_output + and finish_reason is not None + ), + ) return async_generator() # type: ignore[no-untyped-call, no-any-return] @@ -527,7 +1125,7 @@ def create( self, *, model: str, - config: Optional[GenerateContentConfigOrDict] = None, + config: Optional[ChatConfig] = None, history: Optional[list[ContentOrDict]] = None, ) -> AsyncChat: """Creates a new chat session. diff --git a/google/genai/models.py b/google/genai/models.py index 3e34a0bf0..9c5803951 100644 --- a/google/genai/models.py +++ b/google/genai/models.py @@ -6542,9 +6542,6 @@ def generate_content( # scones. """ - incompatible_tools_indexes = ( - _extra_utils.find_afc_incompatible_tool_indexes(config) - ) parsed_config = _extra_utils.parse_config_for_mcp_usage(config) if ( parsed_config @@ -6554,96 +6551,9 @@ def generate_content( raise errors.UnsupportedFunctionError( 'MCP sessions are not supported in synchronous methods.' ) - if _extra_utils.should_disable_afc(parsed_config): - return self._generate_content( - model=model, contents=contents, config=parsed_config - ) - if incompatible_tools_indexes: - original_tools_length = 0 - if isinstance(config, types.GenerateContentConfig): - if config.tools: - original_tools_length = len(config.tools) - elif isinstance(config, dict): - tools = config.get('tools', []) - if tools: - original_tools_length = len(tools) - if len(incompatible_tools_indexes) != original_tools_length: - indices_str = ', '.join(map(str, incompatible_tools_indexes)) - logger.warning( - 'Tools at indices [%s] are not compatible with automatic function ' - 'calling (AFC). AFC is disabled. If AFC is intended, please ' - 'include python callables in the tool list, and do not include ' - 'function declaration and MCP server in the tool list.', - indices_str, - ) - return self._generate_content( - model=model, contents=contents, config=parsed_config - ) - - remaining_remote_calls_afc = _extra_utils.get_max_remote_calls_afc( - parsed_config - ) - logger.info( - f'AFC is enabled with max remote calls: {remaining_remote_calls_afc}.' + return self._generate_content( + model=model, contents=contents, config=parsed_config ) - automatic_function_calling_history: list[types.Content] = [] - response = types.GenerateContentResponse() - i = 0 - while remaining_remote_calls_afc > 0: - parsed_config_to_call = ( - parsed_config.model_copy(deep=True) if parsed_config else None - ) - function_map = _extra_utils.get_function_map(parsed_config) - if function_map: - parsed_config_to_call = _extra_utils.get_usage_header( - parsed_config_to_call - ) - i += 1 - response = self._generate_content( - model=model, contents=contents, config=parsed_config_to_call - ) - - if not function_map: - break - if not response: - break - if ( - not response.candidates - or not response.candidates[0].content - or not response.candidates[0].content.parts - ): - break - func_response_parts = _extra_utils.get_function_response_parts( - response, function_map - ) - if not func_response_parts: - break - logger.info(f'AFC remote call {i} is done.') - remaining_remote_calls_afc -= 1 - if remaining_remote_calls_afc == 0: - logger.info('Reached max remote calls for automatic function calling.') - - func_call_content = response.candidates[0].content - func_response_content = types.Content( - role='user', - parts=func_response_parts, - ) - contents = t.t_contents(contents) # type: ignore[assignment] - if not automatic_function_calling_history: - automatic_function_calling_history.extend(contents) # type: ignore[arg-type] - if isinstance(contents, list): - contents.append(func_call_content) # type: ignore[arg-type] - contents.append(func_response_content) # type: ignore[arg-type] - automatic_function_calling_history.append(func_call_content) - automatic_function_calling_history.append(func_response_content) - if ( - _extra_utils.should_append_afc_history(parsed_config) - and response is not None - ): - response.automatic_function_calling_history = ( - automatic_function_calling_history - ) - return response def generate_content_stream( self, @@ -6711,9 +6621,6 @@ def generate_content_stream( # scones. """ - incompatible_tools_indexes = ( - _extra_utils.find_afc_incompatible_tool_indexes(config) - ) parsed_config = _extra_utils.parse_config_for_mcp_usage(config) if ( parsed_config @@ -6723,128 +6630,10 @@ def generate_content_stream( raise errors.UnsupportedFunctionError( 'MCP sessions are not supported in synchronous methods.' ) - if _extra_utils.should_disable_afc(parsed_config): - yield from self._generate_content_stream( - model=model, contents=contents, config=parsed_config - ) - return - - if incompatible_tools_indexes: - original_tools_length = 0 - if isinstance(config, types.GenerateContentConfig): - if config.tools: - original_tools_length = len(config.tools) - elif isinstance(config, dict): - tools = config.get('tools', []) - if tools: - original_tools_length = len(tools) - if len(incompatible_tools_indexes) != original_tools_length: - indices_str = ', '.join(map(str, incompatible_tools_indexes)) - logger.warning( - 'Tools at indices [%s] are not compatible with automatic function ' - 'calling. AFC will be disabled.', - indices_str, - ) - yield from self._generate_content_stream( - model=model, contents=contents, config=parsed_config - ) - return - - # With tool compatibility confirmed, validate that the configuration are - # compatible with each other and raise an error if invalid. - _extra_utils.raise_error_for_afc_incompatible_config(parsed_config) - - remaining_remote_calls_afc = _extra_utils.get_max_remote_calls_afc( - parsed_config + yield from self._generate_content_stream( + model=model, contents=contents, config=parsed_config ) - logger.info( - f'AFC is enabled with max remote calls: {remaining_remote_calls_afc}.' - ) - automatic_function_calling_history: list[types.Content] = [] - chunk = None - func_response_parts = None - i = 0 - while remaining_remote_calls_afc > 0: - parsed_config_to_call = ( - parsed_config.model_copy(deep=True) if parsed_config else None - ) - function_map = _extra_utils.get_function_map(parsed_config) - if function_map: - parsed_config_to_call = _extra_utils.get_usage_header( - parsed_config_to_call - ) - i += 1 - response = self._generate_content_stream( - model=model, contents=contents, config=parsed_config_to_call - ) - - if i == 1: - # First request gets a function call. - # Then get function response parts. - # Yield chunks only if there's no function response parts. - for chunk in response: - if not function_map: - contents = _extra_utils.append_chunk_contents(contents, chunk) # type: ignore[assignment] - yield chunk - else: - if ( - not chunk.candidates - or not chunk.candidates[0].content - or not chunk.candidates[0].content.parts - ): - break - func_response_parts = _extra_utils.get_function_response_parts( - chunk, function_map - ) - if not func_response_parts: - contents = _extra_utils.append_chunk_contents(contents, chunk) # type: ignore[assignment] - yield chunk - - else: - # Second request and beyond, yield chunks. - for chunk in response: - if _extra_utils.should_append_afc_history(parsed_config): - chunk.automatic_function_calling_history = ( - automatic_function_calling_history - ) - contents = _extra_utils.append_chunk_contents(contents, chunk) # type: ignore[assignment] - yield chunk - if ( - chunk is None - or not chunk.candidates - or not chunk.candidates[0].content - or not chunk.candidates[0].content.parts - ): - break - func_response_parts = _extra_utils.get_function_response_parts( - chunk, function_map - ) - - if not function_map: - break - if not func_response_parts: - break - logger.info(f'AFC remote call {i} is done.') - remaining_remote_calls_afc -= 1 - if remaining_remote_calls_afc == 0: - logger.info('Reached max remote calls for automatic function calling.') - - # Append function response parts to contents for the next request. - if chunk is not None and chunk.candidates is not None: - func_call_content = chunk.candidates[0].content - func_response_content = types.Content( - role='user', - parts=func_response_parts, - ) - contents = t.t_contents(contents) # type: ignore[assignment] - if not automatic_function_calling_history: - automatic_function_calling_history.extend(contents) # type: ignore[arg-type] - if isinstance(contents, list) and func_call_content is not None: - contents.append(func_call_content) # type: ignore[arg-type] - contents.append(func_response_content) # type: ignore[arg-type] - if func_call_content is not None: - automatic_function_calling_history.append(func_call_content) - automatic_function_calling_history.append(func_response_content) + return @_common.experimental_warning( 'The generate_images method is deprecated and will be removed in the ' @@ -8680,14 +8469,6 @@ async def generate_content( print(response.text) # J'aime les bagels. """ - # Retrieve and cache any MCP sessions if provided. - incompatible_tools_indexes = ( - _extra_utils.find_afc_incompatible_tool_indexes( - config, - is_agent_platform=getattr(self._api_client, 'vertexai', False), - ) - ) - if not config: parsed_config = None elif isinstance(config, dict): @@ -8755,106 +8536,9 @@ async def generate_content( ) ) - if _extra_utils.should_disable_afc(final_parsed_config): - return await self._generate_content( - model=model, contents=contents, config=final_parsed_config - ) - - if incompatible_tools_indexes: - original_tools_length = 0 - if isinstance(config, types.GenerateContentConfig): - if config.tools: - original_tools_length = len(config.tools) - elif isinstance(config, dict): - tools = config.get('tools', []) - if tools: - original_tools_length = len(tools) - if len(incompatible_tools_indexes) != original_tools_length: - indices_str = ', '.join(map(str, incompatible_tools_indexes)) - logger.warning( - 'Tools at indices [%s] are not compatible with automatic function' - ' calling (AFC). AFC is disabled. If AFC is intended, please' - ' include python callables in the tool list, and do not include' - ' function declaration and MCP server in the tool list.', - indices_str, - ) - return await self._generate_content( - model=model, contents=contents, config=final_parsed_config - ) - - remaining_remote_calls_afc = _extra_utils.get_max_remote_calls_afc( - final_parsed_config - ) - logger.info( - f'AFC is enabled with max remote calls: {remaining_remote_calls_afc}.' + return await self._generate_content( + model=model, contents=contents, config=final_parsed_config ) - automatic_function_calling_history: list[types.Content] = [] - response = types.GenerateContentResponse() - - while remaining_remote_calls_afc > 0: - function_map = _extra_utils.get_function_map( - final_parsed_config, - mcp_to_genai_tool_adapters, - is_caller_method_async=True, - ) - final_parsed_config_to_call = ( - final_parsed_config.model_copy(deep=True) - if final_parsed_config - else None - ) - if function_map: - final_parsed_config_to_call = _extra_utils.get_usage_header( - final_parsed_config_to_call - ) - response = await self._generate_content( - model=model, contents=contents, config=final_parsed_config_to_call - ) - remaining_remote_calls_afc -= 1 - if remaining_remote_calls_afc == 0: - logger.info( - 'Reached max remote calls for automatic function calling.' - ) - - if not function_map: - break - if not response: - break - if ( - not response.candidates - or not response.candidates[0].content - or not response.candidates[0].content.parts - ): - break - func_response_parts = ( - await _extra_utils.get_function_response_parts_async( - response, function_map - ) - ) - if not func_response_parts: - break - func_call_content = response.candidates[0].content - func_response_content = types.Content( - role='user', - parts=func_response_parts, - ) - contents = t.t_contents(contents) # type: ignore[assignment] - if not automatic_function_calling_history: - automatic_function_calling_history.extend(contents) # type: ignore[arg-type] - if isinstance(contents, list): - contents.append(func_call_content) # type: ignore[arg-type] - contents.append(func_response_content) # type: ignore[arg-type] - automatic_function_calling_history.append(func_call_content) - automatic_function_calling_history.append(func_response_content) - - if ( - _extra_utils.should_append_afc_history(final_parsed_config) - and response is not None - ): - response.automatic_function_calling_history = ( - automatic_function_calling_history - ) - - return response async def generate_content_stream( self, @@ -8928,13 +8612,6 @@ async def generate_content_stream( else: parsed_config = config.model_copy(deep=True) - incompatible_tools_indexes = ( - _extra_utils.find_afc_incompatible_tool_indexes( - parsed_config, - is_agent_platform=getattr(self._api_client, 'vertexai', False), - ) - ) - async def stream_generator(): # type: ignore[no-untyped-def] # Use AsyncExitStack to keep MCP connections alive across the entire stream async with contextlib.AsyncExitStack() as stack: @@ -8992,159 +8669,12 @@ async def stream_generator(): # type: ignore[no-untyped-def] ) ) - if _extra_utils.should_disable_afc(final_parsed_config): - response = await self._generate_content_stream( - model=model, contents=contents, config=final_parsed_config - ) - async for chunk in response: # type: ignore[attr-defined] - yield chunk - return - - if incompatible_tools_indexes: - original_tools_length = 0 - if isinstance(config, types.GenerateContentConfig): - if config.tools: - original_tools_length = len(config.tools) - elif isinstance(config, dict): - tools = config.get('tools', []) - if tools: - original_tools_length = len(tools) - if len(incompatible_tools_indexes) != original_tools_length: - indices_str = ', '.join(map(str, incompatible_tools_indexes)) - logger.warning( - 'Tools at indices [%s] are not compatible with automatic' - ' function calling (AFC). AFC is disabled. If AFC is intended,' - ' please include python callables in the tool list, and do not' - ' include function declaration and MCP server in the tool' - ' list.', - indices_str, - ) - response = await self._generate_content_stream( - model=model, contents=contents, config=final_parsed_config - ) - async for chunk in response: # type: ignore[attr-defined] - yield chunk - return - - _extra_utils.raise_error_for_afc_incompatible_config( - final_parsed_config - ) - - remaining_remote_calls_afc = _extra_utils.get_max_remote_calls_afc( - final_parsed_config - ) - logger.info( - 'AFC is enabled with max remote calls:' - f' {remaining_remote_calls_afc}.' + response = await self._generate_content_stream( + model=model, contents=contents, config=final_parsed_config ) - automatic_function_calling_history: list[types.Content] = [] - func_response_parts = None - chunk = None - i = 0 - loop_contents = contents - - while remaining_remote_calls_afc > 0: - - function_map = _extra_utils.get_function_map( - final_parsed_config, - mcp_to_genai_tool_adapters, - is_caller_method_async=True, - ) - - final_parsed_config_to_call = ( - final_parsed_config.model_copy(deep=True) - if final_parsed_config - else None - ) - if function_map: - final_parsed_config_to_call = _extra_utils.get_usage_header( - final_parsed_config_to_call - ) - - i += 1 - - response = await self._generate_content_stream( - model=model, - contents=loop_contents, - config=final_parsed_config_to_call, - ) - - if i > 1: - logger.info(f'AFC remote call {i} is done.') - remaining_remote_calls_afc -= 1 - if i > 1 and remaining_remote_calls_afc == 0: - logger.info( - 'Reached max remote calls for automatic function calling.' - ) - - if i == 1: - async for chunk in response: # type: ignore[attr-defined] - if not function_map: - loop_contents = _extra_utils.append_chunk_contents( - loop_contents, chunk - ) - yield chunk - else: - if ( - not chunk.candidates - or not chunk.candidates[0].content - or not chunk.candidates[0].content.parts - ): - break - func_response_parts = ( - await _extra_utils.get_function_response_parts_async( - chunk, function_map - ) - ) - if not func_response_parts: - loop_contents = _extra_utils.append_chunk_contents( - loop_contents, chunk - ) - yield chunk - else: - async for chunk in response: # type: ignore[attr-defined] - if _extra_utils.should_append_afc_history(final_parsed_config): - chunk.automatic_function_calling_history = ( - automatic_function_calling_history - ) - loop_contents = _extra_utils.append_chunk_contents( - loop_contents, chunk - ) - yield chunk - if ( - chunk is None - or not chunk.candidates - or not chunk.candidates[0].content - or not chunk.candidates[0].content.parts - ): - break - func_response_parts = ( - await _extra_utils.get_function_response_parts_async( - chunk, function_map - ) - ) - - if not function_map or not func_response_parts: - break - - if chunk is None: - continue - - # Append function response parts to contents for the next request. - func_call_content = chunk.candidates[0].content - func_response_content = types.Content( - role='user', - parts=func_response_parts, - ) - loop_contents = t.t_contents(loop_contents) # type: ignore[assignment] - if not automatic_function_calling_history: - automatic_function_calling_history.extend(loop_contents) # type: ignore[arg-type] - if isinstance(loop_contents, list) and func_call_content is not None: - loop_contents.append(func_call_content) # type: ignore[arg-type] - loop_contents.append(func_response_content) # type: ignore[arg-type] - if func_call_content is not None: - automatic_function_calling_history.append(func_call_content) - automatic_function_calling_history.append(func_response_content) + async for chunk in response: # type: ignore[attr-defined] + yield chunk + return return stream_generator() # type: ignore[no-untyped-call, no-any-return] diff --git a/google/genai/tests/afc/test_generate_content_stream_afc_thoughts.py b/google/genai/tests/afc/test_generate_content_stream_afc_thoughts.py index 434259352..9331ac8fd 100644 --- a/google/genai/tests/afc/test_generate_content_stream_afc_thoughts.py +++ b/google/genai/tests/afc/test_generate_content_stream_afc_thoughts.py @@ -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. @@ -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, diff --git a/google/genai/tests/afc/test_generate_content_stream_afc.py b/google/genai/tests/afc/test_generate_content_stream_fc.py similarity index 62% rename from google/genai/tests/afc/test_generate_content_stream_afc.py rename to google/genai/tests/afc/test_generate_content_stream_fc.py index b9ec0e282..3952872f2 100644 --- a/google/genai/tests/afc/test_generate_content_stream_afc.py +++ b/google/genai/tests/afc/test_generate_content_stream_fc.py @@ -179,7 +179,7 @@ def mock_generate_content_stream_no_afc(): @pytest.fixture -def mock_generate_content_stream_with_afc(): +def mock_generate_content_stream_with_fc(): with mock.patch.object( models.Models, '_generate_content_stream' ) as mock_stream_with_afc: @@ -189,11 +189,6 @@ def mock_generate_content_stream_with_afc(): candidates=[types.Candidate(content=TEST_FUNCTION_CALL_CONTENT)] ) ], - [ - types.GenerateContentResponse( - candidates=[types.Candidate(content=TEST_AFC_TEXT_CONTENT)] - ) - ], ] yield mock_stream_with_afc @@ -214,7 +209,7 @@ async def async_generator(): @pytest.fixture -def mock_generate_content_stream_with_afc_async(): +def mock_generate_content_stream_with_fc_async(): with mock.patch.object( models.AsyncModels, '_generate_content_stream' ) as mock_stream_with_afc: @@ -224,14 +219,8 @@ async def async_generator_1(): candidates=[types.Candidate(content=TEST_FUNCTION_CALL_CONTENT)] ) - async def async_generator_2(): - yield types.GenerateContentResponse( - candidates=[types.Candidate(content=TEST_AFC_TEXT_CONTENT)] - ) - mock_stream_with_afc.side_effect = [ async_generator_1(), - async_generator_2(), ] yield mock_stream_with_afc @@ -255,36 +244,8 @@ def test_generate_content_stream_no_function_map( assert mock_get_function_response_parts_none.call_count == 0 -def test_generate_content_stream_afc_disabled( - mock_generate_content_stream_with_afc, - mock_get_function_response_parts_none, -): - """Test when function tools are provided but AFC is disabled. - - Expected to respond with function call. - """ - models_instance = models.Models(api_client_=mock_api_client) - stream = models_instance.generate_content_stream( - model='test_model', - contents='what is the weather in San Francisco?', - config=types.GenerateContentConfig( - tools=[get_current_weather], - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=True - ), - ), - ) - for chunk in stream: - # Work as manual function calling. - assert chunk.candidates[0].content.parts[0].function_call - - assert mock_generate_content_stream_with_afc.call_count == 1 - assert mock_get_function_response_parts_none.call_count == 0 - - def test_generate_content_stream_no_function_response( mock_generate_content_stream_no_afc, - mock_get_function_response_parts_none, ): """Test when function tools are provided and function responses are not returned. @@ -301,42 +262,10 @@ def test_generate_content_stream_no_function_response( assert chunk.text == TEST_NO_AFC_PART.text assert mock_generate_content_stream_no_afc.call_count == 1 - assert mock_get_function_response_parts_none.call_count == 1 - - -def test_generate_content_stream_with_function_tools_used( - mock_generate_content_stream_with_afc, - mock_get_function_response_parts, -): - """Test when function tools are provided and function responses are returned. - - Expected to answer weather based on function response. - """ - models_instance = models.Models(api_client_=mock_api_client) - config = types.GenerateContentConfig(tools=[get_current_weather]) - stream = models_instance.generate_content_stream( - model='test_model', - contents='what is the weather in San Francisco?', - config=config, - ) - - chunk = None - for chunk in stream: - assert chunk.text == TEST_AFC_TEXT_PART.text - - assert mock_generate_content_stream_with_afc.call_count == 2 - assert mock_get_function_response_parts.call_count == 2 - - assert chunk is not None - for i in range(len(chunk.automatic_function_calling_history)): - assert chunk.automatic_function_calling_history[i].model_dump( - exclude_none=True - ) == TEST_AFC_HISTORY[i].model_dump(exclude_none=True) def test_generate_content_stream_with_thought_summaries( - mock_generate_content_stream_with_afc, - mock_get_function_response_parts, + mock_generate_content_stream_with_fc, ): """Test when function tools are provided and thought summaries are enabled. @@ -355,16 +284,9 @@ def test_generate_content_stream_with_thought_summaries( chunk = None for chunk in stream: - assert chunk.text == TEST_AFC_TEXT_PART.text + assert chunk.function_calls == [TEST_FUNCTION_CALL_PART.function_call] - assert mock_generate_content_stream_with_afc.call_count == 2 - assert mock_get_function_response_parts.call_count == 2 - - assert chunk is not None - for i in range(len(chunk.automatic_function_calling_history)): - assert chunk.automatic_function_calling_history[i].model_dump( - exclude_none=True - ) == TEST_AFC_HISTORY[i].model_dump(exclude_none=True) + assert mock_generate_content_stream_with_fc.call_count == 1 @pytest.mark.asyncio @@ -387,38 +309,9 @@ async def test_generate_content_stream_no_function_map_async( assert mock_get_function_response_parts_none.call_count == 0 -@pytest.mark.asyncio -async def test_generate_content_stream_afc_disabled_async( - mock_generate_content_stream_with_afc_async, - mock_get_function_response_parts_none, -): - """Test when function tools are provided but AFC is disabled. - - Expected to respond with function call. - """ - models_instance = models.AsyncModels(api_client_=mock_api_client) - stream = await models_instance.generate_content_stream( - model='test_model', - contents='what is the weather in San Francisco?', - config=types.GenerateContentConfig( - tools=[get_current_weather], - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=True - ), - ), - ) - async for chunk in stream: - # Work as manual function calling. - assert chunk.candidates[0].content.parts[0].function_call - - assert mock_generate_content_stream_with_afc_async.call_count == 1 - assert mock_get_function_response_parts_none.call_count == 0 - - @pytest.mark.asyncio async def test_generate_content_stream_no_function_response_async( mock_generate_content_stream_no_afc_async, - mock_get_function_response_parts_none_async, ): """Test when function tools are provided and function responses are not returned. @@ -436,13 +329,10 @@ async def test_generate_content_stream_no_function_response_async( assert mock_generate_content_stream_no_afc_async.call_count == 1 - assert mock_get_function_response_parts_none_async.call_count == 1 - @pytest.mark.asyncio async def test_generate_content_stream_with_function_tools_used_async( - mock_generate_content_stream_with_afc_async, - mock_get_function_response_parts_async, + mock_generate_content_stream_with_fc_async, ): """Test when function tools are provided and function responses are returned. @@ -458,55 +348,14 @@ async def test_generate_content_stream_with_function_tools_used_async( chunk = None async for chunk in stream: - assert chunk.text == TEST_AFC_TEXT_PART.text - - assert mock_generate_content_stream_with_afc_async.call_count == 2 - - assert mock_get_function_response_parts_async.call_count == 2 - - assert chunk is not None - for i in range(len(chunk.automatic_function_calling_history)): - assert chunk.automatic_function_calling_history[i].model_dump( - exclude_none=True - ) == TEST_AFC_HISTORY[i].model_dump(exclude_none=True) - - -@pytest.mark.asyncio -async def test_generate_content_stream_with_function_async_function_used_async( - mock_generate_content_stream_with_afc_async, - mock_get_function_response_parts_async, -): - """Test when function tools are provided and function responses are returned. - - Expected to answer weather based on function response. - """ - models_instance = models.AsyncModels(api_client_=mock_api_client) - config = types.GenerateContentConfig(tools=[get_current_weather_async]) - stream = await models_instance.generate_content_stream( - model='test_model', - contents='what is the weather in San Francisco?', - config=config, - ) - - chunk = None - async for chunk in stream: - assert chunk.text == TEST_AFC_TEXT_PART.text - - assert mock_generate_content_stream_with_afc_async.call_count == 2 + assert chunk.function_calls == [TEST_FUNCTION_CALL_PART.function_call] - assert mock_get_function_response_parts_async.call_count == 2 - - assert chunk is not None - for i in range(len(chunk.automatic_function_calling_history)): - assert chunk.automatic_function_calling_history[i].model_dump( - exclude_none=True - ) == TEST_AFC_HISTORY[i].model_dump(exclude_none=True) + assert mock_generate_content_stream_with_fc_async.call_count == 1 @pytest.mark.asyncio async def test_generate_content_stream_with_thought_summaries_async( - mock_generate_content_stream_with_afc_async, - mock_get_function_response_parts_async, + mock_generate_content_stream_with_fc_async, ): """Test when function tools are provided and thought summaries are enabled. @@ -525,14 +374,7 @@ async def test_generate_content_stream_with_thought_summaries_async( chunk = None async for chunk in stream: - assert chunk.text == TEST_AFC_TEXT_PART.text - - assert mock_generate_content_stream_with_afc_async.call_count == 2 + assert chunk.function_calls == [TEST_FUNCTION_CALL_PART.function_call] - assert mock_get_function_response_parts_async.call_count == 2 + assert mock_generate_content_stream_with_fc_async.call_count == 1 - assert chunk is not None - for i in range(len(chunk.automatic_function_calling_history)): - assert chunk.automatic_function_calling_history[i].model_dump( - exclude_none=True - ) == TEST_AFC_HISTORY[i].model_dump(exclude_none=True) diff --git a/google/genai/tests/afc/test_get_max_remote_calls_for_afc.py b/google/genai/tests/afc/test_get_max_remote_calls_for_afc.py index 5a72020ed..14529c24b 100644 --- a/google/genai/tests/afc/test_get_max_remote_calls_for_afc.py +++ b/google/genai/tests/afc/test_get_max_remote_calls_for_afc.py @@ -26,28 +26,16 @@ def test_config_is_none(): def test_afc_unset_max_unset(): - assert get_max_remote_calls_afc(types.GenerateContentConfig()) == 10 + with pytest.raises(ValueError): + get_max_remote_calls_afc(types.ChatConfig()) == 10 def test_afc_unset_max_set(): - assert ( - get_max_remote_calls_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - maximum_remote_calls=20, - ), - ) - ) - == 20 - ) - - -def test_afc_disabled_max_unset(): with pytest.raises(ValueError): get_max_remote_calls_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=True, + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + maximum_remote_calls=20, ), ) ) @@ -56,8 +44,8 @@ def test_afc_disabled_max_unset(): def test_afc_disabled_max_set(): with pytest.raises(ValueError): get_max_remote_calls_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( disable=True, maximum_remote_calls=20, ), @@ -68,9 +56,9 @@ def test_afc_disabled_max_set(): def test_afc_d_max_unset(): assert ( get_max_remote_calls_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=False, + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True, ), ) ) @@ -81,9 +69,9 @@ def test_afc_d_max_unset(): def test_afc_d_max_set(): assert ( get_max_remote_calls_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=False, + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True, maximum_remote_calls=5, ), ) @@ -95,8 +83,8 @@ def test_afc_d_max_set(): def test_afc_enabled_max_set_to_zero(): with pytest.raises(ValueError): get_max_remote_calls_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( disable=False, maximum_remote_calls=0, ), @@ -107,8 +95,8 @@ def test_afc_enabled_max_set_to_zero(): def test_afc_enabled_max_set_to_negative(): with pytest.raises(ValueError): get_max_remote_calls_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( disable=False, maximum_remote_calls=-1, ), @@ -119,9 +107,9 @@ def test_afc_enabled_max_set_to_negative(): def test_afc_enabled_max_set_to_float(): assert ( get_max_remote_calls_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=False, + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True, maximum_remote_calls=5.0, ), ) diff --git a/google/genai/tests/afc/test_raise_error_for_afc_incompatible_config.py b/google/genai/tests/afc/test_raise_error_for_afc_incompatible_config.py index eadadb2b6..a6eabda58 100644 --- a/google/genai/tests/afc/test_raise_error_for_afc_incompatible_config.py +++ b/google/genai/tests/afc/test_raise_error_for_afc_incompatible_config.py @@ -27,9 +27,9 @@ def test_config_is_none(): def test_tool_config_config_unset(): assert ( - raise_error_for_afc_incompatible_config(types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=False, + raise_error_for_afc_incompatible_config(types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True, maximum_remote_calls=1, ), )) @@ -39,9 +39,9 @@ def test_tool_config_config_unset(): def test_function_calling_config_unset(): assert ( - raise_error_for_afc_incompatible_config(types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=False, + raise_error_for_afc_incompatible_config(types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True, maximum_remote_calls=1, ), tool_config=types.ToolConfig(), @@ -50,14 +50,10 @@ def test_function_calling_config_unset(): ) - def test_compatible_config_afc_disabled(): assert ( raise_error_for_afc_incompatible_config( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=True, - ), + types.ChatConfig( tool_config=types.ToolConfig( function_calling_config=types.FunctionCallingConfig( stream_function_call_arguments=False, @@ -72,7 +68,7 @@ def test_compatible_config_afc_disabled(): def test_compatible_config_stream_function_call_arguments_unset_afc_unset(): assert ( raise_error_for_afc_incompatible_config( - types.GenerateContentConfig( + types.ChatConfig( tool_config=types.ToolConfig( function_calling_config=types.FunctionCallingConfig( ), @@ -83,11 +79,12 @@ def test_compatible_config_stream_function_call_arguments_unset_afc_unset(): ) -def test_compatible_config_stream_function_call_arguments_unset_no_disable_afc(): +def test_compatible_config_stream_function_call_arguments_unset_enable_afc_false(): assert ( raise_error_for_afc_incompatible_config( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=False, ), tool_config=types.ToolConfig( function_calling_config=types.FunctionCallingConfig( @@ -102,9 +99,8 @@ def test_compatible_config_stream_function_call_arguments_unset_no_disable_afc() def test_compatible_config_stream_function_call_arguments_unset_disable_afc_true(): assert ( raise_error_for_afc_incompatible_config( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=True, + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( ), tool_config=types.ToolConfig( function_calling_config=types.FunctionCallingConfig( @@ -119,9 +115,9 @@ def test_compatible_config_stream_function_call_arguments_unset_disable_afc_true def test_incompatible_config_stream_function_call_arguments_set_enable_afc(): with pytest.raises(ValueError): raise_error_for_afc_incompatible_config( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=False, + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True, ), tool_config=types.ToolConfig( function_calling_config=types.FunctionCallingConfig( @@ -133,27 +129,13 @@ def test_incompatible_config_stream_function_call_arguments_set_enable_afc(): def test_incompatible_config_stream_function_call_arguments_set_no_afc_config(): - with pytest.raises(ValueError): - raise_error_for_afc_incompatible_config( - types.GenerateContentConfig( + assert raise_error_for_afc_incompatible_config( + types.ChatConfig( tool_config=types.ToolConfig( function_calling_config=types.FunctionCallingConfig( stream_function_call_arguments=True, ), ), ) - ) - + ) is None -def test_incompatible_config_stream_function_call_arguments_set_no_disable_afc(): - with pytest.raises(ValueError): - raise_error_for_afc_incompatible_config( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig(), - tool_config=types.ToolConfig( - function_calling_config=types.FunctionCallingConfig( - stream_function_call_arguments=True, - ), - ), - ) - ) diff --git a/google/genai/tests/afc/test_should_append_afc_history.py b/google/genai/tests/afc/test_should_append_afc_history.py deleted file mode 100644 index f146f3e28..000000000 --- a/google/genai/tests/afc/test_should_append_afc_history.py +++ /dev/null @@ -1,53 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -"""Tests for _extra_utils.should_append_afc_history.""" - -from ... import types -from ..._extra_utils import should_append_afc_history - - -def test_should_append_afc_history_with_default_config(): - config = types.GenerateContentConfig() - - assert should_append_afc_history(config) == True - - -def test_should_append_afc_history_with_empty_afc_config(): - config = types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig() - ) - - assert should_append_afc_history(config) == True - - -def test_should_append_afc_history_with_ignore_call_history_true(): - config = types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - ignore_call_history=True - ) - ) - - assert should_append_afc_history(config) == False - - -def test_should_append_afc_history_with_ignore_call_history_false(): - config = types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - ignore_call_history=False - ) - ) - - assert should_append_afc_history(config) == True diff --git a/google/genai/tests/afc/test_should_disable_afc.py b/google/genai/tests/afc/test_should_disable_afc.py deleted file mode 100644 index bfc002972..000000000 --- a/google/genai/tests/afc/test_should_disable_afc.py +++ /dev/null @@ -1,222 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - - -"""Tests for should_disable_afc.""" - -import pytest -from .. import pytest_helper -from ... import types -from ..._extra_utils import should_disable_afc - -pytestmark = [ - pytest.mark.skipif( - "config.getoption('--private')", - reason="AFC re-written for private SDK", - ), -] - - -def test_config_is_none(): - assert should_disable_afc(None) is False - - -def test_afc_config_unset(): - assert should_disable_afc(types.GenerateContentConfig()) is False - - -def test_afc_enable_unset_max_0(): - assert ( - should_disable_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - maximum_remote_calls=0, - ), - ) - ) - is True - ) - - -def test_afc_enable_unset_max_negative(): - assert ( - should_disable_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - maximum_remote_calls=-1, - ), - ) - ) - is True - ) - - -def test_afc_enable_unset_max_0_0(): - assert ( - should_disable_afc( - {'automatic_function_calling': {'maximum_remote_calls': 0.0}} - ) - is True - ) - - -def test_afc_enable_unset_max_1(): - assert ( - should_disable_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - maximum_remote_calls=1, - ), - ) - ) - is False - ) - - -def test_afc_enable_unset_max_1_0(): - assert ( - should_disable_afc( - {'automatic_function_calling': {'maximum_remote_calls': 1.0}} - ) - is False - ) - - -def test_afc_enable_false_max_unset(): - assert ( - should_disable_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=True, - ), - ) - ) - is True - ) - - -def test_afc_enable_false_max_0(): - assert ( - should_disable_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=True, - maximum_remote_calls=0, - ), - ) - ) - is True - ) - - -def test_afc_enable_false_max_negative(): - assert ( - should_disable_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=True, - maximum_remote_calls=-1, - ), - ) - ) - is True - ) - - -def test_afc_enable_false_max_0_0(): - assert ( - should_disable_afc( - {'automatic_function_calling': {'maximum_remote_calls': 0.0}} - ) - is True - ) - - -def test_afc_enable_false_max_1(): - assert ( - should_disable_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=True, - maximum_remote_calls=1, - ), - ) - ) - is True - ) - - -def test_afc_enable_true_max_unset(): - assert ( - should_disable_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=False, - ), - ) - ) - is False - ) - - -def test_afc_enable_true_max_0(): - assert ( - should_disable_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=False, - maximum_remote_calls=0, - ), - ) - ) - is True - ) - - -def test_afc_enable_true_max_negative(): - assert ( - should_disable_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=False, - maximum_remote_calls=-1, - ), - ) - ) - is True - ) - - -def test_afc_enable_true_max_0_0(): - assert ( - should_disable_afc( - {'automatic_function_calling': {'maximum_remote_calls': 0.0}} - ) - is True - ) - - -def test_afc_enable_true_max_1(): - assert ( - should_disable_afc( - types.GenerateContentConfig( - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=False, - maximum_remote_calls=1, - ), - ) - ) - is False - ) diff --git a/google/genai/tests/afc/test_should_enable_afc.py b/google/genai/tests/afc/test_should_enable_afc.py new file mode 100644 index 000000000..c6d018c0a --- /dev/null +++ b/google/genai/tests/afc/test_should_enable_afc.py @@ -0,0 +1,158 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + + +"""Tests for should_enable_afc.""" + +import pytest +from .. import pytest_helper +from ... import types +from ..._extra_utils import should_enable_afc + +pytestmark = [ + pytest.mark.skipif( + "config.getoption('--private')", + reason="AFC re-written for private SDK", + ), +] + + +def test_config_is_none(): + assert should_enable_afc(None) is False + + +def test_afc_config_unset(): + assert should_enable_afc(types.ChatConfig()) is False + + +def test_afc_enable_unset_max_0(): + assert ( + should_enable_afc( + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + maximum_remote_calls=0, + ), + ) + ) + is False + ) + + +def test_afc_enable_unset_max_negative(): + assert ( + should_enable_afc( + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + maximum_remote_calls=-1, + ), + ) + ) + is False + ) + + +def test_afc_enable_unset_max_1(): + assert ( + should_enable_afc( + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + maximum_remote_calls=1, + ), + ) + ) + is False + ) + + +def test_afc_enable_false_max_unset(): + assert ( + should_enable_afc( + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True, + ), + ) + ) + is True + ) + + +def test_afc_enable_false_max_0(): + assert ( + should_enable_afc( + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True, + maximum_remote_calls=0, + ), + ) + ) + is False + ) + + +def test_afc_enable_false(): + assert ( + should_enable_afc( + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=False, + ), + ) + ) + is False + ) + + +def test_afc_enable_false_max_1(): + assert ( + should_enable_afc( + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=False, + maximum_remote_calls=1, + ), + ) + ) + is False + ) + + +def test_afc_enable_true_max_negative(): + assert ( + should_enable_afc( + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True, + maximum_remote_calls=-1, + ), + ) + ) + is False + ) + + +def test_afc_enable_true_max_1(): + assert ( + should_enable_afc( + types.ChatConfig( + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True, + maximum_remote_calls=1, + ), + ) + ) + is True + ) diff --git a/google/genai/tests/chats/test_get_history.py b/google/genai/tests/chats/test_get_history.py deleted file mode 100644 index df2ae53d9..000000000 --- a/google/genai/tests/chats/test_get_history.py +++ /dev/null @@ -1,606 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - - -from unittest import mock - -import pytest - -from ... import chats -from ... import client -from ... import models -from ... import types - -AFC_HISTORY = [ - types.Content( - role='user', - parts=[types.Part.from_text(text='afc input')], - ), - types.Content( - role='model', - parts=[ - types.Part( - function_call=types.FunctionCall( - name='foo', args={'bar': 'baz'} - ) - ) - ], - ), -] - - -pytest_plugins = 'pytest_asyncio' - - -@pytest.fixture -def mock_api_client(vertexai=False): - api_client = mock.MagicMock(spec=client.ApiClient) - api_client.api_key = 'TEST_API_KEY' - api_client._host = lambda: 'test_host' - api_client._http_options = {'headers': {}} # Ensure headers exist - api_client.vertexai = vertexai - return api_client - - -@pytest.fixture -def mock_generate_content_with_empty_text_part(): - with mock.patch.object( - models.Models, 'generate_content' - ) as mock_generate_content: - mock_generate_content.return_value = types.GenerateContentResponse( - candidates=[ - types.Candidate( - content=types.Content( - role='model', - parts=[types.Part(text='')], - ) - ) - ] - ) - yield mock_generate_content - - -@pytest.fixture -def mock_generate_content_empty_content(): - with mock.patch.object( - models.Models, 'generate_content' - ) as mock_generate_content: - mock_generate_content.return_value = types.GenerateContentResponse( - candidates=[] - ) - yield mock_generate_content - - -@pytest.fixture -def mock_generate_content_stream_with_empty_text_part(): - with mock.patch.object( - models.Models, 'generate_content_stream' - ) as mock_generate_content: - mock_generate_content.return_value = [ - types.GenerateContentResponse( - candidates=[ - types.Candidate( - content=types.Content( - role='model', - parts=[types.Part(text='')], - ), - finish_reason=types.FinishReason.STOP, - ) - ] - ) - ] - yield mock_generate_content - - -@pytest.fixture -def mock_generate_content_stream_empty_content(): - with mock.patch.object( - models.Models, 'generate_content_stream' - ) as mock_generate_content: - mock_generate_content.return_value = [ - types.GenerateContentResponse(candidates=[]) - ] - yield mock_generate_content - - -@pytest.fixture -def mock_generate_content_afc_history(): - with mock.patch.object( - models.Models, 'generate_content' - ) as mock_generate_content: - mock_generate_content.return_value = types.GenerateContentResponse( - candidates=[ - types.Candidate( - content=types.Content( - role='model', - parts=[types.Part.from_text(text='afc output')], - ) - ) - ], - automatic_function_calling_history=AFC_HISTORY, - ) - yield mock_generate_content - - -@pytest.fixture -def mock_generate_content_stream_afc_history(): - with mock.patch.object( - models.Models, 'generate_content_stream' - ) as mock_generate_content: - mock_generate_content.return_value = [ - types.GenerateContentResponse( - candidates=[ - types.Candidate( - content=types.Content( - role='model', - parts=[types.Part.from_text(text='afc output')], - ), - finish_reason=types.FinishReason.STOP, - ) - ], - automatic_function_calling_history=AFC_HISTORY, - ) - ] - yield mock_generate_content - - -def test_history_start_with_valid_model_content(): - history = [ - types.Content( - role='model', - parts=[types.Part.from_text(text='Hello there! how can I help you?')], - ), - types.Content(role='user', parts=[types.Part.from_text(text='Hello')]), - ] - - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash', history=history) - - assert chat.get_history() == history - assert chat.get_history(curated=True) == history - - -def test_history_start_with_invalid_model_content(): - history = [ - types.Content( - role='model', - parts=[], - ), - types.Content(role='user', parts=[types.Part.from_text(text='Hello')]), - ] - - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash', history=history) - - assert chat.get_history() == history - assert chat.get_history(curated=True) == [types.Content(role='user', parts=[types.Part.from_text(text='Hello')])] - - -def test_history_with_consecutive_valid_user_inputs(): - history = [ - types.Content( - role='user', - parts=[types.Part.from_text(text='user input 1')], - ), - types.Content( - role='user', - parts=[types.Part.from_text(text='user input 2')], - ), - ] - - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash', history=history) - - assert chat.get_history() == history - assert chat.get_history(curated=True) == history - - -def test_history_with_valid_and_invalid_user_inputs(): - history = [ - types.Content( - role='user', - parts=[types.Part.from_text(text='user input 1')], - ), - types.Content( - role='user', - parts=[], # invalid content - ), - types.Content( - role='user', - parts=[types.Part.from_text(text='user input 2')], - ), - ] - - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash', history=history) - - assert chat.get_history() == history - assert chat.get_history(curated=True) == history - - -def test_history_with_consecutive_valid_model_outputs(): - history = [ - types.Content( - role='model', - parts=[types.Part.from_text(text='model output 1')], - ), - types.Content( - role='model', - parts=[types.Part.from_text(text='model output 2')], - ), - ] - - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash', history=history) - - assert chat.get_history() == history - assert chat.get_history(curated=True) == history - - -def test_history_with_valid_and_invalid_model_output(): - history = [ - types.Content( - role='model', - parts=[types.Part.from_text(text='model output 1')], - ), - types.Content( - role='model', - parts=[], # invalid content - ), - types.Content( - role='model', - parts=[types.Part.from_text(text='model output 2')], - ), - ] - - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash', history=history) - - assert chat.get_history() == history - assert chat.get_history(curated=True) == [] - - -def test_history_end_with_user_input(): - history = [ - types.Content( - role='user', - parts=[types.Part.from_text(text='user input 1')], - ), - types.Content( - role='model', - parts=[types.Part.from_text(text='model output')], - ), - types.Content( - role='user', - parts=[types.Part.from_text(text='user input 2')], - ), - ] - - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash', history=history) - - assert chat.get_history() == history - assert chat.get_history(curated=True) == history - - -def test_unrecognized_role_in_history(): - history = [ - types.Content(role='user', parts=[types.Part.from_text(text='Hello')]), - types.Content( - role='invalid_role', - parts=[types.Part.from_text(text='Hello there! how can I help you?')], - ), - ] - - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - with pytest.raises(ValueError) as e: - chat = chats_module.create(model='gemini-2.5-flash', history=history) - - assert 'Role must be user or model' in str(e) - - -def test_sync_chat_create(): - history = [ - types.Content( - role='user', parts=[types.Part.from_text(text='user input turn 1')] - ), - types.Content( - role='model', - parts=[types.Part.from_text(text='model output turn 1')], - ), - types.Content( - role='model', - parts=[types.Part.from_text(text='model output turn 1')], - ), - types.Content( - role='model', - parts=[types.Part.from_text(text='user input turn 2')], - ), - ] - - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash', history=history) - - assert chat.get_history() == history - assert chat.get_history(curated=True) == history - - -def test_async_chat_create(): - history = [ - types.Content( - role='user', parts=[types.Part.from_text(text='user input turn 1')] - ), - types.Content( - role='model', - parts=[types.Part.from_text(text='model output turn 1')], - ), - types.Content( - role='model', - parts=[types.Part.from_text(text='model output turn 1')], - ), - types.Content( - role='model', - parts=[types.Part.from_text(text='user input turn 2')], - ), - types.Content( - role='model', - parts=[types.Part.from_text(text='model output turn 2')], - ), - ] - - models_module = models.AsyncModels(mock_api_client) - chats_module = chats.AsyncChats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash', history=history) - - assert chat.get_history() == history - assert chat.get_history(curated=True) == history - - -def test_sync_chat_create_with_history_dict(): - history = [ - {'role': 'user', 'parts': [{'text': 'user input turn 1'}]}, - {'role': 'model', 'parts': [{'text': 'model output turn 1'}]}, - {'role': 'user', 'parts': [{'text': 'user input turn 2'}]}, - {'role': 'model', 'parts': [{'text': 'model output turn 2'}]}, - ] - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash', history=history) - - expected_history = [ - types.Content( - role='user', parts=[types.Part.from_text(text='user input turn 1')] - ), - types.Content( - role='model', - parts=[types.Part.from_text(text='model output turn 1')], - ), - types.Content( - role='user', parts=[types.Part.from_text(text='user input turn 2')] - ), - types.Content( - role='model', - parts=[types.Part.from_text(text='model output turn 2')], - ), - ] - assert chat.get_history() == expected_history - assert chat.get_history(curated=True) == expected_history - - -def test_async_chat_create_with_history_dict(): - history = [ - {'role': 'user', 'parts': [{'text': 'user input turn 1'}]}, - {'role': 'model', 'parts': [{'text': 'model output turn 1'}]}, - {'role': 'user', 'parts': [{'text': 'user input turn 2'}]}, - {'role': 'model', 'parts': [{'text': 'model output turn 2'}]}, - ] - models_module = models.AsyncModels(mock_api_client) - chats_module = chats.AsyncChats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash', history=history) - - expected_history = [ - types.Content( - role='user', parts=[types.Part.from_text(text='user input turn 1')] - ), - types.Content( - role='model', - parts=[types.Part.from_text(text='model output turn 1')], - ), - types.Content( - role='user', parts=[types.Part.from_text(text='user input turn 2')] - ), - types.Content( - role='model', - parts=[types.Part.from_text(text='model output turn 2')], - ), - ] - assert chat.get_history() == expected_history - assert chat.get_history(curated=True) == expected_history - - -def test_history_with_invalid_turns(): - valid_input = types.Content( - role='user', parts=[types.Part.from_text(text='Hello')] - ) - valid_output = [ - types.Content( - role='model', - parts=[types.Part.from_text(text='Hello there! how can I help you?')], - ), - types.Content( - role='model', - parts=[types.Part.from_text(text='Hello there! how can I help you?')], - ), - ] - invalid_input = types.Content( - role='user', - parts=[ - types.Part.from_text(text='a input will be rejected by the model') - ], - ) - invalid_output = types.Content( - role='model', - parts=[], - ) - comprehensive_history = [] - comprehensive_history.append(valid_input) - comprehensive_history.extend(valid_output) - comprehensive_history.append(invalid_input) - comprehensive_history.append(invalid_output) - curated_history = [] - curated_history.append(valid_input) - curated_history.extend(valid_output) - - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - chat = chats_module.create( - model='gemini-2.5-flash', history=comprehensive_history - ) - - assert chat.get_history() == comprehensive_history - assert chat.get_history(curated=True) == curated_history - - -def test_chat_with_empty_text_part(mock_generate_content_with_empty_text_part): - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash') - - chat.send_message('Hello') - - expected_comprehensive_history = [ - types.UserContent(parts=[types.Part.from_text(text='Hello')]), - types.Content( - parts=[types.Part(text='')], - role='model', - ), - ] - assert chat.get_history() == expected_comprehensive_history - assert chat.get_history(curated=True) == expected_comprehensive_history - - -def test_chat_with_empty_content(mock_generate_content_empty_content): - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash') - - chat.send_message('Hello') - - expected_comprehensive_history = [ - types.UserContent(parts=[types.Part.from_text(text='Hello')]), - types.Content( - parts=[], - role='model', - ), - ] - assert chat.get_history() == expected_comprehensive_history - assert not chat.get_history(curated=True) - - -def test_chat_stream_with_empty_text_part( - mock_generate_content_stream_with_empty_text_part, -): - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash') - - chunks = chat.send_message_stream('Hello') - for chunk in chunks: - pass - - expected_comprehensive_history = [ - types.UserContent(parts=[types.Part.from_text(text='Hello')]), - types.Content( - parts=[types.Part(text='')], - role='model', - ), - ] - assert chat.get_history() == expected_comprehensive_history - assert chat.get_history(curated=True) == expected_comprehensive_history - - -def test_chat_stream_with_empty_content( - mock_generate_content_stream_empty_content, -): - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash') - - chunks = chat.send_message_stream('Hello') - for chunk in chunks: - pass - - expected_comprehensive_history = [ - types.UserContent(parts=[types.Part.from_text(text='Hello')]), - types.Content( - parts=[], - role='model', - ), - ] - assert chat.get_history() == expected_comprehensive_history - assert not chat.get_history(curated=True) - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC logic in private is re-written', -) -def test_chat_with_afc_history(mock_generate_content_afc_history): - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash') - - chat.send_message('Hello') - - expected_history = AFC_HISTORY + [ - types.Content( - role='model', - parts=[types.Part.from_text(text='afc output')], - ), - ] - assert chat.get_history() == expected_history - assert chat.get_history(curated=True) == expected_history - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC logic in private is re-written', -) -def test_chat_stream_with_afc_history(mock_generate_content_stream_afc_history): - models_module = models.Models(mock_api_client) - chats_module = chats.Chats(modules=models_module) - chat = chats_module.create(model='gemini-2.5-flash') - - chunks = chat.send_message_stream('Hello') - for chunk in chunks: - pass - - expected_history = AFC_HISTORY + [ - types.Content( - role='model', - parts=[types.Part.from_text(text='afc output')], - ), - ] - assert chat.get_history() == expected_history - assert chat.get_history(curated=True) == expected_history diff --git a/google/genai/tests/chats/test_send_message.py b/google/genai/tests/chats/test_send_message.py index 0ec0f90b3..a68fd8445 100644 --- a/google/genai/tests/chats/test_send_message.py +++ b/google/genai/tests/chats/test_send_message.py @@ -15,11 +15,15 @@ import json import os +import pydantic + import sys from pydantic import BaseModel from pydantic import ValidationError import pytest +import typing +from typing import Any, Union from .. import pytest_helper from ... import errors @@ -45,32 +49,44 @@ file=__file__, globals_for_file=globals(), ), - pytest.mark.skipif( - "config.getoption('--private')", - reason="AFC re-written for private SDK", - ), ] pytest_plugins = ('pytest_asyncio',) MODEL_NAME = 'gemini-2.5-flash' -def divide_intergers_with_customized_math_rule( - numerator: int, denominator: int -) -> int: - """Divides two integers with customized math rule.""" - return numerator // denominator + 1 + +def get_weather(city: str) -> str: + return f'The weather in {city} is sunny and 100 degrees.' + + +def get_stock_price(symbol: str) -> str: + if symbol == 'GOOG': + return '1000' + else: + return '100' def square_integer(given_integer: int) -> int: return given_integer*given_integer +def divide_floats(numerator: float, denominator: float) -> float: + """Divide two floats.""" + return numerator / denominator + + +def divide_integers(numerator: int, denominator: int) -> int: + """Divide two integers.""" + return numerator // denominator + + def power_disco_ball(power: bool) -> bool: """Powers the spinning disco ball.""" print(f"Disco ball is {'spinning!' if power else 'stopped.'}") return True + def start_music(energetic: bool, loud: bool, bpm: int) -> str: """Play some music matching the specified parameters. @@ -84,6 +100,7 @@ def start_music(energetic: bool, loud: bool, bpm: int) -> str: print(f"Starting music! {energetic=} {loud=}, {bpm=}") return "Never gonna give you up." + def dim_lights(brightness: float) -> bool: """Dim the lights. @@ -93,6 +110,7 @@ def dim_lights(brightness: float) -> bool: print(f"Lights are now set to {brightness:.0%}") return True + def test_text(client): chat = client.chats.create(model=MODEL_NAME) chat.send_message( @@ -131,12 +149,12 @@ def test_thinking_budget(client): """Tests that the thinking budget is respected and generates thoughts.""" chat = client.chats.create( model=MODEL_NAME, - config={ - 'thinking_config': { - 'include_thoughts': True, - 'thinking_budget': 10000, - }, - }, + config=types.ChatConfig( + thinking_config=types.ThinkingConfig( + include_thoughts=True, + thinking_budget=10000, + ), + ), ) response1 = chat.send_message( 'what is the sum of natural numbers from 1 to 100?', @@ -167,12 +185,12 @@ def test_thinking_budget_stream(client): """Tests that the thinking budget is respected and generates thoughts.""" chat = client.chats.create( model=MODEL_NAME, - config={ - 'thinking_config': { - 'include_thoughts': True, - 'thinking_budget': 10000, - }, - }, + config=types.ChatConfig( + thinking_config=types.ThinkingConfig( + include_thoughts=True, + thinking_budget=10000, + ), + ), ) has_thought1 = False for chunk in chat.send_message_stream( @@ -206,7 +224,9 @@ def test_google_cloud_storage_uri(client): [ 'what is the image about?', types.Part.from_uri( - file_uri='gs://unified-genai-dev/imagen-inputs/google_small.png', + file_uri=( + 'gs://unified-genai-dev/imagen-inputs/google_small.png' + ), mime_type='image/png', ), ], @@ -228,9 +248,9 @@ def test_uploaded_file_uri(client): def test_config_override(client): - chat_config = {'candidate_count': 1} + chat_config = types.ChatConfig(candidate_count=1) chat = client.chats.create(model=MODEL_NAME, config=chat_config) - request_config = {'candidate_count': 2} + request_config = types.ChatConfig(candidate_count=2) request_config_response = chat.send_message( 'tell me a story in 100 words', config=request_config) @@ -263,79 +283,22 @@ def test_send_2_messages(client): chat.send_message('write a unit test for the function') -def test_with_afc_history(client): - chat = client.chats.create( - model='gemini-2.0-flash-exp', - config={'tools': [divide_intergers_with_customized_math_rule]}, - ) - _ = chat.send_message('what is the result of 100/2?') - chat_history = chat.get_history() - - assert len(chat_history) == 4 - assert chat_history[0].role == 'user' - assert chat_history[0].parts[0].text == 'what is the result of 100/2?' - - assert chat_history[1].role == 'model' - assert ( - chat_history[1].parts[0].function_call.name - == 'divide_intergers_with_customized_math_rule' - ) - assert chat_history[1].parts[0].function_call.args == { - 'numerator': 100, - 'denominator': 2, - } - - assert chat_history[2].role == 'user' - assert ( - chat_history[2].parts[0].function_response.name - == 'divide_intergers_with_customized_math_rule' - ) - assert chat_history[2].parts[0].function_response.response == {'result': 51} - - assert chat_history[3].role == 'model' - assert '51' in chat_history[3].parts[0].text - - -def test_existing_chat_history_extends_afc_history(client): - chat = client.chats.create( - model='gemini-2.0-flash-exp', - config={'tools': [divide_intergers_with_customized_math_rule]}, - ) - _ = chat.send_message('hello') - _ = chat.send_message('could you help me with a math problem?') - _ = chat.send_message('what is the result of 100/2?') - chat_history = chat.get_history() - content_strings = [] - for content in chat_history: - content_strings.append(content.model_dump_json()) - - # checks that the history is not duplicated - assert len(content_strings) == len(set(content_strings)) - - -@pytest.mark.skipif( - sys.version_info >= (3, 13), - reason=( - 'object type is dumped as as opposed to' - ' "OBJECT" in Python 3.13' - ), -) def test_with_afc_multiple_remote_calls(client): house_fns = [power_disco_ball, start_music, dim_lights] - config = { - 'tools': house_fns, - # Force the model to act (call 'any' function), instead of chatting. - 'tool_config': { - 'function_calling_config': { - 'mode': 'ANY', - } - }, - 'automatic_function_calling': { - 'maximum_remote_calls': 3, - } - } - chat = client.chats.create(model=MODEL_NAME, config=config) + config = types.ChatConfig( + tools=house_fns, + tool_config=types.ToolConfig( + function_calling_config=types.FunctionCallingConfig( + mode=types.FunctionCallingConfigMode.ANY, + ), + ), + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True, + maximum_remote_calls=3, + ), + ) + chat = client.chats.create(model='gemini-3.1-pro-preview', config=config) chat.send_message('Turn this place into a party!') curated_history = chat.get_history() @@ -372,30 +335,24 @@ def test_with_afc_multiple_remote_calls(client): assert part.function_call -@pytest.mark.skipif( - sys.version_info >= (3, 13), - reason=( - 'object type is dumped as as opposed to' - ' "OBJECT" in Python 3.13' - ), -) -def test_with_afc_multiple_remote_calls_async(client): +@pytest.mark.asyncio +async def test_with_afc_multiple_remote_calls_async(client): house_fns = [power_disco_ball, start_music, dim_lights] - config = { - 'tools': house_fns, - # Force the model to act (call 'any' function), instead of chatting. - 'tool_config': { - 'function_calling_config': { - 'mode': 'ANY', - } - }, - 'automatic_function_calling': { - 'maximum_remote_calls': 3, - } - } - chat = client.chats.create(model=MODEL_NAME, config=config) - chat.send_message('Turn this place into a party!') + config = types.ChatConfig( + tools=house_fns, + tool_config=types.ToolConfig( + function_calling_config=types.FunctionCallingConfig( + mode=types.FunctionCallingConfigMode.ANY, + ), + ), + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True, + maximum_remote_calls=3, + ), + ) + chat = client.aio.chats.create(model='gemini-3.1-pro-preview', config=config) + await chat.send_message('Turn this place into a party!') curated_history = chat.get_history() assert len(curated_history) == 8 @@ -430,13 +387,13 @@ def test_with_afc_multiple_remote_calls_async(client): for part in curated_history[7].parts: assert part.function_call + def test_with_afc_disabled(client): chat = client.chats.create( - model='gemini-2.0-flash-exp', - config={ - 'tools': [square_integer], - 'automatic_function_calling': {'disable': True}, - }, + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[square_integer], + ), ) chat.send_message( 'Do the square of 3.', @@ -454,48 +411,13 @@ def test_with_afc_disabled(client): } -@pytest.mark.asyncio -async def test_with_afc_history_async(client): - chat = client.aio.chats.create( - model='gemini-2.0-flash-exp', - config={'tools': [divide_intergers_with_customized_math_rule]}, - ) - _ = await chat.send_message('what is the result of 100/2?') - chat_history = chat.get_history() - - assert len(chat_history) == 4 - assert chat_history[0].role == 'user' - assert chat_history[0].parts[0].text == 'what is the result of 100/2?' - - assert chat_history[1].role == 'model' - assert ( - chat_history[1].parts[0].function_call.name - == 'divide_intergers_with_customized_math_rule' - ) - assert chat_history[1].parts[0].function_call.args == { - 'numerator': 100, - 'denominator': 2, - } - - assert chat_history[2].role == 'user' - assert ( - chat_history[2].parts[0].function_response.name - == 'divide_intergers_with_customized_math_rule' - ) - assert chat_history[2].parts[0].function_response.response == {'result': 51} - - assert chat_history[3].role == 'model' - assert '51' in chat_history[3].parts[0].text - - @pytest.mark.asyncio async def test_with_afc_disabled_async(client): chat = client.aio.chats.create( - model='gemini-2.0-flash-exp', - config={ - 'tools': [square_integer], - 'automatic_function_calling': {'disable': True}, - }, + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[square_integer], + ), ) await chat.send_message( 'Do the square of 3.', @@ -550,9 +472,9 @@ def test_stream_parts(client): def test_stream_config_override(client): - chat_config = {'response_mime_type': 'text/plain'} + chat_config = types.ChatConfig(response_mime_type='text/plain') chat = client.chats.create(model=MODEL_NAME, config=chat_config) - request_config = {'response_mime_type': 'application/json'} + request_config = types.ChatConfig(response_mime_type='application/json') request_config_text = '' for chunk in chat.send_message_stream( 'tell me a story in 100 words', config=request_config @@ -569,32 +491,40 @@ def test_stream_config_override(client): def test_stream_function_calling(client): chat = client.chats.create( - model='gemini-2.0-flash-exp', - config={'tools': [divide_intergers_with_customized_math_rule]}, + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[square_integer], + ), ) - # Now we support AFC. for chunk in chat.send_message_stream( - 'what is the result of 100/2?', + 'do the square of 3', ): pass for chunk in chat.send_message_stream( - 'what is the result of 50/2?', + 'do the square of 4', ): pass chat_history = chat.get_history() + assert len(chat_history) == 6 assert chat_history[0].role == 'user' - assert chat_history[0].parts[0].text == 'what is the result of 100/2?' - + assert chat_history[0].parts[0].text == 'do the square of 3' assert chat_history[1].role == 'model' - assert ( - chat_history[1].parts[0].function_call.name - == 'divide_intergers_with_customized_math_rule' - ) + assert chat_history[1].parts[0].function_call.name == 'square_integer' assert chat_history[1].parts[0].function_call.args == { - 'numerator': 100, - 'denominator': 2, + 'given_integer': 3, } + assert chat_history[2].role == 'model' + assert chat_history[2].parts[0].text == '' + assert chat_history[3].role == 'user' + assert chat_history[3].parts[0].text == 'do the square of 4' + assert chat_history[4].role == 'model' + assert chat_history[4].parts[0].function_call.name == 'square_integer' + assert chat_history[4].parts[0].function_call.args == { + 'given_integer': 4, + } + assert chat_history[5].role == 'model' + assert chat_history[5].parts[0].text == '' def test_stream_send_2_messages(client): @@ -633,9 +563,9 @@ async def test_async_parts(client): @pytest.mark.asyncio async def test_async_config_override(client): - chat_config = {'candidate_count': 1} + chat_config = types.ChatConfig(candidate_count=1) chat = client.aio.chats.create(model=MODEL_NAME, config=chat_config) - request_config = {'candidate_count': 2} + request_config = types.ChatConfig(candidate_count=2) request_config_response = await chat.send_message( 'tell me a story in 100 words', config=request_config) @@ -702,9 +632,9 @@ async def test_async_stream_parts(client): @pytest.mark.asyncio async def test_async_stream_config_override(client): - chat_config = {'response_mime_type': 'text/plain'} + chat_config = types.ChatConfig(response_mime_type='text/plain') chat = client.aio.chats.create(model=MODEL_NAME, config=chat_config) - request_config = {'response_mime_type': 'application/json'} + request_config = types.ChatConfig(response_mime_type='application/json') request_config_text = '' async for chunk in await chat.send_message_stream( 'tell me a story in 100 words', config=request_config @@ -720,33 +650,6 @@ async def test_async_stream_config_override(client): json.loads(default_config_text) -@pytest.mark.asyncio -async def test_async_stream_function_calling(client): - chat = client.aio.chats.create( - model='gemini-2.0-flash-exp', - config={'tools': [divide_intergers_with_customized_math_rule]}, - ) - # Now we support AFC. - async for chunk in await chat.send_message_stream('what is the result of 100/2?'): - pass - async for chunk in await chat.send_message_stream('what is the result of 50/2?'): - pass - chat_history = chat.get_history() - - assert chat_history[0].role == 'user' - assert chat_history[0].parts[0].text == 'what is the result of 100/2?' - - assert chat_history[1].role == 'model' - assert ( - chat_history[1].parts[0].function_call.name - == 'divide_intergers_with_customized_math_rule' - ) - assert chat_history[1].parts[0].function_call.args == { - 'numerator': 100, - 'denominator': 2, - } - - @pytest.mark.asyncio async def test_async_stream_send_2_messages(client): chat = client.aio.chats.create(model=MODEL_NAME) @@ -762,8 +665,9 @@ async def test_async_stream_send_2_messages(client): def test_mcp_tools(client): chat = client.chats.create( - model='gemini-2.0-flash-exp', - config={'tools': [ + model='gemini-3.1-pro-preview', + config = types.ChatConfig( + tools=[ mcp_types.Tool( name='get_weather', description='Get the weather in a city.', @@ -772,7 +676,8 @@ def test_mcp_tools(client): 'properties': {'location': {'type': 'string'}}, }, ) - ],}, + ], + ), ) response = chat.send_message('What is the weather in Boston?') response = chat.send_message('What is the weather in San Francisco?') @@ -780,18 +685,19 @@ def test_mcp_tools(client): def test_mcp_tools_stream(client): chat = client.chats.create( - model='gemini-2.0-flash-exp', - config={'tools': [ - mcp_types.Tool( - name='get_weather', - description='Get the weather in a city.', - inputSchema={ - 'type': 'object', - 'properties': {'location': {'type': 'string'}}, - }, - ) - ], - }, + model='gemini-3.1-pro-preview', + config = types.ChatConfig( + tools=[ + mcp_types.Tool( + name='get_weather', + description='Get the weather in a city.', + inputSchema={ + 'type': 'object', + 'properties': {'location': {'type': 'string'}}, + }, + ) + ], + ), ) for chunk in chat.send_message_stream( 'What is the weather in Boston?' @@ -806,17 +712,19 @@ def test_mcp_tools_stream(client): @pytest.mark.asyncio async def test_async_mcp_tools(client): chat = client.aio.chats.create( - model='gemini-2.0-flash-exp', - config={'tools': [ - mcp_types.Tool( - name='get_weather', - description='Get the weather in a city.', - inputSchema={ - 'type': 'object', - 'properties': {'location': {'type': 'string'}}, - }, - ) - ],}, + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + mcp_types.Tool( + name='get_weather', + description='Get the weather in a city.', + inputSchema={ + 'type': 'object', + 'properties': {'location': {'type': 'string'}}, + }, + ) + ], + ), ) await chat.send_message('What is the weather in Boston?'); await chat.send_message('What is the weather in San Francisco?'); @@ -825,18 +733,19 @@ async def test_async_mcp_tools(client): @pytest.mark.asyncio async def test_async_mcp_tools_stream(client): chat = client.aio.chats.create( - model='gemini-2.0-flash-exp', - config={'tools': [ - mcp_types.Tool( - name='get_weather', - description='Get the weather in a city.', - inputSchema={ - 'type': 'object', - 'properties': {'location': {'type': 'string'}}, - }, - ) - ], - }, + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + mcp_types.Tool( + name='get_weather', + description='Get the weather in a city.', + inputSchema={ + 'type': 'object', + 'properties': {'location': {'type': 'string'}}, + }, + ) + ], + ), ) async for chunk in await chat.send_message_stream( @@ -853,8 +762,8 @@ def test_server_side_mcp_tools(client): with pytest_helper.exception_if_vertex(client, ValueError): chat = client.chats.create( model='gemini-2.5-flash', - config={ - 'tools': [ + config=types.ChatConfig( + tools=[ { 'mcp_servers': [ { @@ -872,7 +781,7 @@ def test_server_side_mcp_tools(client): ], }, ], - }, + ), ) response = chat.send_message('What is the weather in Boston on 02/02/2026?') response = chat.send_message( @@ -884,8 +793,8 @@ def test_server_side_mcp_tools_stream(client): with pytest_helper.exception_if_vertex(client, ValueError): chat = client.chats.create( model='gemini-2.5-flash', - config={ - 'tools': [ + config=types.ChatConfig( + tools=[ { 'mcp_servers': [ { @@ -903,7 +812,7 @@ def test_server_side_mcp_tools_stream(client): ], }, ], - }, + ), ) for chunk in chat.send_message_stream( 'What is the weather in Boston on 02/02/2026?' @@ -920,8 +829,8 @@ async def test_async_server_side_mcp_tools(client): with pytest_helper.exception_if_vertex(client, ValueError): chat = client.aio.chats.create( model='gemini-2.5-flash', - config={ - 'tools': [ + config=types.ChatConfig( + tools=[ { 'mcp_servers': [ { @@ -939,9 +848,728 @@ async def test_async_server_side_mcp_tools(client): ], }, ], - }, + ), ) await chat.send_message('What is the weather in Boston on 02/02/2026?') await chat.send_message( 'What is the weather in San Francisco on 02/02/2026?' ) + + +def test_function_tool_afc(client): + chat = client.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + get_weather, + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ), + ) + chat.send_message('What is the weather in Boston?') + history = chat.get_history() + assert len(history) == 4 + assert history[0].role == 'user' + assert history[0].parts[0].text == 'What is the weather in Boston?' + assert history[1].role == 'model' + assert history[1].parts[0].function_call.name == 'get_weather' + assert history[1].parts[0].function_call.args == {'city': 'Boston'} + assert history[2].role == 'user' + assert history[2].parts[0].function_response.name == 'get_weather' + assert history[3].role == 'model' + assert 'sunny' in history[3].parts[0].text.lower() + + +def test_function_tool_multi_turn_afc(client): + chat = client.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + get_weather, + get_stock_price, + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ), + ) + chat.send_message('What is the weather in Boston?') + history = chat.get_history() + assert len(history) == 4 + assert history[0].role == 'user' + assert history[0].parts[0].text == 'What is the weather in Boston?' + assert history[1].role == 'model' + assert history[1].parts[0].function_call.name == 'get_weather' + assert history[1].parts[0].function_call.args == {'city': 'Boston'} + assert history[2].role == 'user' + assert history[2].parts[0].function_response.name == 'get_weather' + assert history[3].role == 'model' + assert 'sunny' in history[3].parts[0].text.lower() + + chat.send_message('What is the stock price of symbol GOOG?') + history = chat.get_history() + assert len(history) == 8 + assert history[0].role == 'user' + assert history[0].parts[0].text == 'What is the weather in Boston?' + assert history[1].role == 'model' + assert history[1].parts[0].function_call.name == 'get_weather' + assert history[1].parts[0].function_call.args == {'city': 'Boston'} + assert history[2].role == 'user' + assert history[2].parts[0].function_response.name == 'get_weather' + assert history[3].role == 'model' + assert 'sunny' in history[3].parts[0].text.lower() + assert history[4].role == 'user' + assert history[4].parts[0].text == 'What is the stock price of symbol GOOG?' + assert history[5].role == 'model' + assert history[5].parts[0].function_call.name == 'get_stock_price' + assert history[5].parts[0].function_call.args == {'symbol': 'GOOG'} + assert history[6].role == 'user' + assert history[6].parts[0].function_response.name == 'get_stock_price' + assert history[7].role == 'model' + assert '1000' in history[7].parts[0].text + + +def test_multi_turn_afc_enabled_FC_FR_parts(client): + chat = client.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + get_weather, + get_stock_price, + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ), + history=[ + types.Content( + role='user', + parts=[types.Part(text='What is the weather in Boston?')], + ), + types.Content( + role='model', + parts=[ + types.Part( + function_call=types.FunctionCall( + name='get_weather', + args={'city': 'Boston'}, + ), + ), + types.Part( + function_response=types.FunctionResponse( + name='get_weather', + response={'weather': 'sunny and 80 degrees'}, + ), + ), + types.Part(text='The weather is sunny.'), + ], + ), + ] + ) + with pytest_helper.exception_if_vertex(client, errors.ClientError): + chat.send_message('What is the stock price of symbol GOOG?') + + +@pytest.mark.asyncio +async def test_async_function_tool_afc_disabled(client): + chat = client.aio.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + get_weather, + ], + ), + ) + await chat.send_message('What is the weather in Boston?') + history = chat.get_history() + assert len(history) == 2 + assert history[0].role == 'user' + assert history[0].parts[0].text == 'What is the weather in Boston?' + assert history[1].role == 'model' + assert len(history[1].parts) == 1 + assert history[1].parts[0].function_call.name == 'get_weather' + assert history[1].parts[0].function_call.args == {'city': 'Boston'} + + +@pytest.mark.asyncio +async def test_async_function_tool_afc_enabled(client): + chat = client.aio.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + get_weather, + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ), + ) + await chat.send_message('What is the weather in Boston?') + history = chat.get_history() + assert len(history) == 4 + assert history[0].role == 'user' + assert history[0].parts[0].text == 'What is the weather in Boston?' + assert history[1].role == 'model' + assert history[1].parts[0].function_call.name == 'get_weather' + assert history[1].parts[0].function_call.args == {'city': 'Boston'} + assert history[2].role == 'user' + assert history[2].parts[0].function_response.name == 'get_weather' + assert history[3].role == 'model' + assert 'sunny' in history[3].parts[0].text.lower() + + +@pytest.mark.asyncio +async def test_async_function_tool_afc_enabled_multi_turn(client): + chat = client.aio.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + get_weather, + get_stock_price, + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ), + ) + await chat.send_message('What is the weather in Boston?') + history = chat.get_history() + assert len(history) == 4 + assert history[0].role == 'user' + assert history[0].parts[0].text == 'What is the weather in Boston?' + assert history[1].role == 'model' + assert history[1].parts[0].function_call.name == 'get_weather' + assert history[1].parts[0].function_call.args == {'city': 'Boston'} + assert history[2].role == 'user' + assert history[2].parts[0].function_response.name == 'get_weather' + assert history[3].role == 'model' + assert 'sunny' in history[3].parts[0].text.lower() + + await chat.send_message('What is the stock price of symbol GOOG?') + history = chat.get_history() + assert len(history) == 8 + assert history[0].role == 'user' + assert history[0].parts[0].text == 'What is the weather in Boston?' + assert history[1].role == 'model' + assert history[1].parts[0].function_call.name == 'get_weather' + assert history[1].parts[0].function_call.args == {'city': 'Boston'} + assert history[2].role == 'user' + assert history[2].parts[0].function_response.name == 'get_weather' + assert history[3].role == 'model' + assert 'sunny' in history[3].parts[0].text.lower() + assert history[4].role == 'user' + assert history[4].parts[0].text == 'What is the stock price of symbol GOOG?' + assert history[5].role == 'model' + assert history[5].parts[0].function_call.name == 'get_stock_price' + assert history[5].parts[0].function_call.args == {'symbol': 'GOOG'} + assert history[6].role == 'user' + assert history[6].parts[0].function_response.name == 'get_stock_price' + assert history[7].role == 'model' + assert '1000' in history[7].parts[0].text + + +def test_stream_function_tool_afc_disabled(client): + chat = client.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + get_weather, + ], + ), + ) + for chunk in chat.send_message_stream('What is the weather in Boston?'): + pass + history = chat.get_history() + assert len(history) == 3 + assert history[0].role == 'user' + assert history[1].role == 'model' + assert history[1].parts[0].function_call.name == 'get_weather' + assert history[2].role == 'model' + assert history[2].parts[0].text == '' + + +def test_stream_function_tool_afc_enabled(client): + chat = client.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + get_weather, + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ), + ) + for chunk in chat.send_message_stream('What is the weather in Boston?'): + pass + history = chat.get_history() + assert len(history) == 6 + assert history[0].role == 'user' + assert history[1].role == 'model' + assert history[1].parts[0].function_call.name == 'get_weather' + assert history[2].role == 'model' + assert history[2].parts[0].text == '' + assert history[3].role == 'user' + assert history[3].parts[0].function_response.name == 'get_weather' + assert history[4].role == 'model' + assert 'Boston' in history[4].parts[0].text + assert history[5].role == 'model' + assert history[5].parts[0].text == '' + + +def test_stream_function_tool_afc_enabled_multi_turn(client): + chat = client.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + get_weather, + get_stock_price, + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ), + ) + for chunk in chat.send_message_stream('What is the weather in Boston?'): + pass + history = chat.get_history() + + assert len(history) == 6 + assert history[0].role == 'user' + assert history[1].role == 'model' + assert history[1].parts[0].function_call.name == 'get_weather' + assert history[2].role == 'model' + assert history[2].parts[0].text == '' + assert history[3].role == 'user' + assert history[3].parts[0].function_response.name == 'get_weather' + assert history[4].role == 'model' + assert 'Boston' in history[4].parts[0].text + assert history[5].role == 'model' + + for chunk in chat.send_message_stream('What is the stock price of symbol GOOG?'): + pass + history = chat.get_history() + + assert len(history) == 13 + assert history[6].role == 'user' + assert history[7].role == 'model' + assert history[7].parts[0].function_call.name == 'get_stock_price' + assert history[8].role == 'model' + assert history[8].parts[0].text == '' + assert history[9].role == 'user' + assert history[9].parts[0].function_response.name == 'get_stock_price' + assert history[10].role == 'model' + assert 'stock' in history[10].parts[0].text + assert history[11].role == 'model' + assert '1000' in history[11].parts[0].text + assert history[12].role == 'model' + assert history[12].parts[0].text == '' + + +@pytest.mark.asyncio +async def test_async_stream_function_tool_afc_disabled(client): + chat = client.aio.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + get_weather, + ], + ), + ) + async for chunk in await chat.send_message_stream( + 'What is the weather in Boston?' + ): + pass + history = chat.get_history() + assert len(history) == 3 + assert history[0].role == 'user' + assert history[1].role == 'model' + assert history[1].parts[0].function_call.name == 'get_weather' + assert history[2].role == 'model' + assert history[2].parts[0].text == '' + + +@pytest.mark.asyncio +async def test_async_stream_function_tool_afc_enabled(client): + chat = client.aio.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + get_weather, + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ), + ) + async for _ in await chat.send_message_stream( + 'What is the weather in Boston?' + ): + pass + + history = chat.get_history() + assert len(history) == 6 + assert history[0].role == 'user' + assert history[1].role == 'model' + assert history[1].parts[0].function_call.name == 'get_weather' + assert history[2].role == 'model' + assert history[2].parts[0].text == '' + assert history[3].role == 'user' + assert history[3].parts[0].function_response.name == 'get_weather' + assert history[4].role == 'model' + assert 'Boston' in history[4].parts[0].text + assert history[5].parts[0].text == '' + + +@pytest.mark.asyncio +async def test_async_stream_function_tool_afc_enabled_multi_turn( + client, +): + chat = client.aio.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + get_weather, + get_stock_price, + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ), + ) + async for _ in await chat.send_message_stream( + 'What is the weather in Boston?' + ): + pass + history = chat.get_history() + + assert len(history) == 6 + assert history[0].role == 'user' + assert history[1].role == 'model' + assert history[1].parts[0].function_call.name == 'get_weather' + assert history[2].role == 'model' + assert history[2].parts[0].text == '' + assert history[3].role == 'user' + assert history[3].parts[0].function_response.name == 'get_weather' + assert history[4].role == 'model' + assert 'Boston' in history[4].parts[0].text + assert history[5].role == 'model' + assert history[5].parts[0].text == '' + + async for _ in await chat.send_message_stream( + 'What is the stock price of symbol GOOG?' + ): + pass + history = chat.get_history() + + assert len(history) == 13 + assert history[6].role == 'user' + assert history[7].role == 'model' + assert history[7].parts[0].function_call.name == 'get_stock_price' + assert history[8].role == 'model' + assert history[8].parts[0].text == '' + assert history[9].role == 'user' + assert history[9].parts[0].function_response.name == 'get_stock_price' + assert history[10].role == 'model' + assert 'stock' in history[10].parts[0].text + assert history[11].role == 'model' + assert '1000' in history[11].parts[0].text + assert history[12].role == 'model' + assert history[12].parts[0].text == '' + + +def test_union_type_afc(client): + + def add_numbers( + a: Union[int, float], b: Union[int, float] + ) -> Union[int, float]: + """add two numbers.""" + + return a + b + + chat = client.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + add_numbers, + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ), + ) + chat.send_message('add 1 and 2.5') + history = chat.get_history() + assert len(history) == 4 + assert history[0].role == 'user' + assert history[0].parts[0].text == 'add 1 and 2.5' + assert history[1].role == 'model' + assert history[1].parts[0].function_call.name == 'add_numbers' + assert history[2].role == 'user' + assert history[2].parts[0].function_response.name == 'add_numbers' + assert history[3].role == 'model' + assert '3.5' in history[3].parts[0].text + + +@pytest.mark.asyncio +async def test_mcp_sync_call_async_afc(client): + class MockMcpClientSession(McpClientSession): + + def __init__(self): + self._read_stream = None + self._write_stream = None + + async def list_tools(self): + return mcp_types.ListToolsResult( + tools=[ + mcp_types.Tool( + name='get_weather', + description='Get the weather in a city.', + inputSchema={ + 'type': 'object', + 'properties': {'location': {'type': 'string'}}, + }, + ), + mcp_types.Tool( + name='add_numbers', + description='Add two numbers together.', + inputSchema={ + 'type': 'object', + 'properties': { + 'a': {'type': 'number'}, + 'b': {'type': 'number'}, + }, + }, + ), + ] + ) + + async def call_tool( + self, + name: str, + arguments: dict[str, Any], + ): + if name == 'get_weather': + return mcp_types.CallToolResult( + content=[mcp_types.TextContent(type='text', text='Sunny')] + ) + else: + return mcp_types.CallToolResult( + content=[mcp_types.TextContent(type='text', text='100')] + ) + + config = types.ChatConfig( + tools=[ + MockMcpClientSession(), + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ) + chat = client.aio.chats.create( + model='gemini-3.1-pro-preview', + config=config, + ) + + response = await chat.send_message( + 'What is the weather in Boston?' + ) + assert 'sunny' in response.text.lower() + + response_2 = await chat.send_message( + 'What is 50 + 50?' + ) + assert '100' in response_2.text + + +def test_afc_float_without_decimal(client): + chat = client.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + divide_floats, + divide_integers, + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ), + ) + response = chat.send_message('what is the result of 10.0/2?') + assert '5.0' in response.text + + +def test_afc_pydantic_model(client): + class CityObject(pydantic.BaseModel): + city_name: str + + def get_weather_pydantic_model( + city_object: CityObject, is_winter: bool + ) -> str: + if is_winter: + return f'The weather in {city_object.city_name} is cold and 10 degrees.' + else: + return f'The weather in {city_object.city_name} is warm and 25 degrees.' + + chat = client.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + get_weather_pydantic_model, + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ), + ) + response = chat.send_message( + 'it is winter now, what is the weather in Boston?' + ) + assert 'cold' in response.text and 'Boston' in response.text + + +def test_afc_pydantic_model_list_type(client): + class CityObject(pydantic.BaseModel): + city_name: str + + def get_weather_from_list_of_cities( + city_object_list: list[CityObject], + is_winter: bool, + ) -> str: + result = '' + if is_winter: + for city_object in city_object_list: + result += ( + f'The weather in {city_object.city_name} is cold and 10 degrees.\n' + ) + else: + for city_object in city_object_list: + result += ( + f'The weather in {city_object.city_name} is warm and 100 degrees.\n' + ) + return result + + chat = client.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + get_weather_from_list_of_cities, + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ), + ) + response = chat.send_message( + 'it is winter now, what is the weather in Boston and New York?' + ) + assert 'cold' in response.text + assert 'Boston' in response.text + assert 'New York' in response.text + + +@pytest.mark.skip( + 'kokoro pydantic version is too low, can unskip after increasing the' + ' pydantic version to 2.13.4 in requirements.txt later' +) +def test_afc_pydantic_model_union_type(client): + + class AnimalObject(pydantic.BaseModel): + name: str + age: int + species: str + + class PlantObject(pydantic.BaseModel): + name: str + height: float + color: str + + def get_information( + object_of_interest: Union[AnimalObject, PlantObject], + ) -> str: + if isinstance(object_of_interest, AnimalObject): + return ( + f'The animal is of {object_of_interest.species} species and is named' + f' {object_of_interest.name} is {object_of_interest.age} years old' + ) + elif isinstance(object_of_interest, PlantObject): + return ( + f'The plant is named {object_of_interest.name} and is' + f' {object_of_interest.height} meters tall and is' + f' {object_of_interest.color} color' + ) + else: + return 'The animal is not supported' + + chat = client.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + get_information, + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ), + ) + response = chat.send_message( + 'I have a one year old cat named Sundae, can you get the' + ' information of the cat for me?' + ) + assert 'Sundae' in response.text + assert 'cat' in response.text + + +def test_afc_with_coroutine_function(client): + + async def divide_integers_async(a: int, b: int) -> int: + return a // b + + chat = client.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + divide_integers_async, + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ), + ) + + with pytest.raises(errors.UnsupportedFunctionError): + chat.send_message('Divide 1000 by 2.') + + +def test_class_method_afc(client): + + class FunctionHolder: + NAME = 'FunctionHolder' + def is_a_duck(self, number: int) -> str: + return self.NAME + 'says isOdd: ' + str(number % 2 == 1) + def is_a_rabbit(self, number: int) -> str: + return self.NAME + 'says isEven: ' + str(number % 2 == 0) + + function_holder = FunctionHolder() + + chat = client.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( + tools=[ + function_holder.is_a_duck, + function_holder.is_a_rabbit, + ], + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True + ), + ), + ) + + response = chat.send_message( + 'Print the vertatim output of is_a_duck and is_a_rabbit for the number' + ' of 100' + ) + + assert 'functionholder' in response.text.lower() diff --git a/google/genai/tests/models/test_function_call_streaming.py b/google/genai/tests/chats/test_send_message_stream.py similarity index 68% rename from google/genai/tests/models/test_function_call_streaming.py rename to google/genai/tests/chats/test_send_message_stream.py index fbc3ec914..dded02068 100644 --- a/google/genai/tests/models/test_function_call_streaming.py +++ b/google/genai/tests/chats/test_send_message_stream.py @@ -13,13 +13,14 @@ # limitations under the License. # -"""Tests for models.generate_content_stream() with stream_function_call_arguments enabled.""" +"""Tests for chats.send_message_stream() with stream_function_call_arguments enabled.""" import pytest from ... import types from unittest import mock from .. import pytest_helper -from . import test_generate_content_tools +from . import test_send_message + json_function_declarations = [{ 'name': 'get_current_weather', @@ -79,26 +80,16 @@ 'purpose': { 'type': 'STRING', 'description': 'Discribes the purpose of asking the weather', - } + }, }, 'required': ['location', 'unit', 'country'], }, }] -generate_content_prompt = [ - types.Content( - role='user', - parts=[ - types.Part( - text=( - 'get the current weather in boston in celsius, the' - ' country should be US, the purpose is to know' - ' what to wear today?' - ) - ) - ], - ), -] +generate_content_prompt = ( + 'get the current weather in boston in celsius, the country should be US,' + ' the purpose is to know what to wear today?' +) previous_generate_content_history = [ types.Content( role='user', @@ -149,63 +140,55 @@ ) -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='in private it was not able to find the replay file', -) def test_streaming_with_python_native_no_afc_config(client): """Tests streaming function calls with native python AFC without disabling AFC.""" if not client.vertexai: return - with pytest.raises(ValueError) as e: - for chunk in client.models.generate_content_stream( - model='gemini-3-pro-preview', - contents=generate_content_prompt, - config=types.GenerateContentConfig( - tools=[ - test_generate_content_tools.get_weather, - test_generate_content_tools.get_stock_price, - ], - tool_config=types.ToolConfig( - function_calling_config={ - 'stream_function_call_arguments': True, - } - ), - ), - ): - pass + chat = client.chats.create( + model='gemini-3-pro-preview', + config=types.ChatConfig( + tools=[ + test_send_message.get_weather, + test_send_message.get_stock_price, + ], + ), + ) + for _ in chat.send_message_stream( + generate_content_prompt, + ): + pass + history = chat.get_history() + assert len(history) == 2 + assert history[0].role == 'user' + assert history[1].role == 'model' + assert history[1].parts[0].function_call.name == 'get_weather' - assert 'not compatible with automatic function calling (AFC)' in str(e.value) -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='in private it was not able to find the replay file', -) -def test_streaming_with_python_afc_disabled_false(client): +def test_streaming_with_python_afc_enabled(client): """Tests streaming function calls with native python AFC without disabling AFC.""" if not client.vertexai: return with pytest.raises(ValueError) as e: - for chunk in client.models.generate_content_stream( + chat = client.chats.create( model='gemini-3-pro-preview', - contents=( - 'What is the price of GOOG? And what is the weather in Boston?' - ), - config=types.GenerateContentConfig( + config=types.ChatConfig( tools=[ - test_generate_content_tools.get_weather, - test_generate_content_tools.get_stock_price, + test_send_message.get_weather, + test_send_message.get_stock_price, ], - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=False, - ), tool_config=types.ToolConfig( function_calling_config={ 'stream_function_call_arguments': True, } ), + automatic_function_calling_config=types.AutomaticFunctionCallingConfig( + enable=True, + ), ), + ) + for _ in chat.send_message_stream( + 'What is the price of GOOG? And what is the weather in Boston?' ): pass assert 'not compatible with automatic function calling (AFC)' in str(e.value) @@ -215,10 +198,9 @@ def test_streaming_with_json_parameters_without_history(client): """Tests streaming function calls with FunctionDeclaration withJSON parameters.""" with pytest_helper.exception_if_mldev(client, ValueError): - for chunk in client.models.generate_content_stream( - model='gemini-3-pro-preview', - contents=generate_content_prompt, - config=types.GenerateContentConfig( + chat = client.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( tools=[{'function_declarations': json_function_declarations}], tool_config=types.ToolConfig( function_calling_config={ @@ -226,6 +208,9 @@ def test_streaming_with_json_parameters_without_history(client): } ), ), + ) + for chunk in chat.send_message_stream( + generate_content_prompt, ): assert chunk is not None assert chunk.candidates is not None @@ -237,10 +222,9 @@ def test_streaming_with_json_parameters_without_history(client): async def test_streaming_with_json_parameters_async(client): """Tests streaming function calls with FunctionDeclaration withJSON parameters.""" with pytest_helper.exception_if_mldev(client, ValueError): - async for chunk in await client.aio.models.generate_content_stream( - model='gemini-3-pro-preview', - contents=generate_content_prompt, - config=types.GenerateContentConfig( + chat = client.aio.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( tools=[{'function_declarations': json_function_declarations}], tool_config=types.ToolConfig( function_calling_config={ @@ -248,6 +232,9 @@ async def test_streaming_with_json_parameters_async(client): } ), ), + ) + async for chunk in await chat.send_message_stream( + generate_content_prompt, ): assert chunk is not None assert chunk.candidates is not None @@ -258,10 +245,9 @@ async def test_streaming_with_json_parameters_async(client): def test_streaming_with_gemini_parameters_without_history(client): """Tests streaming function calls with FunctionDeclaration withJSON parameters.""" with pytest_helper.exception_if_mldev(client, ValueError): - for chunk in client.models.generate_content_stream( - model='gemini-3-pro-preview', - contents=generate_content_prompt, - config=types.GenerateContentConfig( + chat = client.chats.create( + model='gemini-3.1-pro-preview', + config=types.ChatConfig( tools=[{ 'function_declarations': gemini_function_declarations }], @@ -271,73 +257,15 @@ def test_streaming_with_gemini_parameters_without_history(client): } ), ), + ) + for chunk in chat.send_message_stream( + generate_content_prompt, ): assert chunk is not None assert chunk.candidates is not None assert chunk.candidates[0].content is not None assert chunk.candidates[0].content.parts is not None -def test_streaming_with_gemini_parameters_with_response(client): - """Tests streaming function calls with FunctionDeclaration withJSON parameters.""" - with pytest_helper.exception_if_mldev(client, ValueError): - streaming_function_call_content = [] - for chunk in client.models.generate_content_stream( - model='gemini-3-pro-preview', - contents=[ - types.Content( - role='user', - parts=[ - types.Part( - text=( - 'get the current weather in boston in celsius, the' - ' country should be US, the purpose is to know' - ' what to wear today?' - ) - ) - ], - ), - ], - config=types.GenerateContentConfig( - tools=[{ - 'function_declarations': gemini_function_declarations - }], - tool_config=types.ToolConfig( - function_calling_config={ - 'stream_function_call_arguments': True, - } - ), - ), - ): - streaming_function_call_content.append(chunk.candidates[0].content) - - streaming_function_call_content.append( - types.Content( - role='user', - parts=[ - types.Part.from_function_response( - name='get_current_weather', - response={ - 'temperature': 21, - 'unit': 'C', - }, - ) - ], - ), - ) - - for chunk in client.models.generate_content_stream( - model='gemini-3-pro-preview', - contents=streaming_function_call_content, - config=types.GenerateContentConfig( - tools=[{'function_declarations': json_function_declarations}], - tool_config=types.ToolConfig( - function_calling_config={ - 'stream_function_call_arguments': True, - } - ), - ), - ): - pass def test_chat_streaming_with_json_parameters_with_history(client): """Tests streaming function calls with FunctionDeclaration withJSON parameters.""" @@ -373,9 +301,9 @@ def test_chat_streaming_with_json_parameters_with_history(client): ), ] chat = client.chats.create( - model='gemini-3-pro-preview', + model='gemini-3.1-pro-preview', history=previous_generate_content_history, - config=types.GenerateContentConfig( + config=types.ChatConfig( tools=[{ 'function_declarations': gemini_function_declarations }], @@ -434,7 +362,7 @@ async def test_chat_streaming_with_json_parameters_with_history_async(client): chat = client.aio.chats.create( model='gemini-3-pro-preview', history=previous_generate_content_history, - config=types.GenerateContentConfig( + config=types.ChatConfig( tools=[{'function_declarations': gemini_function_declarations}], tool_config=types.ToolConfig( function_calling_config={ @@ -448,4 +376,4 @@ async def test_chat_streaming_with_json_parameters_with_history_async(client): assert chunk is not None assert chunk.candidates is not None assert chunk.candidates[0].content is not None - assert chunk.candidates[0].content.parts is not None + assert chunk.candidates[0].content.parts is not None \ No newline at end of file diff --git a/google/genai/tests/live/test_live.py b/google/genai/tests/live/test_live.py index 046b08a0a..d298af0fc 100644 --- a/google/genai/tests/live/test_live.py +++ b/google/genai/tests/live/test_live.py @@ -1203,6 +1203,10 @@ async def test_bidi_setup_to_api_with_tools_function_behavior(vertexai): ) +@pytest.mark.skipif( + 'config.getoption("--private")', + reason='private serialized into camelCase but public keeps snake_case', +) @pytest.mark.parametrize('vertexai', [True, False]) @pytest.mark.asyncio async def test_bidi_setup_to_api_with_config_mcp_tools( @@ -1216,11 +1220,11 @@ async def test_bidi_setup_to_api_with_config_mcp_tools( 'model': 'models/test_model', 'tools': [{ 'functionDeclarations': [{ - 'parameters': { - 'type': 'OBJECT', + 'parameters_json_schema': { + 'type': 'object', 'properties': { 'location': { - 'type': 'STRING', + 'type': 'string', }, }, }, @@ -1242,11 +1246,11 @@ async def test_bidi_setup_to_api_with_config_mcp_tools( ), 'tools': [{ 'functionDeclarations': [{ - 'parameters': { - 'type': 'OBJECT', + 'parameters_json_schema': { + 'type': 'object', 'properties': { 'location': { - 'type': 'STRING', + 'type': 'string', }, }, }, @@ -1280,6 +1284,10 @@ async def test_bidi_setup_to_api_with_config_mcp_tools( ) +@pytest.mark.skipif( + 'config.getoption("--private")', + reason='private serialized into camelCase but public keeps snake_case', +) @pytest.mark.parametrize('vertexai', [True, False]) @pytest.mark.asyncio async def test_bidi_setup_to_api_with_config_mcp_session( @@ -1313,11 +1321,11 @@ async def list_tools(self): 'model': 'models/test_model', 'tools': [{ 'functionDeclarations': [{ - 'parameters': { - 'type': 'OBJECT', + 'parameters_json_schema': { + 'type': 'object', 'properties': { 'location': { - 'type': 'STRING', + 'type': 'string', }, }, }, @@ -1339,11 +1347,11 @@ async def list_tools(self): ), 'tools': [{ 'functionDeclarations': [{ - 'parameters': { - 'type': 'OBJECT', + 'parameters_json_schema': { + 'type': 'object', 'properties': { 'location': { - 'type': 'STRING', + 'type': 'string', }, }, }, @@ -1361,6 +1369,7 @@ async def list_tools(self): }, ) + assert ( result == expected_result_vertexai if vertexai diff --git a/google/genai/tests/mcp/test_mcp_to_gemini_tools.py b/google/genai/tests/mcp/test_mcp_to_gemini_tools.py index 61d64102c..38df1eb02 100644 --- a/google/genai/tests/mcp/test_mcp_to_gemini_tools.py +++ b/google/genai/tests/mcp/test_mcp_to_gemini_tools.py @@ -65,10 +65,12 @@ def test_unknown_field_conversion(): types.FunctionDeclaration( name='tool', description='tool-description', - parameters=types.Schema( - type='OBJECT', - properties={}, - ), + parameters_json_schema={ + 'type': 'object', + 'properties': {}, + 'unknown_field': 'unknownField', + 'unknown_object': {}, + }, ), ], ), @@ -104,16 +106,16 @@ def test_items_conversion(): types.FunctionDeclaration( name='tool', description='tool-description', - parameters=types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - properties={ - 'key1': types.Schema(type='STRING'), - 'key2': types.Schema(type='NUMBER'), + parameters_json_schema={ + 'type': 'array', + 'items': { + 'type': 'object', + 'properties': { + 'key1': {'type': 'string'}, + 'key2': {'type': 'number'}, }, - ), - ), + }, + }, ), ], ), @@ -146,13 +148,13 @@ def test_any_of_conversion(): types.FunctionDeclaration( name='tool', description='tool-description', - parameters=types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='STRING'), - types.Schema(type='NUMBER'), + parameters_json_schema={ + 'type': 'object', + 'any_of': [ + {'type': 'string'}, + {'type': 'number'}, ], - ), + }, ), ], ), @@ -185,13 +187,13 @@ def test_properties_conversion(): types.FunctionDeclaration( name='tool', description='tool-description', - parameters=types.Schema( - type='OBJECT', - properties={ - 'key1': types.Schema(type='STRING'), - 'key2': types.Schema(type='NUMBER'), + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'key1': {'type': 'string'}, + 'key2': {'type': 'number'}, }, - ), + }, ), ], ), diff --git a/google/genai/tests/models/test_generate_content.py b/google/genai/tests/models/test_generate_content.py index ae64eb89f..f0b333512 100644 --- a/google/genai/tests/models/test_generate_content.py +++ b/google/genai/tests/models/test_generate_content.py @@ -2196,28 +2196,6 @@ class Foo(BaseModel): ) -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_function(client): - def get_weather(city: str) -> str: - """Returns the weather in a city.""" - return f'The weather in {city} is sunny and 100 degrees.' - - response = client.models.generate_content( - model=GEMINI_FLASH_LATEST, - contents=( - 'What is the weather like in Sunnyvale? Answer in very short' - ' sentence.' - ), - config={ - 'tools': [get_weather], - }, - ) - assert '100' in response.text - - def test_invalid_input_without_transformer(client): with pytest.raises(ValidationError) as e: client.models.generate_content( diff --git a/google/genai/tests/models/test_generate_content_mcp.py b/google/genai/tests/models/test_generate_content_mcp.py index 7e3f0a219..e44ef8c59 100644 --- a/google/genai/tests/models/test_generate_content_mcp.py +++ b/google/genai/tests/models/test_generate_content_mcp.py @@ -46,7 +46,7 @@ @pytest.mark.asyncio async def test_mcp_tools_async(client): response = await client.aio.models.generate_content( - model='gemini-2.5-flash', + model='gemini-3.1-pro-preview', contents=t.t_contents('What is the weather in Boston?'), config={ 'tools': [ @@ -61,12 +61,9 @@ async def test_mcp_tools_async(client): ], }, ) - assert response.function_calls == [ - types.FunctionCall( - name='get_weather', - args={'location': 'Boston'}, - ) - ] + assert len(response.function_calls) == 1 + assert response.function_calls[0].name == 'get_weather' + assert response.function_calls[0].args == {'location': 'Boston'} @pytest.mark.asyncio @@ -89,92 +86,19 @@ async def test_mcp_tools_with_custom_headers_async(client): ], } response = await client.aio.models.generate_content( - model='gemini-2.5-flash', + model='gemini-3.1-pro-preview', contents=t.t_contents('What is the weather in Boston?'), config=config, ) - assert response.function_calls == [ - types.FunctionCall( - name='get_weather', - args={'location': 'Boston'}, - ) - ] + assert len(response.function_calls) == 1 + assert response.function_calls[0].name == 'get_weather' + assert response.function_calls[0].args == {'location': 'Boston'} # Assert config is not modified. assert config['http_options']['headers'] == { 'x-goog-api-client': 'google-genai-sdk/1.0.0 gl-python/1.0.0' } -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC by default is disabled in private models.py', -) -@pytest.mark.asyncio -async def test_mcp_tools_subsequent_calls_async(client): - class MockMcpClientSession(McpClientSession): - - def __init__(self): - self._read_stream = None - self._write_stream = None - - async def list_tools(self): - return mcp_types.ListToolsResult( - tools=[ - mcp_types.Tool( - name='get_weather', - description='Get the weather in a city.', - inputSchema={ - 'type': 'object', - 'properties': {'location': {'type': 'string'}}, - }, - ), - mcp_types.Tool( - name='add_numbers', - description='Add two numbers together.', - inputSchema={ - 'type': 'object', - 'properties': { - 'a': {'type': 'number'}, - 'b': {'type': 'number'}, - }, - }, - ), - ] - ) - - async def call_tool( - self, - name: str, - arguments: dict[str, Any], - ): - if name == 'get_weather': - return mcp_types.CallToolResult( - content=[mcp_types.TextContent(type='text', text='Sunny')] - ) - else: - return mcp_types.CallToolResult( - content=[mcp_types.TextContent(type='text', text='100')] - ) - - config = { - 'tools': [MockMcpClientSession()], - } - - response = await client.aio.models.generate_content( - model='gemini-2.5-flash', - contents=t.t_contents('What is the weather in Boston?'), - config=config, - ) - assert 'sunny' in response.text.lower() - - response_2 = await client.aio.models.generate_content( - model='gemini-2.5-flash', - contents=t.t_contents('What is 50 + 50?'), - config=config, - ) - assert '100' in response_2.text - - @pytest.mark.asyncio async def test_mcp_tools_duplicate_tool_name_raises_error(client): class MockMcpClientSession(McpClientSession): @@ -217,7 +141,7 @@ async def list_tools(self): def test_mcp_tools_synchronous_call(client): response = client.models.generate_content( - model='gemini-2.5-flash', + model='gemini-3.1-pro-preview', contents=t.t_contents('What is the weather in Boston?'), config={ 'tools': [ @@ -232,12 +156,9 @@ def test_mcp_tools_synchronous_call(client): ] }, ) - assert response.function_calls == [ - types.FunctionCall( - name='get_weather', - args={'location': 'Boston'}, - ) - ] + assert len(response.function_calls) == 1 + assert response.function_calls[0].name == 'get_weather' + assert response.function_calls[0].args == {'location': 'Boston'} def test_mcp_session_synchronous_call_raises_error(client): @@ -281,7 +202,7 @@ async def list_tools(self): def test_mcp_tools_synchronous_stream_call(client): response = client.models.generate_content_stream( - model='gemini-2.5-flash', + model='gemini-3.1-pro-preview', contents=t.t_contents('What is the weather in Boston?'), config={ 'tools': [ @@ -296,13 +217,14 @@ def test_mcp_tools_synchronous_stream_call(client): ] }, ) + asserted_function_call = False for chunk in response: - assert chunk.function_calls == [ - types.FunctionCall( - name='get_weather', - args={'location': 'Boston'}, - ) - ] + if chunk.function_calls: + asserted_function_call = True + assert len(chunk.function_calls) == 1 + assert chunk.function_calls[0].name == 'get_weather' + assert chunk.function_calls[0].args == {'location': 'Boston'} + assert asserted_function_call def test_mcp_session_synchronous_stream_call_raises_error(client): diff --git a/google/genai/tests/models/test_generate_content_tools.py b/google/genai/tests/models/test_generate_content_tools.py index 40721bbbf..323f9d6d5 100644 --- a/google/genai/tests/models/test_generate_content_tools.py +++ b/google/genai/tests/models/test_generate_content_tools.py @@ -513,6 +513,10 @@ def divide_floats(a: float, b: float) -> float: }, ), exception_if_vertex='only supported in Gemini Developer API mode', + skip_in_private=( + 'disabled_safety_policies parameter is supported on Vertex AI in' + ' Private SDK' + ), ), pytest_helper.TestTableItem( name='test_computer_use_multi_turn', @@ -684,6 +688,10 @@ def divide_floats(a: float, b: float) -> float: exception_if_vertex=( 'parameter is only supported in Gemini Developer API mode' ), + skip_in_private=( + 'include_server_side_tool_invocations parameter is supported on' + ' Vertex AI in Private SDK' + ), ), pytest_helper.TestTableItem( name='test_include_server_side_tool_invocations_with_tool_call_echo', @@ -757,15 +765,10 @@ def divide_floats(a: float, b: float) -> float: test_method='models.generate_content', test_table=test_table, ), - pytest.mark.skipif( - "config.getoption('--private')", - reason='ComputerUse on Vertex API behaves differently between public and private modules.', - ), ] pytest_plugins = ('pytest_asyncio',) -# Cannot be included in test_table because json serialization fails on function. def test_function_google_search(client): contents = 'What is the price of GOOG?.' config = types.GenerateContentConfig( @@ -779,15 +782,18 @@ def test_function_google_search(client): function_calling_config=types.FunctionCallingConfig(mode='AUTO') ), ) - # bad request to combine function call and google search retrieval - with pytest.raises(errors.ClientError): + with pytest_helper.exception_if_mldev(client, errors.ClientError): client.models.generate_content( - model='gemini-3.5-flash', + model='gemini-3.1-pro-preview', contents=contents, config=config, ) +@pytest.mark.skipif( + "config.getoption('--private')", + reason="include_server_side_tool_invocations is supported on Vertex AI in Private SDK", +) def test_function_google_search_server_side_tool_invocations(client): contents = ( 'What is the weather in Buenos Aires? If it is raining, schedule a' @@ -823,6 +829,10 @@ def test_function_google_search_server_side_tool_invocations(client): ) +@pytest.mark.skipif( + "config.getoption('--private')", + reason="include_server_side_tool_invocations is supported on Vertex AI in Private SDK", +) def test_function_google_search_server_side_tool_invocations_one_tool(client): contents = ( 'What is the weather in Buenos Aires? If it is raining, schedule a' @@ -870,635 +880,14 @@ def test_google_search_stream(client): pass -@pytest.mark.skipif( - sys.version_info >= (3, 13), - reason=( - 'object type is dumped as as opposed to' - ' "OBJECT" in Python 3.13' - ), -) def test_function_calling_without_implementation(client): response = client.models.generate_content( - model='gemini-2.5-flash', + model='gemini-3.1-pro-preview', contents='What is the weather in Boston?', config={ 'tools': [get_weather_declaration_only], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_2_function(client): - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='What is the price of GOOG? And what is the weather in Boston?', - config={ - 'tools': [get_weather, get_stock_price], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - assert '1000' in response.text - assert 'Boston' in response.text - assert 'sunny' in response.text - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -@pytest.mark.asyncio -async def test_2_function_async(client): - response = await client.aio.models.generate_content( - model='gemini-2.5-flash', - contents='What is the price of GOOG? And what is the weather in Boston?', - config={ - 'tools': [get_weather, get_stock_price], - 'automatic_function_calling': {'ignore_call_history': True}, }, ) - assert '1000' in response.text - assert 'Boston' in response.text - assert 'sunny' in response.text - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_automatic_function_calling_with_customized_math_rule(client): - def customized_divide_integers(numerator: int, denominator: int) -> int: - """Divide two integers with customized math rule.""" - return numerator // denominator + 1 - - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [customized_divide_integers], - }, - ) - assert '501' in response.text - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_automatic_function_calling(client): - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - - assert '500' in response.text - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -@pytest.mark.asyncio -async def test_automatic_function_calling_with_async_function(client): - response = await client.aio.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1001.0/2.0?', - config={ - 'tools': [divide_floats_async], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - - assert '500.5' in response.text - - -def test_automatic_function_calling_stream(client): - response = client.models.generate_content_stream( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - chunks = 0 - for part in response: - chunks += 1 - assert part.text is not None or part.candidates[0].finish_reason - - -def test_disable_automatic_function_calling_stream(client): - # If AFC is disabled, the response should contain a function call. - response = client.models.generate_content_stream( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': {'disable': True}, - }, - ) - chunks = 0 - for chunk in response: - chunks += 1 - assert chunk.parts[0].function_call is not None - - -def test_automatic_function_calling_no_function_response_stream(client): - response = client.models.generate_content_stream( - model='gemini-2.5-flash', - contents='what is the weather in Boston?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - chunks = 0 - for part in response: - chunks += 1 - assert part.text is not None or part.candidates[0].finish_reason - - -@pytest.mark.asyncio -async def test_disable_automatic_function_calling_stream_async(client): - # If AFC is disabled, the response should contain a function call. - response = await client.aio.models.generate_content_stream( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': {'disable': True}, - }, - ) - chunks = 0 - async for chunk in response: - chunks += 1 - assert chunk.parts[0].function_call is not None - - -@pytest.mark.asyncio -async def test_automatic_function_calling_no_function_response_stream_async( - client, -): - response = await client.aio.models.generate_content_stream( - model='gemini-2.5-flash', - contents='what is the weather in Boston?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - chunks = 0 - async for chunk in response: - chunks += 1 - assert chunk.text is not None or chunk.candidates[0].finish_reason - - -@pytest.mark.asyncio -async def test_automatic_function_calling_stream_async(client): - response = await client.aio.models.generate_content_stream( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - chunks = 0 - async for chunk in response: - chunks += 1 - assert chunk.text is not None or chunk.candidates[0].finish_reason - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_callable_tools_user_disable_afc(client): - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': { - 'disable': True, - 'ignore_call_history': True, - }, - }, - ) - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_callable_tools_user_disable_afc_with_max_remote_calls(client): - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': { - 'disable': True, - 'maximum_remote_calls': 2, - 'ignore_call_history': True, - }, - }, - ) - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_callable_tools_user_disable_afc_with_max_remote_calls_negative( - client, -): - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': { - 'disable': True, - 'maximum_remote_calls': -1, - 'ignore_call_history': True, - }, - }, - ) - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_callable_tools_user_disable_afc_with_max_remote_calls_zero(client): - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': { - 'disable': True, - 'maximum_remote_calls': 0, - 'ignore_call_history': True, - }, - }, - ) - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_callable_tools_user_enable_afc(client): - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': { - 'disable': False, - 'ignore_call_history': True, - }, - }, - ) - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_callable_tools_user_enable_afc_with_max_remote_calls(client): - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': { - 'disable': False, - 'maximum_remote_calls': 2, - 'ignore_call_history': True, - }, - }, - ) - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_callable_tools_user_enable_afc_with_max_remote_calls_negative( - client, -): - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': { - 'disable': False, - 'maximum_remote_calls': -1, - 'ignore_call_history': True, - }, - }, - ) - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_callable_tools_user_enable_afc_with_max_remote_calls_zero(client): - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': { - 'disable': False, - 'maximum_remote_calls': 0, - 'ignore_call_history': True, - }, - }, - ) - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_automatic_function_calling_with_exception(client): - client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/0?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_automatic_function_calling_float_without_decimal(client): - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000.0/2.0?', - config={ - 'tools': [divide_floats, divide_integers], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - - assert '500.0' in response.text - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_automatic_function_calling_with_pydantic_model(client): - class CityObject(pydantic.BaseModel): - city_name: str - - def get_weather_pydantic_model( - city_object: CityObject, is_winter: bool - ) -> str: - if is_winter: - return f'The weather in {city_object.city_name} is cold and 10 degrees.' - else: - return f'The weather in {city_object.city_name} is sunny and 100 degrees.' - - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='it is winter now, what is the weather in Boston?', - config={ - 'tools': [get_weather_pydantic_model], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - - assert 'cold' in response.text and 'Boston' in response.text - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_automatic_function_calling_with_pydantic_model_in_list_type(client): - class CityObject(pydantic.BaseModel): - city_name: str - - def get_weather_from_list_of_cities( - city_object_list: list[CityObject], is_winter: bool - ) -> str: - result = '' - if is_winter: - for city_object in city_object_list: - result += ( - f'The weather in {city_object.city_name} is cold and 10 degrees.\n' - ) - else: - for city_object in city_object_list: - result += ( - f'The weather in {city_object.city_name} is sunny and 100' - ' degrees.\n' - ) - return result - - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='it is winter now, what is the weather in Boston and New York?', - config={ - 'tools': [get_weather_from_list_of_cities], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - - assert 'cold' in response.text and 'Boston' in response.text - assert 'cold' in response.text and 'New York' in response.text - - -@pytest.mark.skip( - reason=( - 'AFC is in progress of refactoring, this test is failing python 3.14' - ' b/512415555 will update once refactoring from yyyu@ is done' - ), -) -def test_automatic_function_calling_with_pydantic_model_in_union_type(client): - class AnimalObject(pydantic.BaseModel): - name: str - age: int - species: str - - class PlantObject(pydantic.BaseModel): - name: str - height: float - color: str - - def get_information( - object_of_interest: typing.Union[AnimalObject, PlantObject], - ) -> str: - if isinstance(object_of_interest, AnimalObject): - return ( - f'The animal is of {object_of_interest.species} species and is named' - f' {object_of_interest.name} is {object_of_interest.age} years old' - ) - elif isinstance(object_of_interest, PlantObject): - return ( - f'The plant is named {object_of_interest.name} and is' - f' {object_of_interest.height} meters tall and is' - f' {object_of_interest.color} color' - ) - else: - return 'The animal is not supported' - - with pytest_helper.exception_if_vertex(client, errors.ClientError): - response = client.models.generate_content( - model='gemini-2.5-flash', - contents=( - 'I have a one year old cat named Sundae, can you get the' - ' information of the cat for me?' - ), - config={ - 'system_instruction': ( - 'you answer questions based on the tools provided' - ), - 'tools': [get_information], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - assert 'Sundae' in response.text - assert 'cat' in response.text - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_automatic_function_calling_with_union_operator(client): - class AnimalObject(pydantic.BaseModel): - name: str - age: int - species: str - - def get_information( - object_of_interest: str | AnimalObject, - ) -> str: - if isinstance(object_of_interest, AnimalObject): - return ( - f'The animal is of {object_of_interest.species} species and is named' - f' {object_of_interest.name} is {object_of_interest.age} years old' - ) - else: - return f'The object of interest is {object_of_interest}' - - response = client.models.generate_content( - model='gemini-2.5-flash', - contents=( - 'I have a one year old cat named Sundae, can you get the' - ' information of the cat for me?' - ), - config={ - 'tools': [get_information], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - assert response.text - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_automatic_function_calling_with_tuple_param(client): - def output_latlng( - latlng: tuple[float, float], - ) -> str: - return f'The latitude is {latlng[0]} and the longitude is {latlng[1]}' - - response = client.models.generate_content( - model='gemini-2.5-flash', - contents=( - 'The coordinates are (51.509, -0.118). What is the latitude and longitude?' - ), - config={ - 'tools': [output_latlng], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - assert response.text - - -@pytest.mark.skipif( - sys.version_info < (3, 10), - reason='| is only supported in Python 3.10 and above.', -) -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_automatic_function_calling_with_union_operator_return_type(client): - def get_cheese_age(cheese: int) -> int | float: - """ - Retrieves data about the age of the cheese given its ID. - - Args: - cheese_id: The ID of the cheese. - - Returns: - An int or float of the age of the cheese. - """ - if cheese == 1: - return 2.5 - elif cheese == 2: - return 3 - else: - return 0.0 - - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='How old is the cheese with id 2?', - config={ - 'tools': [get_cheese_age], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - assert '3' in response.text - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_automatic_function_calling_with_parameterized_generic_union_type( - client, -): - def describe_cities( - country: str, - cities: typing.Optional[list[str]] = None, - ) -> str: - 'Given a country and an optional list of cities, describe the cities.' - if cities is None: - return 'There are no cities to describe.' - else: - return ( - f'The cities in {country} are: {", ".join(cities)} and they are nice.' - ) - - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='Can you describe the city of San Francisco, USA?', - config={ - 'tools': [describe_cities], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - assert 'San Francisco' in response.text @pytest.mark.asyncio @@ -1522,19 +911,13 @@ def test_empty_tools(client): ) -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) def test_with_1_empty_tool(client): - # Bad request for empty tool. with pytest_helper.exception_if_vertex(client, errors.ClientError): client.models.generate_content( - model='gemini-2.5-flash', + model='gemini-3.1-pro-preview', contents='What is the price of GOOG?.', config={ 'tools': [{}, get_stock_price], - 'automatic_function_calling': {'ignore_call_history': True}, }, ) @@ -1589,383 +972,6 @@ async def test_vai_search_stream_async(client): assert 'retrieval' in str(e) -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_automatic_function_calling_with_coroutine_function(client): - async def divide_integers(a: int, b: int) -> int: - return a // b - - with pytest.raises(errors.UnsupportedFunctionError): - client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -@pytest.mark.asyncio -async def test_automatic_function_calling_with_coroutine_function_async( - client, -): - async def divide_integers(a: int, b: int) -> int: - return a // b - - response = await client.aio.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - - assert '500' in response.text - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC by default is disabled in private models.py', -) -@pytest.mark.asyncio -async def test_automatic_function_calling_async(client): - def divide_integers(a: int, b: int) -> int: - return a // b - - response = await client.aio.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - - assert '500' in response.text - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC by default is disabled in private models.py', -) -@pytest.mark.asyncio -async def test_automatic_function_calling_async_with_exception(client): - def mystery_function(a: int, b: int) -> int: - return a // b - - response = await client.aio.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/0?', - config={ - 'tools': [divide_integers], - 'system_instruction': ( - 'you must first look at the tools and then think about answers' - ), - }, - ) - assert response.automatic_function_calling_history - assert ( - response.automatic_function_calling_history[-1] - .parts[0] - .function_response.response['error'] - ) - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC by default is disabled in private models.py', -) -@pytest.mark.asyncio -async def test_automatic_function_calling_async_float_without_decimal(client): - response = await client.aio.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000.0/2.0?', - config={ - 'tools': [divide_floats, divide_integers], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - - assert '500.0' in response.text - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC by default is disabled in private models.py', -) -@pytest.mark.asyncio -async def test_automatic_function_calling_async_with_pydantic_model(client): - class CityObject(pydantic.BaseModel): - city_name: str - - def get_weather_pydantic_model( - city_object: CityObject, is_winter: bool - ) -> str: - if is_winter: - return f'The weather in {city_object.city_name} is cold and 10 degrees.' - else: - return f'The weather in {city_object.city_name} is sunny and 100 degrees.' - - response = await client.aio.models.generate_content( - model='gemini-2.5-flash', - contents='it is winter now, what is the weather in Boston?', - config={ - 'tools': [get_weather_pydantic_model], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - - # ML Dev couldn't understand pydantic model - if client.vertexai: - assert 'cold' in response.text and 'Boston' in response.text - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC by default is disabled in private models.py', -) -@pytest.mark.asyncio -async def test_automatic_function_calling_async_with_async_function(client): - async def get_current_weather_async(city: str) -> str: - """Returns the current weather in the city.""" - - return 'windy' - - response = await client.aio.models.generate_content( - model='gemini-2.5-flash', - contents='what is the weather in San Francisco?', - config={ - 'tools': [get_current_weather_async], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - - assert 'windy' in response.text - assert 'San Francisco' in response.text - - -@pytest.mark.asyncio -async def test_automatic_function_calling_async_with_async_function_stream( - client, -): - async def get_current_weather_async(city: str) -> str: - """Returns the current weather in the city.""" - - return 'windy' - - response = await client.aio.models.generate_content_stream( - model='gemini-2.5-flash', - contents='what is the weather in San Francisco?', - config={ - 'tools': [get_current_weather_async], - 'automatic_function_calling': {'ignore_call_history': True}, - }, - ) - - chunk = None - async for chunk in response: - if chunk.parts[0].function_call: - assert chunk.parts[0].function_call.name == 'get_current_weather_async' - assert chunk.parts[0].function_call.args['city'] == 'San Francisco' - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_2_function_with_history(client): - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='What is the price of GOOG? And what is the weather in Boston?', - config={ - 'tools': [get_weather, get_stock_price], - 'automatic_function_calling': {'ignore_call_history': False}, - }, - ) - - actual_history = response.automatic_function_calling_history - - assert actual_history[0].role == 'user' - assert ( - actual_history[0].parts[0].text - == 'What is the price of GOOG? And what is the weather in Boston?' - ) - - assert actual_history[1].role == 'model' - assert actual_history[1].parts[0].function_call.model_dump_json( - exclude_none=True - ) == types.FunctionCall( - name='get_stock_price', - args={'symbol': 'GOOG'}, - ).model_dump_json( - exclude_none=True - ) - assert actual_history[1].parts[1].function_call.model_dump_json( - exclude_none=True - ) == types.FunctionCall( - name='get_weather', - args={'city': 'Boston'}, - ).model_dump_json( - exclude_none=True - ) - - assert actual_history[2].role == 'user' - assert actual_history[2].parts[0].function_response.model_dump_json( - exclude_none=True - ) == types.FunctionResponse( - name='get_stock_price', response={'result': '1000'} - ).model_dump_json( - exclude_none=True - ) - assert actual_history[2].parts[1].function_response.model_dump_json( - exclude_none=True - ) == types.FunctionResponse( - name='get_weather', - response={'result': 'The weather in Boston is sunny and 100 degrees.'}, - ).model_dump_json( - exclude_none=True - ) - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC by default is disabled in private models.py', -) -@pytest.mark.asyncio -async def test_2_function_with_history_async(client): - response = await client.aio.models.generate_content( - model='gemini-2.5-flash', - contents='What is the price of GOOG? And what is the weather in Boston?', - config={ - 'tools': [get_weather, get_stock_price], - 'automatic_function_calling': {'ignore_call_history': False}, - }, - ) - - actual_history = response.automatic_function_calling_history - - assert actual_history[0].role == 'user' - assert ( - actual_history[0].parts[0].text - == 'What is the price of GOOG? And what is the weather in Boston?' - ) - - assert actual_history[1].role == 'model' - assert actual_history[1].parts[0].function_call.model_dump_json( - exclude_none=True - ) == types.FunctionCall( - name='get_stock_price', - args={'symbol': 'GOOG'}, - ).model_dump_json( - exclude_none=True - ) - assert actual_history[1].parts[1].function_call.model_dump_json( - exclude_none=True - ) == types.FunctionCall( - name='get_weather', - args={'city': 'Boston'}, - ).model_dump_json( - exclude_none=True - ) - - assert actual_history[2].role == 'user' - assert actual_history[2].parts[0].function_response.model_dump_json( - exclude_none=True - ) == types.FunctionResponse( - name='get_stock_price', response={'result': '1000'} - ).model_dump_json( - exclude_none=True - ) - assert actual_history[2].parts[1].function_response.model_dump_json( - exclude_none=True - ) == types.FunctionResponse( - name='get_weather', - response={'result': 'The weather in Boston is sunny and 100 degrees.'}, - ).model_dump_json( - exclude_none=True - ) - - -class FunctionHolder: - NAME = 'FunctionHolder' - - def is_a_duck(self, number: int) -> str: - return self.NAME + 'says isOdd: ' + str(number % 2 == 1) - - def is_a_rabbit(self, number: int) -> str: - return self.NAME + 'says isEven: ' + str(number % 2 == 0) - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_class_method_tools(client): - # This test is to make sure that instance method tools can be used in - # the generate_content request. - - function_holder = FunctionHolder() - response = client.models.generate_content( - model='gemini-2.0-flash-exp', - contents=( - 'Print the verbatim output of is_a_duck and is_a_rabbit for the' - ' number 100.' - ), - config={ - 'tools': [function_holder.is_a_duck, function_holder.is_a_rabbit], - }, - ) - assert 'FunctionHolder' in response.text - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_disable_afc_in_any_mode(client): - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config=types.GenerateContentConfig( - tools=[divide_integers], - automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=True - ), - tool_config=types.ToolConfig( - function_calling_config=types.FunctionCallingConfig(mode='ANY') - ), - ), - ) - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_afc_once_in_any_mode(client): - response = client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config=types.GenerateContentConfig( - tools=[divide_integers], - automatic_function_calling=types.AutomaticFunctionCallingConfig( - maximum_remote_calls=2 - ), - tool_config=types.ToolConfig( - function_calling_config=types.FunctionCallingConfig(mode='ANY') - ), - ), - ) - - def test_code_execution_tool(client): response = client.models.generate_content( model='gemini-2.0-flash-exp', @@ -1985,64 +991,18 @@ def test_code_execution_tool(client): ) -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_afc_logs_to_logger_instance(client, caplog): - caplog.set_level(logging.DEBUG, logger='google_genai.models') - client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': { - 'disable': False, - 'maximum_remote_calls': 1, - 'ignore_call_history': True, - }, - }, - ) - for log in caplog.records: - assert log.levelname == 'INFO' - assert log.name == 'google_genai.models' - - assert 'AFC is enabled with max remote calls: 1' in caplog.text - assert 'remote call 1 is done' in caplog.text - assert 'Reached max remote calls' in caplog.text - - -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) -def test_suppress_logs_with_sdk_logger(client, caplog): - caplog.set_level(logging.DEBUG, logger='google_genai.models') - sdk_logger = logging.getLogger('google_genai.models') - sdk_logger.setLevel(logging.ERROR) - client.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/2?', - config={ - 'tools': [divide_integers], - 'automatic_function_calling': { - 'disable': False, - 'maximum_remote_calls': 2, - 'ignore_call_history': True, - }, - }, - ) - assert not caplog.text - - def test_tools_chat_curation(client, caplog): caplog.set_level(logging.DEBUG, logger='google_genai.models') sdk_logger = logging.getLogger('google_genai.models') sdk_logger.setLevel(logging.ERROR) - config = { - 'tools': [{'function_declarations': function_declarations}], - } + config = types.ChatConfig( + tools=[ + types.Tool( + function_declarations=function_declarations, + ) + ], + ) chat = client.chats.create( model='gemini-2.5-flash', @@ -2061,13 +1021,9 @@ def test_tools_chat_curation(client, caplog): assert len(history) == 4 -@pytest.mark.skipif( - 'config.getoption("--private")', - reason='AFC removed from private models.py', -) def test_function_declaration_with_callable(client): response = client.models.generate_content( - model='gemini-2.5-pro', + model='gemini-3.1-pro-preview', contents=( 'Divide 1000 by 2. And tell' ' me the weather in London.' @@ -2084,7 +1040,7 @@ def test_function_declaration_with_callable(client): def test_function_declaration_with_callable_stream_now(client): for chunk in client.models.generate_content_stream( - model='gemini-2.5-pro', + model='gemini-3.1-pro-preview', contents='Divide 1000 by 2. And tell me the weather in London.', config={ 'tools': [ @@ -2099,7 +1055,7 @@ def test_function_declaration_with_callable_stream_now(client): @pytest.mark.asyncio async def test_function_declaration_with_callable_async(client): response = await client.aio.models.generate_content( - model='gemini-2.5-pro', + model='gemini-3.1-pro-preview', contents=( 'Divide 1000 by 2. And tell' ' me the weather in London.' @@ -2117,7 +1073,7 @@ async def test_function_declaration_with_callable_async(client): @pytest.mark.asyncio async def test_function_declaration_with_callable_async_stream(client): async for chunk in await client.aio.models.generate_content_stream( - model='gemini-2.5-pro', + model='gemini-3.1-pro-preview', contents='Divide 1000 by 2. And tell me the weather in London.', config={ 'tools': [ diff --git a/google/genai/tests/private/__init__.py b/google/genai/tests/private/__init__.py deleted file mode 100644 index 5d5d078c2..000000000 --- a/google/genai/tests/private/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - - -"""Tests for the Google GenAI SDK's private module.""" diff --git a/google/genai/tests/private/test_send_message_private.py b/google/genai/tests/private/test_send_message_private.py deleted file mode 100644 index 768d1dada..000000000 --- a/google/genai/tests/private/test_send_message_private.py +++ /dev/null @@ -1,250 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -"""Replay tests for private chats.send_message().""" - -import pytest - -from .. import pytest_helper -from ...errors import ClientError -from ..models import test_generate_content_tools -from ...types import Content -from ...types import FunctionCall -from ...types import FunctionResponse -from ...types import Part - - -pytestmark = [ - pytest_helper.setup( - file=__file__, - globals_for_file=globals(), - ), - pytest.mark.skipif( - "not config.getoption('--private')", - reason="This test file is only intended for the private SDK", - ), -] - - -MODEL_NAME = 'gemini-3.1-pro-preview' -get_weather = test_generate_content_tools.get_weather -get_stock_price = test_generate_content_tools.get_stock_price - - -def test_send_message_function_tool_afc_disabled(client): - chat = client.chats.create( - model=MODEL_NAME, - config={ - 'tools': [get_weather], - }, - ) - chat.send_message('What is the weather in Boston?') - history = chat.get_history() - assert len(history) == 2 - assert history[0].role == 'user' - assert history[0].parts[0].text == 'What is the weather in Boston?' - assert history[1].role == 'model' - assert len(history[1].parts) == 1 - assert history[1].parts[0].function_call.name == 'get_weather' - assert history[1].parts[0].function_call.args == {'city': 'Boston'} - - -def test_send_message_function_tool_afc_enabled(client): - chat = client.chats.create( - model=MODEL_NAME, - config={ - 'tools': [get_weather], - 'automatic_function_calling': {'enable': True}, - }, - ) - chat.send_message('What is the weather in Boston?') - history = chat.get_history() - assert len(history) == 4 - assert history[0].role == 'user' - assert history[0].parts[0].text == 'What is the weather in Boston?' - assert history[1].role == 'model' - assert history[1].parts[0].function_call.name == 'get_weather' - assert history[1].parts[0].function_call.args == {'city': 'Boston'} - assert history[2].role == 'user' - assert history[2].parts[0].function_response.name == 'get_weather' - assert history[3].role == 'model' - assert 'sunny' in history[3].parts[0].text.lower() - - -def test_send_message_function_tool_afc_enabled_multi_turn(client): - chat = client.chats.create( - model=MODEL_NAME, - config={ - 'tools': [get_weather, get_stock_price], - 'automatic_function_calling': {'enable': True}, - }, - ) - chat.send_message('What is the weather in Boston?') - history = chat.get_history() - assert len(history) == 4 - assert history[0].role == 'user' - assert history[0].parts[0].text == 'What is the weather in Boston?' - assert history[1].role == 'model' - assert history[1].parts[0].function_call.name == 'get_weather' - assert history[1].parts[0].function_call.args == {'city': 'Boston'} - assert history[2].role == 'user' - assert history[2].parts[0].function_response.name == 'get_weather' - assert history[3].role == 'model' - assert 'sunny' in history[3].parts[0].text.lower() - - chat.send_message('What is the stock price of symbol GOOG?') - history = chat.get_history() - assert len(history) == 8 - assert history[0].role == 'user' - assert history[0].parts[0].text == 'What is the weather in Boston?' - assert history[1].role == 'model' - assert history[1].parts[0].function_call.name == 'get_weather' - assert history[1].parts[0].function_call.args == {'city': 'Boston'} - assert history[2].role == 'user' - assert history[2].parts[0].function_response.name == 'get_weather' - assert history[3].role == 'model' - assert 'sunny' in history[3].parts[0].text.lower() - assert history[4].role == 'user' - assert history[4].parts[0].text == 'What is the stock price of symbol GOOG?' - assert history[5].role == 'model' - assert history[5].parts[0].function_call.name == 'get_stock_price' - assert history[5].parts[0].function_call.args == {'symbol': 'GOOG'} - assert history[6].role == 'user' - assert history[6].parts[0].function_response.name == 'get_stock_price' - assert history[7].role == 'model' - assert '1000' in history[7].parts[0].text - - -def test_send_message_multi_turn_afc_enabled_FC_FR_parts(client): - chat = client.chats.create( - model=MODEL_NAME, - config={ - 'tools': [get_weather, get_stock_price], - 'automatic_function_calling': {'enable': True}, - }, - history=[ - Content( - role='user', - parts=[Part(text='What is the weather in Boston?')], - ), - Content( - role='model', - parts=[ - Part( - function_call=FunctionCall( - name='get_weather', - args={'city': 'Boston'}, - ), - ), - Part( - function_response=FunctionResponse( - name='get_weather', - response={'weather': 'sunny and 80 degrees'}, - ), - ), - Part(text='The weather is sunny.'), - ], - ), - ] - ) - with pytest_helper.exception_if_vertex(client, ClientError): - chat.send_message('What is the stock price of symbol GOOG?') - - -@pytest.mark.asyncio -async def test_async_send_message_function_tool_afc_disabled(client): - chat = client.aio.chats.create( - model=MODEL_NAME, - config={ - 'tools': [get_weather], - }, - ) - await chat.send_message('What is the weather in Boston?') - history = chat.get_history() - assert len(history) == 2 - assert history[0].role == 'user' - assert history[0].parts[0].text == 'What is the weather in Boston?' - assert history[1].role == 'model' - assert len(history[1].parts) == 1 - assert history[1].parts[0].function_call.name == 'get_weather' - assert history[1].parts[0].function_call.args == {'city': 'Boston'} - - -@pytest.mark.asyncio -async def test_async_send_message_function_tool_afc_enabled(client): - chat = client.aio.chats.create( - model=MODEL_NAME, - config={ - 'tools': [get_weather], - 'automatic_function_calling': {'enable': True}, - }, - ) - await chat.send_message('What is the weather in Boston?') - history = chat.get_history() - assert len(history) == 4 - assert history[0].role == 'user' - assert history[0].parts[0].text == 'What is the weather in Boston?' - assert history[1].role == 'model' - assert history[1].parts[0].function_call.name == 'get_weather' - assert history[1].parts[0].function_call.args == {'city': 'Boston'} - assert history[2].role == 'user' - assert history[2].parts[0].function_response.name == 'get_weather' - assert history[3].role == 'model' - assert 'sunny' in history[3].parts[0].text.lower() - - -@pytest.mark.asyncio -async def test_async_send_message_function_tool_afc_enabled_multi_turn(client): - chat = client.aio.chats.create( - model=MODEL_NAME, - config={ - 'tools': [get_weather, get_stock_price], - 'automatic_function_calling': {'enable': True}, - }, - ) - await chat.send_message('What is the weather in Boston?') - history = chat.get_history() - assert len(history) == 4 - assert history[0].role == 'user' - assert history[0].parts[0].text == 'What is the weather in Boston?' - assert history[1].role == 'model' - assert history[1].parts[0].function_call.name == 'get_weather' - assert history[1].parts[0].function_call.args == {'city': 'Boston'} - assert history[2].role == 'user' - assert history[2].parts[0].function_response.name == 'get_weather' - assert history[3].role == 'model' - assert 'sunny' in history[3].parts[0].text.lower() - - await chat.send_message('What is the stock price of symbol GOOG?') - history = chat.get_history() - assert len(history) == 8 - assert history[0].role == 'user' - assert history[0].parts[0].text == 'What is the weather in Boston?' - assert history[1].role == 'model' - assert history[1].parts[0].function_call.name == 'get_weather' - assert history[1].parts[0].function_call.args == {'city': 'Boston'} - assert history[2].role == 'user' - assert history[2].parts[0].function_response.name == 'get_weather' - assert history[3].role == 'model' - assert 'sunny' in history[3].parts[0].text.lower() - assert history[4].role == 'user' - assert history[4].parts[0].text == 'What is the stock price of symbol GOOG?' - assert history[5].role == 'model' - assert history[5].parts[0].function_call.name == 'get_stock_price' - assert history[5].parts[0].function_call.args == {'symbol': 'GOOG'} - assert history[6].role == 'user' - assert history[6].parts[0].function_response.name == 'get_stock_price' - assert history[7].role == 'model' - assert '1000' in history[7].parts[0].text diff --git a/google/genai/tests/private/test_send_message_stream_private.py b/google/genai/tests/private/test_send_message_stream_private.py deleted file mode 100644 index 5743e6704..000000000 --- a/google/genai/tests/private/test_send_message_stream_private.py +++ /dev/null @@ -1,273 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -"""Tests for private send_message_stream.""" - -import pytest - -from .. import pytest_helper -from ..models import test_generate_content_tools - - -pytestmark = [ - pytest_helper.setup( - file=__file__, - globals_for_file=globals(), - ), - pytest.mark.skipif( - "not config.getoption('--private')", - reason="This test file is only intended for the private SDK", - ), -] - - -MODEL_NAME = 'gemini-3.1-pro-preview' -get_weather = test_generate_content_tools.get_weather -get_stock_price = test_generate_content_tools.get_stock_price - - -def test_send_message_stream_function_tool_afc_disabled(client): - chat = client.chats.create( - model=MODEL_NAME, - config={ - 'tools': [get_weather], - }, - ) - for chunk in chat.send_message_stream('What is the weather in Boston?'): - pass - history = chat.get_history() - assert len(history) == 3 - assert history[0].role == 'user' - assert history[1].role == 'model' - assert history[1].parts[0].function_call.name == 'get_weather' - assert history[2].role == 'model' - assert history[2].parts[0].text == '' - - -def test_send_message_stream_function_tool_afc_enabled(client): - chat = client.chats.create( - model=MODEL_NAME, - config={ - 'tools': [get_weather], - 'automatic_function_calling': {'enable': True}, - }, - ) - for chunk in chat.send_message_stream('What is the weather in Boston?'): - pass - history = chat.get_history() - assert len(history) == 6 - assert history[0].role == 'user' - assert history[1].role == 'model' - assert history[1].parts[0].function_call.name == 'get_weather' - assert history[2].role == 'model' - assert history[2].parts[0].text == '' - assert history[3].role == 'user' - assert history[3].parts[0].function_response.name == 'get_weather' - assert history[4].role == 'model' - assert 'Boston' in history[4].parts[0].text - assert history[5].role == 'model' - assert history[5].parts[0].text == '' - - -def test_send_message_stream_function_tool_afc_enabled_multi_turn(client): - chat = client.chats.create( - model=MODEL_NAME, - config={ - 'tools': [get_weather, get_stock_price], - 'automatic_function_calling': {'enable': True}, - }, - ) - for chunk in chat.send_message_stream('What is the weather in Boston?'): - pass - history = chat.get_history() - - if client.vertexai: - assert len(history) == 7 - assert history[0].role == 'user' - assert history[1].role == 'model' - assert history[1].parts[0].function_call.name == 'get_weather' - assert history[2].role == 'model' - assert history[2].parts[0].text == '' - assert history[3].role == 'user' - assert history[3].parts[0].function_response.name == 'get_weather' - assert history[4].role == 'model' - assert 'Boston' in history[4].parts[0].text - assert history[5].role == 'model' - assert '100' in history[5].parts[0].text - assert history[6].role == 'model' - assert history[6].parts[0].text == '' - else: - assert len(history) == 6 - assert history[0].role == 'user' - assert history[1].role == 'model' - assert history[1].parts[0].function_call.name == 'get_weather' - assert history[2].role == 'model' - assert history[2].parts[0].text == '' - assert history[3].role == 'user' - assert history[3].parts[0].function_response.name == 'get_weather' - assert history[4].role == 'model' - assert 'Boston' in history[4].parts[0].text - assert history[5].role == 'model' - - for chunk in chat.send_message_stream('What is the stock price of symbol GOOG?'): - pass - history = chat.get_history() - - if client.vertexai: - assert len(history) == 14 - assert history[7].role == 'user' - assert history[8].role == 'model' - assert history[8].parts[0].function_call.name == 'get_stock_price' - assert history[9].role == 'model' - assert history[9].parts[0].text == '' - assert history[10].role == 'user' - assert history[10].parts[0].function_response.name == 'get_stock_price' - assert history[11].role == 'model' - assert 'GOOG' in history[11].parts[0].text - assert history[12].role == 'model' - assert '1000' in history[12].parts[0].text - assert history[13].role == 'model' - assert history[13].parts[0].text == '' - else: - assert len(history) == 13 - assert history[6].role == 'user' - assert history[7].role == 'model' - assert history[7].parts[0].function_call.name == 'get_stock_price' - assert history[8].role == 'model' - assert history[8].parts[0].text == '' - assert history[9].role == 'user' - assert history[9].parts[0].function_response.name == 'get_stock_price' - assert history[10].role == 'model' - assert 'GOOG' in history[10].parts[0].text - assert history[11].role == 'model' - assert '1000' in history[11].parts[0].text - assert history[12].role == 'model' - assert history[12].parts[0].text == '' - - -@pytest.mark.asyncio -async def test_async_send_message_stream_function_tool_afc_disabled(client): - chat = client.aio.chats.create( - model=MODEL_NAME, - config={ - 'tools': [get_weather], - }, - ) - async for chunk in await chat.send_message_stream( - 'What is the weather in Boston?' - ): - pass - history = chat.get_history() - assert len(history) == 3 - assert history[0].role == 'user' - assert history[1].role == 'model' - assert history[1].parts[0].function_call.name == 'get_weather' - assert history[2].role == 'model' - assert history[2].parts[0].text == '' - - -@pytest.mark.asyncio -async def test_async_send_message_stream_function_tool_afc_enabled(client): - chat = client.aio.chats.create( - model=MODEL_NAME, - config={ - 'tools': [get_weather], - 'automatic_function_calling': {'enable': True}, - }, - ) - async for chunk in await chat.send_message_stream( - 'What is the weather in Boston?' - ): - pass - history = chat.get_history() - if client.vertexai: - assert len(history) == 6 - assert history[0].role == 'user' - assert history[1].role == 'model' - assert history[1].parts[0].function_call.name == 'get_weather' - assert history[2].role == 'model' - assert history[2].parts[0].text == '' - assert history[3].role == 'user' - assert history[3].parts[0].function_response.name == 'get_weather' - assert history[4].role == 'model' - assert 'Boston' in history[4].parts[0].text - assert history[5].parts[0].text == '' - else: - assert len(history) == 7 - assert history[0].role == 'user' - assert history[1].role == 'model' - assert history[1].parts[0].function_call.name == 'get_weather' - assert history[2].role == 'model' - assert history[2].parts[0].text == '' - assert history[3].role == 'user' - assert history[3].parts[0].function_response.name == 'get_weather' - assert history[4].role == 'model' - assert 'Boston' in history[4].parts[0].text - assert history[5].role == 'model' - assert 'degrees' in history[5].parts[0].text - assert history[6].role == 'model' - assert history[6].parts[0].text == '' - - -@pytest.mark.asyncio -async def test_async_send_message_stream_function_tool_afc_enabled_multi_turn( - client, -): - chat = client.aio.chats.create( - model=MODEL_NAME, - config={ - 'tools': [get_weather, get_stock_price], - 'automatic_function_calling': {'enable': True}, - }, - ) - async for _ in await chat.send_message_stream( - 'What is the weather in Boston?' - ): - pass - history = chat.get_history() - - assert len(history) == 6 - assert history[0].role == 'user' - assert history[1].role == 'model' - assert history[1].parts[0].function_call.name == 'get_weather' - assert history[2].role == 'model' - assert history[2].parts[0].text == '' - assert history[3].role == 'user' - assert history[3].parts[0].function_response.name == 'get_weather' - assert history[4].role == 'model' - assert 'Boston' in history[4].parts[0].text - assert history[5].role == 'model' - assert history[5].parts[0].text == '' - - async for _ in await chat.send_message_stream( - 'What is the stock price of symbol GOOG?' - ): - pass - history = chat.get_history() - - assert len(history) == 13 - assert history[6].role == 'user' - assert history[7].role == 'model' - assert history[7].parts[0].function_call.name == 'get_stock_price' - assert history[8].role == 'model' - assert history[8].parts[0].text == '' - assert history[9].role == 'user' - assert history[9].parts[0].function_response.name == 'get_stock_price' - assert history[10].role == 'model' - assert 'stock' in history[10].parts[0].text - assert history[11].role == 'model' - assert '1000' in history[11].parts[0].text - assert history[12].role == 'model' - assert history[12].parts[0].text == '' diff --git a/google/genai/tests/pytest_helper.py b/google/genai/tests/pytest_helper.py index f22f1e8bb..84e7bb9b3 100644 --- a/google/genai/tests/pytest_helper.py +++ b/google/genai/tests/pytest_helper.py @@ -38,6 +38,10 @@ class TestTableItem(types.TestTableItem): parameters: SerializeAsAny[BaseModel] = Field( description="""The parameters to the test. Use pydantic models.""", ) + skip_in_private: Optional[str] = Field( + default=None, + description="""When set to a reason string, this test will be skipped in private SDK mode.""", + ) def base_test_function( @@ -48,6 +52,9 @@ def base_test_function( test_table_item: TestTableItem, globals_for_file: dict[str, Any], ): + if getattr(client._api_client, '_private', False) and test_table_item.skip_in_private: + pytest.skip(test_table_item.skip_in_private) + replay_id = ( test_table_item.override_replay_id if test_table_item.override_replay_id diff --git a/google/genai/tests/transformers/test_t_tool.py b/google/genai/tests/transformers/test_t_tool.py index 32d86221b..b7c3ee259 100644 --- a/google/genai/tests/transformers/test_t_tool.py +++ b/google/genai/tests/transformers/test_t_tool.py @@ -62,14 +62,14 @@ def test_func(arg1: str, arg2: int): function_declarations=[ types.FunctionDeclaration( name='test_func', - parameters=types.Schema( - type='OBJECT', - properties={ - 'arg1': types.Schema(type='STRING'), - 'arg2': types.Schema(type='INTEGER'), + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'arg1': {'type': 'string'}, + 'arg2': {'type': 'integer'}, }, - required=['arg1', 'arg2'], - ), + 'required': ['arg1', 'arg2'], + }, ) ] ) @@ -125,13 +125,13 @@ def test_mcp_tool(client): types.FunctionDeclaration( name='tool', description='tool-description', - parameters=types.Schema( - type='OBJECT', - properties={ - 'key1': types.Schema(type='STRING'), - 'key2': types.Schema(type='NUMBER'), + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'key1': {'type': 'string'}, + 'key2': {'type': 'number'}, }, - ), + }, ) ] ) diff --git a/google/genai/tests/transformers/test_t_tools.py b/google/genai/tests/transformers/test_t_tools.py index 05d8ce263..6f0ff1900 100644 --- a/google/genai/tests/transformers/test_t_tools.py +++ b/google/genai/tests/transformers/test_t_tools.py @@ -65,14 +65,14 @@ def test_func(arg1: str, arg2: int): function_declarations=[ types.FunctionDeclaration( name='test_func', - parameters=types.Schema( - type='OBJECT', - properties={ - 'arg1': types.Schema(type='STRING'), - 'arg2': types.Schema(type='INTEGER'), + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'arg1': {'type': 'string'}, + 'arg2': {'type': 'integer'}, }, - required=['arg1', 'arg2'], - ), + 'required': ['arg1', 'arg2'], + }, ) ] ) @@ -125,13 +125,13 @@ def test_mcp_tool(client): types.FunctionDeclaration( name='tool', description='tool-description', - parameters=types.Schema( - type='OBJECT', - properties={ - 'key1': types.Schema(type='STRING'), - 'key2': types.Schema(type='NUMBER'), + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'key1': {'type': 'string'}, + 'key2': {'type': 'number'}, }, - ), + }, ) ] ) @@ -173,22 +173,22 @@ def test_multiple_tools(client): types.FunctionDeclaration( name='tool1', description='tool1-description', - parameters=types.Schema( - type='OBJECT', - properties={ - 'key1': types.Schema(type='STRING'), + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'key1': {'type': 'string'}, }, - ), + }, ), types.FunctionDeclaration( name='tool2', description='tool2-description', - parameters=types.Schema( - type='OBJECT', - properties={ - 'key1': types.Schema(type='NUMBER'), + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'key1': {'type': 'number'}, }, - ), + }, ), ] ), diff --git a/google/genai/tests/types/test_future.py b/google/genai/tests/types/test_future.py index dc5eb597c..f75736a8a 100644 --- a/google/genai/tests/types/test_future.py +++ b/google/genai/tests/types/test_future.py @@ -42,53 +42,62 @@ def test_future_annotation_simple_type(): def func_under_test(param_1: str, param_2: int) -> str: return '123' - expected_schema_mldev = types.FunctionDeclaration( + expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'param_1': types.Schema(type='STRING'), - 'param_2': types.Schema(type='INTEGER'), + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'param_1': {'type': 'string'}, + 'param_2': {'type': 'integer'}, }, - required=['param_1', 'param_2'], - ), + 'required': ['param_1', 'param_2'], + }, + response_json_schema={ + 'type': 'string', + }, ) - expected_schema_vertex = copy.deepcopy(expected_schema_mldev) - expected_schema_vertex.response = types.Schema(type='STRING') actual_schema_mldev = types.FunctionDeclaration.from_callable( client=mldev_client, callable=func_under_test ) actual_schema_vertex = types.FunctionDeclaration.from_callable( client=vertex_client, callable=func_under_test ) - assert actual_schema_mldev == expected_schema_mldev - assert actual_schema_vertex == expected_schema_vertex + assert actual_schema_mldev == expected_schema + assert actual_schema_vertex == expected_schema def test_future_annotation_complex_type(): def func_under_test(param_1: ComplexType, param_2: int) -> str: return '123' - expected_schema_mldev = types.FunctionDeclaration( + expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'param_1': types.Schema( - type='OBJECT', - properties={ - 'param_x': types.Schema(type='STRING'), - 'param_y': types.Schema(type='INTEGER'), + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'param_1': { + 'properties': { + 'param_x': { + 'title': 'Param X', + 'type': 'string', + }, + 'param_y': { + 'title': 'Param Y', + 'type': 'integer', + }, }, - required=['param_x', 'param_y'], - ), - 'param_2': types.Schema(type='INTEGER'), + 'required': ['param_x', 'param_y'], + 'title': 'ComplexType', + 'type': 'object', + }, + 'param_2': {'type': 'integer'}, }, - required=['param_1', 'param_2'], - ) + 'required': ['param_1', 'param_2'], + }, + response_json_schema={ + 'type': 'string', + }, ) - expected_schema_vertex = copy.deepcopy(expected_schema_mldev) - expected_schema_vertex.response = types.Schema(type='STRING') actual_schema_mldev = types.FunctionDeclaration.from_callable( client=mldev_client, callable=func_under_test @@ -96,6 +105,5 @@ def func_under_test(param_1: ComplexType, param_2: int) -> str: actual_schema_vertex = types.FunctionDeclaration.from_callable( client=vertex_client, callable=func_under_test ) - - assert actual_schema_mldev == expected_schema_mldev - assert actual_schema_vertex == expected_schema_vertex + assert actual_schema_mldev == expected_schema + assert actual_schema_vertex == expected_schema diff --git a/google/genai/tests/types/test_schema_from_json_schema.py b/google/genai/tests/types/test_schema_from_json_schema.py deleted file mode 100644 index 04a72c45e..000000000 --- a/google/genai/tests/types/test_schema_from_json_schema.py +++ /dev/null @@ -1,417 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the 'License'); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an 'AS IS' BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import logging -import pydantic - -from ... import types - - -def _get_not_none_fields(model: pydantic.BaseModel) -> list[str]: - """Returns field names in a Pydantic model whose values are not None.""" - return [ - field for field, value in model.model_dump().items() if value is not None - ] - - -def test_empty_json_schema_conversion(): - """Test conversion of empty JSONSchema to Schema.""" - json_schema = types.JSONSchema() - gemini_api_schema = types.Schema.from_json_schema(json_schema=json_schema) - vertex_ai_schema = types.Schema.from_json_schema( - json_schema=json_schema, api_option='VERTEX_AI' - ) - - assert gemini_api_schema == types.Schema() - assert vertex_ai_schema == types.Schema() - - -def test_not_null_type_conversion(): - """Test conversion of JSONSchema.type to Schema.type""" - json_schema_types = [ - 'string', - 'number', - 'integer', - 'boolean', - 'array', - 'object', - ] - schema_types = [ - 'STRING', - 'NUMBER', - 'INTEGER', - 'BOOLEAN', - 'ARRAY', - 'OBJECT', - ] - for json_schema_type, expected_type in zip(json_schema_types, schema_types): - json_schema1 = types.JSONSchema(type=types.JSONSchemaType(json_schema_type)) - json_schema2 = types.JSONSchema(type=json_schema_type) - gemini_api_schema1 = types.Schema.from_json_schema(json_schema=json_schema1) - vertex_ai_schema1 = types.Schema.from_json_schema( - json_schema=json_schema1, api_option='VERTEX_AI' - ) - gemini_api_schema2 = types.Schema.from_json_schema(json_schema=json_schema2) - vertex_ai_schema2 = types.Schema.from_json_schema( - json_schema=json_schema2, api_option='VERTEX_AI' - ) - - gemini_api_not_none_field_name1 = _get_not_none_fields(gemini_api_schema1) - vertex_api_not_none_field_name1 = _get_not_none_fields(vertex_ai_schema1) - gemini_api_not_none_field_name2 = _get_not_none_fields(gemini_api_schema2) - vertex_ai_not_none_field_name2 = _get_not_none_fields(vertex_ai_schema2) - - assert gemini_api_schema1.type == expected_type - assert vertex_ai_schema1.type == expected_type - assert gemini_api_schema2.type == expected_type - assert vertex_ai_schema2.type == expected_type - assert gemini_api_not_none_field_name1 == ['type'] - assert vertex_api_not_none_field_name1 == ['type'] - assert gemini_api_not_none_field_name2 == ['type'] - assert vertex_ai_not_none_field_name2 == ['type'] - - -def test_nullable_conversion(): - """Test conversion of JSONSchema.nullable to Schema.nullable""" - json_schema1 = types.JSONSchema( - type=[types.JSONSchemaType('string'), types.JSONSchemaType('null')], - ) - json_schema2 = types.JSONSchema( - type=['string', 'null'], - ) - gemini_api_schema1 = types.Schema.from_json_schema(json_schema=json_schema1) - vertex_ai_schema1 = types.Schema.from_json_schema( - json_schema=json_schema1, api_option='VERTEX_AI' - ) - gemini_api_schema2 = types.Schema.from_json_schema(json_schema=json_schema2) - vertex_ai_schema2 = types.Schema.from_json_schema( - json_schema=json_schema2, api_option='VERTEX_AI' - ) - gemini_api_not_none_field_names1 = _get_not_none_fields(gemini_api_schema1) - vertex_ai_not_none_field_names1 = _get_not_none_fields(vertex_ai_schema1) - gemini_api_not_none_field_names2 = _get_not_none_fields(gemini_api_schema2) - vertex_ai_not_none_field_names2 = _get_not_none_fields(vertex_ai_schema2) - - assert gemini_api_schema1.nullable - assert vertex_ai_schema1.nullable - assert gemini_api_schema2.nullable - assert vertex_ai_schema2.nullable - assert set(gemini_api_not_none_field_names1) == set(['type', 'nullable']) - assert set(vertex_ai_not_none_field_names1) == set(['type', 'nullable']) - assert set(gemini_api_not_none_field_names2) == set(['type', 'nullable']) - assert set(vertex_ai_not_none_field_names2) == set(['type', 'nullable']) - - -def test_nullable_in_union_like_type_conversion(): - """Test conversion of JSONSchema.nullable to Schema.nullable""" - json_schema1 = types.JSONSchema( - type=[ - types.JSONSchemaType('string'), - types.JSONSchemaType('null'), - types.JSONSchemaType('object'), - types.JSONSchemaType('number'), - types.JSONSchemaType('array'), - types.JSONSchemaType('boolean'), - types.JSONSchemaType('integer'), - ], - ) - gemini_api_schema1 = types.Schema.from_json_schema(json_schema=json_schema1) - vertex_ai_schema1 = types.Schema.from_json_schema( - json_schema=json_schema1, api_option='VERTEX_AI' - ) - gemini_api_not_none_field_names1 = _get_not_none_fields(gemini_api_schema1) - vertex_ai_not_none_field_names1 = _get_not_none_fields(vertex_ai_schema1) - json_schema2 = types.JSONSchema( - type=[ - 'string', - 'null', - 'object', - 'number', - 'array', - 'boolean', - 'integer', - ] - ) - gemini_api_schema2 = types.Schema.from_json_schema(json_schema=json_schema2) - vertex_ai_schema2 = types.Schema.from_json_schema( - json_schema=json_schema2, api_option='VERTEX_AI' - ) - expected_schema = types.Schema( - nullable=True, - any_of=[ - types.Schema(type='STRING'), - types.Schema(type='OBJECT'), - types.Schema(type='NUMBER'), - types.Schema(type='ARRAY'), - types.Schema(type='BOOLEAN'), - types.Schema(type='INTEGER'), - ], - ) - - assert gemini_api_schema1 == expected_schema - assert vertex_ai_schema1 == expected_schema - assert gemini_api_schema2 == expected_schema - assert vertex_ai_schema2 == expected_schema - - -def test_union_like_type_conversion_suite1(): - """Test conversion of JSONSchema.type to Schema.any_of""" - json_schema = types.JSONSchema( - type=[ - types.JSONSchemaType('string'), - types.JSONSchemaType('object'), - types.JSONSchemaType('null'), - ], - description='description', - default='default', - max_length=10, - min_length=5, - enum=['value1', 'value2'], - format='format', - pattern='pattern', - title='title', - min_properties=1, - max_properties=2, - required=['field1', 'field2'], - properties={ - 'field1': types.JSONSchema(type='string'), - 'field2': types.JSONSchema(type='integer'), - }, - ) - actual_gemini_api_schema = types.Schema.from_json_schema( - json_schema=json_schema - ) - actual_vertex_ai_schema = types.Schema.from_json_schema( - json_schema=json_schema, api_option='VERTEX_AI' - ) - expected_schema = types.Schema( - nullable=True, - any_of=[ - types.Schema( - type='STRING', - description='description', - max_length=10, - min_length=5, - enum=['value1', 'value2'], - format='format', - pattern='pattern', - title='title', - ), - types.Schema( - type='OBJECT', - properties={ - 'field1': types.Schema(type='STRING'), - 'field2': types.Schema(type='INTEGER'), - }, - required=['field1', 'field2'], - min_properties=1, - max_properties=2, - title='title', - description='description', - ), - ], - ) - - assert actual_gemini_api_schema == expected_schema - assert actual_vertex_ai_schema == expected_schema - - -def test_union_like_type_conversion_suite2(): - """Test conversion of JSONSchema.type to Schema.any_of""" - json_schema = types.JSONSchema( - type=[ - types.JSONSchemaType('integer'), - types.JSONSchemaType('array'), - ], - description='description', - items=types.JSONSchema(type='integer', maximum=2, minimum=1), - min_items=1, - max_items=2, - title='title', - enum=['1', '2'], - maximum=2, - minimum=1, - ) - actual_gemini_api_schema = types.Schema.from_json_schema( - json_schema=json_schema - ) - actual_vertex_ai_schema = types.Schema.from_json_schema( - json_schema=json_schema, api_option='VERTEX_AI' - ) - expected_schema = types.Schema( - any_of=[ - types.Schema( - type='INTEGER', - description='description', - maximum=2, - minimum=1, - enum=['1', '2'], - title='title', - ), - types.Schema( - type='ARRAY', - items=types.Schema(type='INTEGER', maximum=2, minimum=1), - min_items=1, - max_items=2, - title='title', - description='description', - ), - ], - ) - - assert actual_gemini_api_schema == expected_schema - assert actual_vertex_ai_schema == expected_schema - - -def test_array_type_conversion(): - """Test conversion of JSONSchema.items to Schema.items""" - json_schema = types.JSONSchema( - type=types.JSONSchemaType('array'), - items=types.JSONSchema( - type='object', - properties={ - 'field1': types.JSONSchema(type='string'), - 'field2': types.JSONSchema(type='integer'), - }, - required=['field1', 'field2'], - min_properties=1, - max_properties=2, - title='title', - description='description', - ), - ) - gemini_api_schema = types.Schema.from_json_schema(json_schema=json_schema) - vertex_ai_schema = types.Schema.from_json_schema( - json_schema=json_schema, api_option='VERTEX_AI' - ) - expected_schema = types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - properties={ - 'field1': types.Schema(type='STRING'), - 'field2': types.Schema(type='INTEGER'), - }, - required=['field1', 'field2'], - min_properties=1, - max_properties=2, - title='title', - description='description', - ), - ) - - assert gemini_api_schema == expected_schema - assert vertex_ai_schema == expected_schema - - -def test_complex_object_type_conversion(): - """Test conversion of JSONSchema.properties to Schema.properties""" - json_schema = types.JSONSchema( - type=types.JSONSchemaType('object'), - properties={ - 'field1': types.JSONSchema( - type=['string', 'array', 'null'], - description='description1', - max_length=20, - min_length=15, - enum=['value1', 'value2'], - format='format', - pattern='pattern', - title='title1', - items=types.JSONSchema(type='integer', maximum=2, minimum=1), - min_items=1, - max_items=2, - ), - 'field2': types.JSONSchema(type='integer'), - }, - required=['field1', 'field2'], - min_properties=1, - max_properties=2, - title='title', - description='description', - ) - gemini_api_schema = types.Schema.from_json_schema(json_schema=json_schema) - vertex_ai_schema = types.Schema.from_json_schema( - json_schema=json_schema, api_option='VERTEX_AI' - ) - expected_schema = types.Schema( - type='OBJECT', - properties={ - 'field1': types.Schema( - nullable=True, - any_of=[ - types.Schema( - type='STRING', - description='description1', - max_length=20, - min_length=15, - enum=['value1', 'value2'], - format='format', - pattern='pattern', - title='title1', - ), - types.Schema( - type='ARRAY', - items=types.Schema(type='INTEGER', maximum=2, minimum=1), - min_items=1, - max_items=2, - title='title1', - description='description1', - ), - ], - ), - 'field2': types.Schema(type='INTEGER'), - }, - required=['field1', 'field2'], - min_properties=1, - max_properties=2, - title='title', - description='description', - ) - - assert gemini_api_schema == expected_schema - assert vertex_ai_schema == expected_schema - - -def test_from_json_schema_logs_only_once(caplog): - """Test that the info message is logged only once across multiple from_json_schema calls.""" - from ... import types as types_module - - types_module._from_json_schema_warning_logged = False - - caplog.set_level(logging.INFO, logger='google_genai.types') - - json_schema1 = types_module.JSONSchema(type='string') - schema1 = types_module.Schema.from_json_schema(json_schema=json_schema1) - - assert len(caplog.records) == 1 - assert 'Json Schema is now supported natively' in caplog.text - assert 'response_json_schema' in caplog.text - - json_schema2 = types_module.JSONSchema(type='number') - schema2 = types_module.Schema.from_json_schema(json_schema=json_schema2) - - assert len(caplog.records) == 1 - - json_schema3 = types_module.JSONSchema(type='object') - schema3 = types_module.Schema.from_json_schema(json_schema=json_schema3) - - assert len(caplog.records) == 1 - - assert schema1.type == types_module.Type('STRING') - assert schema2.type == types_module.Type('NUMBER') - assert schema3.type == types.Type('OBJECT') - - types_module._from_json_schema_warning_logged = False diff --git a/google/genai/tests/types/test_schema_json_schema.py b/google/genai/tests/types/test_schema_json_schema.py deleted file mode 100644 index 576c87ee2..000000000 --- a/google/genai/tests/types/test_schema_json_schema.py +++ /dev/null @@ -1,468 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - - -import logging -import pydantic - -from ... import types - - -def _get_not_none_fields(model: pydantic.BaseModel) -> list[str]: - """Returns field names in a Pydantic model whose values are not None.""" - return [ - field for field, value in model.model_dump().items() if value is not None - ] - - -def test_empty_schema_conversion(): - """Test conversion of empty Schema to JSONSchema.""" - schema = types.Schema() - json_schema = schema.json_schema - not_none_field_names = _get_not_none_fields(json_schema) - - assert json_schema == types.JSONSchema() - assert not_none_field_names == [] - - -def test_not_null_type_conversion(): - """Test conversion of Schema.type to JSONSchema.type.""" - schema_types = [ - 'OBJECT', - 'ARRAY', - 'STRING', - 'NUMBER', - 'BOOLEAN', - 'INTEGER', - ] - json_schema_types = [ - 'object', - 'array', - 'string', - 'number', - 'boolean', - 'integer', - ] - for schema_type, expected_type in zip(schema_types, json_schema_types): - schema = types.Schema(type=schema_type) - json_schema = schema.json_schema - not_none_field_names = _get_not_none_fields(json_schema) - assert json_schema.type == types.JSONSchemaType(expected_type) - assert not_none_field_names == ['type'] - - -def test_unspecified_type_conversion(): - """Test conversion of Schema.type to JSONSchema.type.""" - schema = types.Schema(type='TYPE_UNSPECIFIED') - json_schema = schema.json_schema - not_none_field_names = _get_not_none_fields(json_schema) - - assert json_schema.type is None - assert not_none_field_names == [] - - -def test_nullable_conversion(): - """Test conversion of Schema.nullable to JSONSchema.type.""" - schema = types.Schema(type='STRING', nullable=True) - json_schema = schema.json_schema - not_none_field_names = _get_not_none_fields(json_schema) - - assert set(json_schema.type) == set([ - types.JSONSchemaType('null'), - types.JSONSchemaType('string') - ]) - assert not_none_field_names == ['type'] - - -def test_property_conversion(): - """Test conversion of Schema.properties to JSONSchema.properties.""" - schema = types.Schema( - type='OBJECT', - properties={ - 'key1': types.Schema(type='STRING'), - 'key2': types.Schema(type='NUMBER'), - }, - ) - json_schema = schema.json_schema - not_none_field_names = _get_not_none_fields(json_schema) - - assert json_schema.properties == { - 'key1': types.JSONSchema(type=types.JSONSchemaType('string')), - 'key2': types.JSONSchema(type=types.JSONSchemaType('number')), - } - assert json_schema.type == types.JSONSchemaType('object') - assert not_none_field_names == ['type', 'properties'] - - -def test_complex_property_conversion(): - """Test conversion of complex Schema.properties to JSONSchema.properties.""" - schema = types.Schema( - type='OBJECT', - properties={ - 'key1': types.Schema( - type='OBJECT', - properties={ - 'key2': types.Schema(type='STRING'), - 'key3': types.Schema(type='NUMBER'), - }, - ), - 'key2': types.Schema(type='ARRAY', items=types.Schema(type='STRING')), - }, - ) - json_schema = schema.json_schema - not_none_field_names = _get_not_none_fields(json_schema) - - assert json_schema.properties == { - 'key1': types.JSONSchema( - type=types.JSONSchemaType('object'), - properties={ - 'key2': types.JSONSchema(type=types.JSONSchemaType('string')), - 'key3': types.JSONSchema(type=types.JSONSchemaType('number')), - }, - ), - 'key2': types.JSONSchema( - type=types.JSONSchemaType('array'), - items=types.JSONSchema(type=types.JSONSchemaType('string')), - ), - } - assert json_schema.type == types.JSONSchemaType('object') - assert not_none_field_names == ['type', 'properties'] - - -def test_items_conversion(): - """Test conversion of Schema.items to JSONSchema.items.""" - schema = types.Schema( - type='ARRAY', - items=types.Schema(type='STRING'), - ) - json_schema = schema.json_schema - not_none_field_names = _get_not_none_fields(json_schema) - - assert json_schema.type == types.JSONSchemaType('array') - assert json_schema.items == types.JSONSchema( - type=types.JSONSchemaType('string') - ) - assert not_none_field_names == ['type', 'items'] - - -def test_complex_items_conversion(): - """Test conversion of complex Schema.items to JSONSchema.items.""" - schema = types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - properties={ - 'key1': types.Schema(type='STRING'), - 'key2': types.Schema(type='NUMBER'), - }, - ), - ) - json_schema = schema.json_schema - not_none_field_names = _get_not_none_fields(json_schema) - - assert json_schema.type == types.JSONSchemaType('array') - assert json_schema.items == types.JSONSchema( - type=types.JSONSchemaType('object'), - properties={ - 'key1': types.JSONSchema(type=types.JSONSchemaType('string')), - 'key2': types.JSONSchema(type=types.JSONSchemaType('number')), - }, - ) - assert not_none_field_names == ['type', 'items'] - - -def test_any_of_conversion(): - """Test conversion of Schema.any_of to JSONSchema.any_of.""" - schema = types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='STRING'), - types.Schema(type='NUMBER'), - ], - ) - json_schema = schema.json_schema - not_none_field_names = _get_not_none_fields(json_schema) - - assert json_schema.type == types.JSONSchemaType('object') - assert json_schema.any_of == [ - types.JSONSchema(type=types.JSONSchemaType('string')), - types.JSONSchema(type=types.JSONSchemaType('number')), - ] - assert not_none_field_names == ['type', 'any_of'] - - -def test_complex_any_of_conversion(): - """Test conversion of complex Schema.any_of to JSONSchema.any_of.""" - schema = types.Schema( - type='OBJECT', - any_of=[ - types.Schema( - type='OBJECT', - properties={ - 'key1': types.Schema(type='STRING'), - 'key2': types.Schema(type='NUMBER'), - }, - ), - types.Schema(type='ARRAY', items=types.Schema(type='STRING')), - ], - ) - json_schema = schema.json_schema - not_none_field_names = _get_not_none_fields(json_schema) - - assert json_schema.type == types.JSONSchemaType('object') - assert json_schema.any_of == [ - types.JSONSchema( - type=types.JSONSchemaType('object'), - properties={ - 'key1': types.JSONSchema(type=types.JSONSchemaType('string')), - 'key2': types.JSONSchema(type=types.JSONSchemaType('number')), - }, - ), - types.JSONSchema( - type=types.JSONSchemaType('array'), - items=types.JSONSchema(type=types.JSONSchemaType('string')), - ), - ] - assert not_none_field_names == ['type', 'any_of'] - - -def test_example_conversion(): - """Test conversion of Schema.direct to JSONSchema.direct.""" - schema = types.Schema( - example='this is an example', - ) - json_schema = schema.json_schema - not_none_field_names = _get_not_none_fields(json_schema) - - assert not_none_field_names == [] - - -def test_property_ordering_conversion(): - """Test conversion of Schema.property_ordering to JSONSchema.property_ordering.""" - schema = types.Schema( - property_ordering=['a', 'b'], - ) - json_schema = schema.json_schema - not_none_field_names = _get_not_none_fields(json_schema) - - assert not_none_field_names == [] - - -def test_direct_conversion(): - """Test Schema fiedls that do not need to be converted.""" - schema = types.Schema( - pattern='^[a-z]+$', - default=1, - max_length=10, - title='title', - min_length=2, - min_properties=3, - max_properties=7, - description='description', - enum=['enum1', 'enum2'], - format='email', - max_items=199, - maximum=300, - min_items=6, - minimum=40, - required=['required1', 'required2'], - ) - json_schema = schema.json_schema - not_none_field_names = _get_not_none_fields(json_schema) - - assert json_schema.pattern == '^[a-z]+$' - assert json_schema.default == 1 - assert json_schema.max_length == 10 - assert json_schema.title == 'title' - assert json_schema.min_length == 2 - assert json_schema.min_properties == 3 - assert json_schema.max_properties == 7 - assert json_schema.description == 'description' - assert json_schema.enum == ['enum1', 'enum2'] - assert json_schema.format == 'email' - assert json_schema.max_items == 199 - assert json_schema.maximum == 300 - assert json_schema.min_items == 6 - assert json_schema.minimum == 40 - assert json_schema.required == ['required1', 'required2'] - assert not_none_field_names.sort() == [ - 'pattern', - 'default', - 'max_length', - 'title', - 'min_length', - 'min_properties', - 'max_properties', - 'description', - 'enum', - 'format', - 'max_items', - 'maximum', - 'min_items', - 'minimum', - 'required', - ].sort() - - -def test_complex_any_of_conversion(): - schema = types.Schema( - type=types.Type.OBJECT, - title='Fruit Basket', - description='A structured representation of a fruit basket', - properties={ - 'fruit': types.Schema( - type=types.Type.ARRAY, - description='An ordered list of the fruit in the basket', - items=types.Schema( - any_of=[ - types.Schema( - title='Apple', - description='Describes an apple', - type=types.Type.OBJECT, - properties={ - 'type': types.Schema( - type=types.Type.STRING, - description='Always "apple"', - ), - 'variety': types.Schema( - type=types.Type.STRING, - description=( - 'The variety of apple (e.g., "Granny' - ' Smith")' - ), - ), - }, - property_ordering=['type', 'variety'], - required=['type', 'variety'], - ), - types.Schema( - title='Orange', - description='Describes an orange', - type=types.Type.OBJECT, - properties={ - 'type': types.Schema( - type=types.Type.STRING, - description='Always "orange"', - ), - 'variety': types.Schema( - type=types.Type.STRING, - description=( - 'The variety of orange (e.g.,"Navel' - ' orange")' - ), - ), - }, - property_ordering=['type', 'variety'], - required=['type', 'variety'], - ), - ], - ), - ), - }, - required=['fruit'], - ) - json_schema = schema.json_schema - not_none_field_names = _get_not_none_fields(json_schema) - - assert json_schema.type == types.JSONSchemaType('object') - assert json_schema.title == 'Fruit Basket' - assert json_schema.description == 'A structured representation of a fruit basket' - assert json_schema.properties == { - 'fruit': types.JSONSchema( - type=types.JSONSchemaType('array'), - description='An ordered list of the fruit in the basket', - items=types.JSONSchema( - any_of=[ - types.JSONSchema( - title='Apple', - description='Describes an apple', - type=types.JSONSchemaType('object'), - properties={ - 'type': types.JSONSchema( - type=types.JSONSchemaType('string'), - description='Always "apple"', - ), - 'variety': types.JSONSchema( - type=types.JSONSchemaType('string'), - description=( - 'The variety of apple (e.g., "Granny' - ' Smith")' - ), - ), - }, - required=['type', 'variety'], - ), - types.JSONSchema( - title='Orange', - description='Describes an orange', - type=types.JSONSchemaType('object'), - properties={ - 'type': types.JSONSchema( - type=types.JSONSchemaType('string'), - description='Always "orange"', - ), - 'variety': types.JSONSchema( - type=types.JSONSchemaType('string'), - description=( - 'The variety of orange (e.g.,"Navel orange")' - ), - ), - }, - required=['type', 'variety'], - ), - ], - ), - ), - } - assert json_schema.required == ['fruit'] - assert not_none_field_names == [ - 'type', - 'title', - 'description', - 'properties', - 'required', - ] - - -def test_json_schema_logs_only_once(caplog): - """Test that the info message is logged only once across multiple json_schema calls.""" - from ... import types as types_module - - types_module._json_schema_warning_logged = False - - caplog.set_level(logging.INFO, logger='google_genai.types') - - schema1 = types_module.Schema(type='STRING') - json_schema1 = schema1.json_schema - - assert len(caplog.records) == 1 - assert 'Json Schema is now supported natively' in caplog.text - assert 'response_json_schema' in caplog.text - - schema2 = types_module.Schema(type='NUMBER') - json_schema2 = schema2.json_schema - - assert len(caplog.records) == 1 - - schema3 = types_module.Schema(type='OBJECT') - json_schema3 = schema3.json_schema - - assert len(caplog.records) == 1 - - assert json_schema1.type == types_module.JSONSchemaType('string') - assert json_schema2.type == types_module.JSONSchemaType('number') - assert json_schema3.type == types_module.JSONSchemaType('object') - - types_module._json_schema_warning_logged = False diff --git a/google/genai/tests/types/test_types.py b/google/genai/tests/types/test_types.py index 50a3f9058..a117d18a6 100644 --- a/google/genai/tests/types/test_types.py +++ b/google/genai/tests/types/test_types.py @@ -402,18 +402,18 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema(type='INTEGER'), - 'b': types.Schema(type='NUMBER'), - 'c': types.Schema(type='BOOLEAN'), - 'd': types.Schema(type='STRING'), - 'e': types.Schema(type='ARRAY'), - 'f': types.Schema(type='OBJECT'), + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': {'type': 'integer'}, + 'b': {'type': 'number'}, + 'c': {'type': 'boolean'}, + 'd': {'type': 'string'}, + 'e': {'type': 'array', 'items': {}}, + 'f': {'type': 'object', 'additionalProperties': True}, }, - required=['a', 'b', 'c', 'd', 'e', 'f'], - ), + 'required': ['a', 'b', 'c', 'd', 'e', 'f'], + }, description='test built in primitives and compounds.', ) @@ -448,15 +448,15 @@ def func_under_test(a: str, b: int = 1, c: list = []): expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema(type='STRING'), - 'b': types.Schema(type='INTEGER', default=1), - 'c': types.Schema(type='ARRAY', default=[]), + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': {'type': 'string'}, + 'b': {'type': 'integer', 'default': 1}, + 'c': {'type': 'array', 'items': {}, 'default': []}, }, - required=['a'], - ), + 'required': ['a'], + }, description='test default value.', ) @@ -471,62 +471,6 @@ def func_under_test(a: str, b: int = 1, c: list = []): assert actual_schema_mldev == expected_schema -@pytest.mark.skipif( - sys.version_info < (3, 10), - reason='| is only supported in Python 3.10 and above.', -) -def test_built_in_primitives_compounds(): - def func_under_test1(a: bytes): - pass - - def func_under_test2(a: set): - pass - - def func_under_test3(a: frozenset): - pass - - def func_under_test4(a: type(None)): - pass - - def func_under_test5(a: int | bytes): - pass - - def func_under_test6(a: int | set): - pass - - def func_under_test7(a: int | frozenset): - pass - - def func_under_test8(a: typing.Union[int, bytes]): - pass - - def func_under_test9(a: typing.Union[int, set]): - pass - - def func_under_test10(a: typing.Union[int, frozenset]): - pass - - all_func_under_test = [ - func_under_test1, - func_under_test2, - func_under_test3, - func_under_test4, - func_under_test5, - func_under_test6, - func_under_test7, - func_under_test8, - func_under_test9, - func_under_test10, - ] - for func_under_test in all_func_under_test: - types.FunctionDeclaration.from_callable( - client=mldev_client, callable=func_under_test - ) - types.FunctionDeclaration.from_callable( - client=vertex_client, callable=func_under_test - ) - - @pytest.mark.skipif( sys.version_info < (3, 10), reason='| is only supported in Python 3.10 and above.', @@ -542,29 +486,29 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='STRING'), - types.Schema(type='NUMBER'), - types.Schema(type='BOOLEAN'), + description='test built in union type.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'anyOf': [ + {'type': 'integer'}, + {'type': 'string'}, + {'type': 'number'}, + {'type': 'boolean'}, ], - ), - 'b': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), + 'type': 'object', + }, + 'b': { + 'anyOf': [ + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, ], - ), + 'type': 'object', + }, }, - required=['a', 'b'], - ), - description='test built in union type.', + 'required': ['a', 'b'], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -589,29 +533,29 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='STRING'), - types.Schema(type='NUMBER'), - types.Schema(type='BOOLEAN'), + description='test built in union type.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'anyOf': [ + {'type': 'integer'}, + {'type': 'string'}, + {'type': 'number'}, + {'type': 'boolean'}, ], - ), - 'b': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), + 'type': 'object', + }, + 'b': { + 'anyOf': [ + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, ], - ), + 'type': 'object', + }, }, - required=['a', 'b'], - ), - description='test built in union type.', + 'required': ['a', 'b'], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -625,40 +569,6 @@ def func_under_test( assert actual_schema_mldev == expected_schema -@pytest.mark.skipif( - sys.version_info < (3, 10), - reason='| is only supported in Python 3.10 and above.', -) -def test_default_value_built_in_union_type(): - def func_under_test( - a: int | str = 1.1, - ): - """test default value not compatible built in union type.""" - pass - - types.FunctionDeclaration.from_callable( - client=mldev_client, callable=func_under_test - ) - types.FunctionDeclaration.from_callable( - client=vertex_client, callable=func_under_test - ) - - -def test_default_value_built_in_union_type_all_py_versions(): - def func_under_test( - a: typing.Union[int, str] = 1.1, - ): - """test default value not compatible built in union type.""" - pass - - types.FunctionDeclaration.from_callable( - client=mldev_client, callable=func_under_test - ) - types.FunctionDeclaration.from_callable( - client=vertex_client, callable=func_under_test - ) - - @pytest.mark.skipif( sys.version_info < (3, 10), reason='| is only supported in Python 3.10 and above.', @@ -675,37 +585,37 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='STRING'), + description='test default value built in union type.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'anyOf': [ + {'type': 'integer'}, + {'type': 'string'}, ], - default='1', - ), - 'b': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), + 'type': 'object', + 'default': '1', + }, + 'b': { + 'anyOf': [ + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, ], - default=[], - ), - 'c': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), + 'type': 'object', + 'default': [], + }, + 'c': { + 'anyOf': [ + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, ], - default={}, - ), + 'type': 'object', + 'default': {}, + }, }, - required=[], - ), - description='test default value built in union type.', + 'required': [], + }, ) actual_schema_vertex = types.FunctionDeclaration.from_callable( @@ -731,37 +641,37 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='STRING'), + description='test default value built in union type.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'anyOf': [ + {'type': 'integer'}, + {'type': 'string'}, ], - default='1', - ), - 'b': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), + 'type': 'object', + 'default': '1', + }, + 'b': { + 'anyOf': [ + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, ], - default=[], - ), - 'c': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), + 'type': 'object', + 'default': [], + }, + 'c': { + 'anyOf': [ + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, ], - default={}, - ), + 'type': 'object', + 'default': {}, + }, }, - required=[], - ), - description='test default value built in union type.', + 'required': [], + }, ) actual_schema_vertex = types.FunctionDeclaration.from_callable( @@ -783,17 +693,17 @@ def func_under_test(a: typing.Literal['a', 'b', 'c']): expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='STRING', - enum=['a', 'b', 'c'], - ), - }, - required=['a'], - ), description='test generic alias literal.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'enum': ['a', 'b', 'c'], + 'type': 'string', + }, + }, + 'required': ['a'], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -876,16 +786,17 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='ARRAY', items=types.Schema(type='INTEGER') - ), - }, - required=['a'], - ), description='test generic alias array.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'type': 'array', + 'items': {'type': 'integer'}, + }, + }, + 'required': ['a'], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -914,35 +825,33 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='STRING'), - types.Schema(type='NUMBER'), - types.Schema(type='BOOLEAN'), + description='test generic alias complex array.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'type': 'array', + 'items': { + 'anyOf': [ + {'type': 'integer'}, + {'type': 'string'}, + {'type': 'number'}, + {'type': 'boolean'}, ], - ), - ), - 'b': types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), + }, + }, + 'b': { + 'type': 'array', + 'items': { + 'anyOf': [ + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, ], - ), - ), + }, + }, }, - required=['a', 'b'], - ), - description='test generic alias complex array.', + 'required': ['a', 'b'], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -966,35 +875,33 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='STRING'), - types.Schema(type='NUMBER'), - types.Schema(type='BOOLEAN'), + description='test generic alias complex array.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'type': 'array', + 'items': { + 'anyOf': [ + {'type': 'integer'}, + {'type': 'string'}, + {'type': 'number'}, + {'type': 'boolean'}, ], - ), - ), - 'b': types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), + }, + }, + 'b': { + 'type': 'array', + 'items': { + 'anyOf': [ + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, ], - ), - ), + }, + }, }, - required=['a', 'b'], - ), - description='test generic alias complex array.', + 'required': ['a', 'b'], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -1007,12 +914,6 @@ def func_under_test( assert actual_schema_mldev == expected_schema -@pytest.mark.skip( - reason=( - 'AFC is in progress of refactoring, this test is failing python 3.14' - ' b/512415555 will update once refactoring from yyyu@ is done' - ), -) def test_generic_alias_complex_array_with_default_value(): def func_under_test( @@ -1035,53 +936,47 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='STRING'), - types.Schema(type='NUMBER'), - types.Schema(type='BOOLEAN'), + description='test generic alias complex array with default value.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'type': 'array', + 'items': { + 'anyOf': [ + {'type': 'integer'}, + {'type': 'string'}, + {'type': 'number'}, + {'type': 'boolean'}, ], - ), - default=[1, 'a', 1.1, True], - ), - 'b': types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='STRING'), - types.Schema(type='NUMBER'), - types.Schema(type='BOOLEAN'), + }, + 'default': [1, 'a', 1.1, True], + }, + 'b': { + 'type': 'array', + 'items': { + 'anyOf': [ + {'type': 'integer'}, + {'type': 'string'}, + {'type': 'number'}, + {'type': 'boolean'}, ], - ), - default=[11, 'aa', 1.11, False], - ), - 'c': types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - any_of=[ - types.Schema( - type='ARRAY', - items=types.Schema(type='INTEGER'), - ), - types.Schema(type='INTEGER'), + }, + 'default': [11, 'aa', 1.11, False], + }, + 'c': { + 'type': 'array', + 'items': { + 'anyOf': [ + {'type': 'array', 'items': {'type': 'integer'}}, + {'type': 'integer'}, ], - ), - default=[[1], 2], - ), + }, + 'default': [[1], 2], + }, }, - required=[], - ), - description='test generic alias complex array with default value.', + 'required': [], + }, ) actual_schema_vertex = types.FunctionDeclaration.from_callable( @@ -1095,13 +990,6 @@ def func_under_test( assert actual_schema_mldev == expected_schema - -@pytest.mark.skip( - reason=( - 'AFC is in progress of refactoring, this test is failing python 3.14' - ' b/512415555 will update once refactoring from yyyu@ is done' - ), -) def test_generic_alias_complex_array_with_default_value_all_py_versions(): def func_under_test( @@ -1124,53 +1012,47 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='STRING'), - types.Schema(type='NUMBER'), - types.Schema(type='BOOLEAN'), + description='test generic alias complex array with default value.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'type': 'array', + 'items': { + 'anyOf': [ + {'type': 'integer'}, + {'type': 'string'}, + {'type': 'number'}, + {'type': 'boolean'}, ], - ), - default=[1, 'a', 1.1, True], - ), - 'b': types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='STRING'), - types.Schema(type='NUMBER'), - types.Schema(type='BOOLEAN'), + }, + 'default': [1, 'a', 1.1, True], + }, + 'b': { + 'type': 'array', + 'items': { + 'anyOf': [ + {'type': 'integer'}, + {'type': 'string'}, + {'type': 'number'}, + {'type': 'boolean'}, ], - ), - default=[11, 'aa', 1.11, False], - ), - 'c': types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - any_of=[ - types.Schema( - type='ARRAY', - items=types.Schema(type='INTEGER'), - ), - types.Schema(type='INTEGER'), + }, + 'default': [11, 'aa', 1.11, False], + }, + 'c': { + 'type': 'array', + 'items': { + 'anyOf': [ + {'type': 'array', 'items': {'type': 'integer'}}, + {'type': 'integer'}, ], - ), - default=[[1], 2], - ), + }, + 'default': [[1], 2], + }, }, - required=[], - ), - description='test generic alias complex array with default value.', + 'required': [], + }, ) actual_schema_vertex = types.FunctionDeclaration.from_callable( @@ -1250,14 +1132,19 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema(type='OBJECT'), - }, - required=['a'], - ), description='test generic alias object.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'type': 'object', + 'additionalProperties': { + 'type': 'integer', + }, + } + }, + 'required': ['a'], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -1362,17 +1249,20 @@ def func_under_test(a: typing.Dict[str, int] = {'a': 1}): expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='OBJECT', - default={'a': 1}, - ), - }, - required=[], - ), description='test generic alias object with default value.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'type': 'object', + 'additionalProperties': { + 'type': 'integer', + }, + 'default': {'a': 1}, + } + }, + 'required': [], + }, ) actual_schema_vertex = types.FunctionDeclaration.from_callable( @@ -1417,46 +1307,74 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='OBJECT', - properties={ - 'a_simple': types.Schema(type='INTEGER'), - 'b_simple': types.Schema(type='STRING'), + description='test pydantic model.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'properties': { + 'a_simple': { + 'title': 'A Simple', + 'type': 'integer', + }, + 'b_simple': { + 'title': 'B Simple', + 'type': 'string', + }, }, - required=['a_simple', 'b_simple'], - ), - 'b': types.Schema( - type='OBJECT', - properties={ - 'a_complex': types.Schema( - type='OBJECT', - properties={ - 'a_simple': types.Schema(type='INTEGER'), - 'b_simple': types.Schema(type='STRING'), + 'required': [ + 'a_simple', + 'b_simple', + ], + 'title': 'MySimplePydanticModel', + 'type': 'object', + }, + 'b': { + 'properties': { + 'a_complex': { + '$ref': '#/$defs/MySimplePydanticModel', + }, + 'b_complex': { + 'items': { + '$ref': '#/$defs/MySimplePydanticModel', }, - required=['a_simple', 'b_simple'], - ), - 'b_complex': types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - properties={ - 'a_simple': types.Schema(type='INTEGER'), - 'b_simple': types.Schema(type='STRING'), - }, - required=['a_simple', 'b_simple'], - ), - ), + 'title': 'B Complex', + 'type': 'array', + }, }, - required=['a_complex', 'b_complex'], - ), + 'required': [ + 'a_complex', + 'b_complex', + ], + 'title': 'MyComplexPydanticModel', + 'type': 'object', + }, }, - required=['a', 'b'], - ), - description='test pydantic model.', + 'required': [ + 'a', + 'b', + ], + '$defs': { + 'MySimplePydanticModel': { + 'properties': { + 'a_simple': { + 'title': 'A Simple', + 'type': 'integer', + }, + 'b_simple': { + 'title': 'B Simple', + 'type': 'string', + }, + }, + 'required': [ + 'a_simple', + 'b_simple', + ], + 'title': 'MySimplePydanticModel', + 'type': 'object', + } + }, + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -1482,25 +1400,42 @@ def func_under_test( pass expected_schema = types.FunctionDeclaration( - name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - properties={ - 'a_simple': types.Schema(type='INTEGER'), - 'b_simple': types.Schema(type='STRING'), + name='func_under_test', + description='test pydantic model in list type.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'items': { + '$ref': '#/$defs/MySimplePydanticModel', + }, + 'type': 'array', + } + }, + '$defs': { + 'MySimplePydanticModel': { + 'properties': { + 'a_simple': { + 'title': 'A Simple', + 'type': 'integer', }, - required=['a_simple', 'b_simple'], - ), - ), + 'b_simple': { + 'title': 'B Simple', + 'type': 'string', + }, + }, + 'required': [ + 'a_simple', + 'b_simple', + ], + 'title': 'MySimplePydanticModel', + 'type': 'object', + } }, - required=['a'], - ), - description='test pydantic model in list type.', + 'required': [ + 'a', + ], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -1514,12 +1449,6 @@ def func_under_test( assert actual_schema_vertex == expected_schema -@pytest.mark.skip( - reason=( - 'AFC is in progress of refactoring, this test is failing python 3.14' - ' b/512415555 will update once refactoring from yyyu@ is done' - ), -) def test_pydantic_model_in_union_type(): class CatInformationObject(pydantic.BaseModel): name: str @@ -1539,45 +1468,75 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'animal': types.Schema( - type='OBJECT', - any_of=[ - types.Schema( - type='OBJECT', - properties={ - 'name': types.Schema(type='STRING'), - 'age': types.Schema(type='INTEGER'), - 'like_purring': types.Schema(type='BOOLEAN'), - }, - ), - types.Schema( - type='OBJECT', - properties={ - 'name': types.Schema(type='STRING'), - 'age': types.Schema(type='INTEGER'), - 'like_barking': types.Schema(type='BOOLEAN'), - }, - ), + description='test pydantic model in union type.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'animal': { + 'anyOf': [ + { + '$ref': '#/$defs/CatInformationObject', + }, + { + '$ref': '#/$defs/DogInformationObject', + }, ], - ), + 'type': 'object', + } }, - required=['animal'], - ), - description='test pydantic model in union type.', + 'required': [ + 'animal', + ], + '$defs': { + 'CatInformationObject': { + 'title': 'CatInformationObject', + 'type': 'object', + 'properties': { + 'name': { + 'title': 'Name', + 'type': 'string', + }, + 'age': { + 'title': 'Age', + 'type': 'integer', + }, + 'like_purring': { + 'title': 'Like Purring', + 'type': 'boolean', + }, + }, + 'required': [ + 'name', + 'age', + 'like_purring', + ], + }, + 'DogInformationObject': { + 'title': 'DogInformationObject', + 'type': 'object', + 'properties': { + 'name': { + 'title': 'Name', + 'type': 'string', + }, + 'age': { + 'title': 'Age', + 'type': 'integer', + }, + 'like_barking': { + 'title': 'Like Barking', + 'type': 'boolean', + }, + }, + 'required': [ + 'name', + 'age', + 'like_barking', + ], + }, + }, + }, ) - expected_schema.parameters.properties['animal'].any_of[0].required = [ - 'name', - 'age', - 'like_purring', - ] - expected_schema.parameters.properties['animal'].any_of[1].required = [ - 'name', - 'age', - 'like_barking', - ] actual_schema_mldev = types.FunctionDeclaration.from_callable( client=mldev_client, callable=func_under_test @@ -1604,27 +1563,29 @@ def func_under_test(a: MySimplePydanticModel = mySimplePydanticModel): expected_schema = types.FunctionDeclaration( description='test pydantic model with default value.', name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - default=MySimplePydanticModel(a_simple=1, b_simple='a'), - type='OBJECT', - properties={ - 'a_simple': types.Schema( - nullable=True, - type='INTEGER', - ), - 'b_simple': types.Schema( - nullable=True, - type='STRING', - ), + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'title': 'MySimplePydanticModel', + 'type': 'object', + 'properties': { + 'a_simple': { + 'title': 'A Simple', + 'anyOf': [{'type': 'integer'}, {'type': 'null'}], + }, + 'b_simple': { + 'title': 'B Simple', + 'anyOf': [{'type': 'string'}, {'type': 'null'}], + }, }, - required=[], - ) + 'required': ['a_simple', 'b_simple'], + 'default': MySimplePydanticModel(a_simple=1, b_simple='a'), + } }, - required=[], - ), + 'required': [], + + } ) actual_schema_vertex = types.FunctionDeclaration.from_callable( @@ -1662,12 +1623,6 @@ def func_under_test(a: MyClass): ) -@pytest.mark.skip( - reason=( - 'AFC is in progress of refactoring, this test is failing python 3.14' - ' b/512415555 will update once refactoring from yyyu@ is done' - ), -) def test_type_union(): def func_under_test( @@ -1681,52 +1636,43 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='STRING'), - ], - ), - 'b': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), + description='test type union.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'anyOf': [{'type': 'integer'}, {'type': 'string'}], + 'type': 'object', + }, + 'b': { + 'anyOf': [ + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, ], - ), - 'c': types.Schema( - type='OBJECT', - any_of=[ - types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='NUMBER'), - ], - ), - ), - types.Schema( - type='OBJECT', - ), + 'type': 'object', + }, + 'c': { + 'anyOf': [ + { + 'type': 'array', + 'items': { + 'anyOf': [{'type': 'integer'}, {'type': 'number'}] + }, + }, + {'type': 'object', 'additionalProperties': True}, ], - ), - 'd': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), + 'type': 'object', + }, + 'd': { + 'anyOf': [ + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, ], - ), + 'type': 'object', + }, }, - required=['a', 'b', 'c', 'd'], - ), - description='test type union.', + 'required': ['a', 'b', 'c', 'd'], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -1740,12 +1686,6 @@ def func_under_test( assert actual_schema_mldev == expected_schema -@pytest.mark.skip( - reason=( - 'AFC is in progress of refactoring, this test is failing python 3.14' - ' b/512415555 will update once refactoring from yyyu@ is done' - ), -) def test_type_union_all_py_versions(): def func_under_test( @@ -1758,45 +1698,36 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='STRING'), - ], - ), - 'b': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), + description='test type union.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'anyOf': [{'type': 'integer'}, {'type': 'string'}], + 'type': 'object', + }, + 'b': { + 'anyOf': [ + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, ], - ), - 'c': types.Schema( - type='OBJECT', - any_of=[ - types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='NUMBER'), - ], - ), - ), - types.Schema( - type='OBJECT', - ), + 'type': 'object', + }, + 'c': { + 'anyOf': [ + { + 'type': 'array', + 'items': { + 'anyOf': [{'type': 'integer'}, {'type': 'number'}] + }, + }, + {'type': 'object', 'additionalProperties': True}, ], - ), + 'type': 'object', + }, }, - required=['a', 'b', 'c'], - ), - description='test type union.', + 'required': ['a', 'b', 'c'], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -1821,17 +1752,22 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema(type='STRING'), - 'b': types.Schema( - nullable=True, type='ARRAY', items=types.Schema(type='STRING') - ), - }, - required=['a'], - ), description='test type optional with list.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': {'type': 'string'}, + 'b': { + 'anyOf': [ + {'type': 'array', 'items': {'type': 'string'}}, + {'type': 'null'}, + ], + 'type': 'object', + 'default': None, + }, + }, + 'required': ['a'], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -1925,12 +1861,6 @@ def func_under_test( assert actual_schema_mldev == expected_schema -@pytest.mark.skip( - reason=( - 'AFC is in progress of refactoring, this test is failing python 3.14' - ' b/512415555 will update once refactoring from yyyu@ is done' - ), -) def test_type_union_with_default_value_all_py_versions(): def func_under_test( @@ -1943,48 +1873,45 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='STRING'), + description='test type union with default value.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'type': 'object', + 'anyOf': [ + {'type': 'integer'}, + {'type': 'string'}, ], - default=1, - ), - 'b': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), + 'default': 1, + }, + 'b': { + 'type': 'object', + 'anyOf': [ + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, ], - default=[1], - ), - 'c': types.Schema( - type='OBJECT', - any_of=[ - types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='NUMBER'), + 'default': [1], + }, + 'c': { + 'type': 'object', + 'anyOf': [ + { + 'type': 'array', + 'items': { + 'anyOf': [ + {'type': 'integer'}, + {'type': 'number'}, ], - ), - ), - types.Schema( - type='OBJECT', - ), + }, + }, + {'type': 'object', 'additionalProperties': True}, ], - default={}, - ), + 'default': {}, + }, }, - required=[], - ), - description='test type union with default value.', + 'required': [], + }, ) actual_schema_vertex = types.FunctionDeclaration.from_callable( @@ -2069,38 +1996,44 @@ def func_under_test( expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='NUMBER'), + description='test type nullable.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'type': 'object', + 'anyOf': [ + {'type': 'integer'}, + {'type': 'number'}, + {'type': 'null'}, ], - nullable=True, - ), - 'b': types.Schema( - type='ARRAY', - nullable=True, - ), - 'c': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), + }, + 'b': { + 'type': 'object', + 'anyOf': [ + {'type': 'array', 'items': {}}, + {'type': 'null'}, ], - nullable=True, - ), - 'd': types.Schema( - type='INTEGER', - nullable=True, - default=None, - ), + }, + 'c': { + 'type': 'object', + 'anyOf': [ + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, + {'type': 'null'}, + ], + }, + 'd': { + 'type': 'object', + 'anyOf': [ + {'type': 'integer'}, + {'type': 'null'}, + ], + 'default': None, + }, }, - required=[], - ), - description='test type nullable.', + 'required': ['a', 'b', 'c'], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -2124,32 +2057,38 @@ def func_under_test( """test type nullable.""" pass - expected_schema = types.FunctionDeclaration( - name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'b': types.Schema( - type='ARRAY', - nullable=True, - ), - 'c': types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), + expected_schema = types.FunctionDeclaration( + name='func_under_test', + description='test type nullable.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'b': { + 'type': 'object', + 'anyOf': [ + {'type': 'array', 'items': {}}, + {'type': 'null'}, ], - nullable=True, - ), - 'd': types.Schema( - type='INTEGER', - nullable=True, - default=None, - ), + }, + 'c': { + 'type': 'object', + 'anyOf': [ + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, + {'type': 'null'}, + ], + }, + 'd': { + 'type': 'object', + 'anyOf': [ + {'type': 'integer'}, + {'type': 'null'}, + ], + 'default': None, + }, }, - required=[], - ), - description='test type nullable.', + 'required': ['b', 'c'], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -2168,12 +2107,13 @@ def func_under_test() -> int: """test empty function with return type.""" return 1 - expected_schema_mldev = types.FunctionDeclaration( + expected_schema = types.FunctionDeclaration( name='func_under_test', description='test empty function with return type.', + response_json_schema={ + 'type': 'integer', + }, ) - expected_schema_vertex = copy.deepcopy(expected_schema_mldev) - expected_schema_vertex.response = types.Schema(type='INTEGER') actual_schema_mldev = types.FunctionDeclaration.from_callable( client=mldev_client, callable=func_under_test @@ -2182,8 +2122,8 @@ def func_under_test() -> int: client=vertex_client, callable=func_under_test ) - assert actual_schema_mldev == expected_schema_mldev - assert actual_schema_vertex == expected_schema_vertex + assert actual_schema_mldev == expected_schema + assert actual_schema_vertex == expected_schema def test_simple_function_with_return_type(): @@ -2191,19 +2131,22 @@ def func_under_test(a: int) -> str: """test return type.""" return '' - expected_schema_mldev = types.FunctionDeclaration( + expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema(type='INTEGER'), - }, - required=['a'], - ), description='test return type.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'type': 'integer', + }, + }, + 'required': ['a'], + }, + response_json_schema={ + 'type': 'string', + }, ) - expected_schema_vertex = copy.deepcopy(expected_schema_mldev) - expected_schema_vertex.response = types.Schema(type='STRING') actual_schema_mldev = types.FunctionDeclaration.from_callable( client=mldev_client, callable=func_under_test @@ -2212,8 +2155,8 @@ def func_under_test(a: int) -> str: client=vertex_client, callable=func_under_test ) - assert actual_schema_mldev == expected_schema_mldev - assert actual_schema_vertex == expected_schema_vertex + assert actual_schema_mldev == expected_schema + assert actual_schema_vertex == expected_schema @pytest.mark.skipif( @@ -2226,22 +2169,21 @@ def func_under_test() -> int | str | float | bool | list | dict | None: """test builtin union return type.""" pass - expected_schema_mldev = types.FunctionDeclaration( + expected_schema = types.FunctionDeclaration( name='func_under_test', description='test builtin union return type.', - ) - expected_schema_vertex = copy.deepcopy(expected_schema_mldev) - expected_schema_vertex.response_json_schema = types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='STRING'), - types.Schema(type='NUMBER'), - types.Schema(type='BOOLEAN'), - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), - ], - nullable=True, + response_json_schema={ + 'type': 'object', + 'anyOf': [ + {'type': 'integer'}, + {'type': 'string'}, + {'type': 'number'}, + {'type': 'boolean'}, + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, + {'type': 'null'}, + ], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -2251,8 +2193,8 @@ def func_under_test() -> int | str | float | bool | list | dict | None: client=vertex_client, callable=func_under_test ) - assert actual_schema_mldev == expected_schema_mldev - assert actual_schema_vertex == expected_schema_vertex + assert actual_schema_mldev == expected_schema + assert actual_schema_vertex == expected_schema def test_builtin_union_return_type_all_py_versions(): @@ -2263,22 +2205,21 @@ def func_under_test() -> ( """test builtin union return type.""" pass - expected_schema_mldev = types.FunctionDeclaration( + expected_schema = types.FunctionDeclaration( name='func_under_test', description='test builtin union return type.', - ) - expected_schema_vertex = copy.deepcopy(expected_schema_mldev) - expected_schema_vertex.response_json_schema = types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='STRING'), - types.Schema(type='NUMBER'), - types.Schema(type='BOOLEAN'), - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), - ], - nullable=True, + response_json_schema={ + 'type': 'object', + 'anyOf': [ + {'type': 'integer'}, + {'type': 'string'}, + {'type': 'number'}, + {'type': 'boolean'}, + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, + {'type': 'null'}, + ], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -2288,8 +2229,8 @@ def func_under_test() -> ( client=vertex_client, callable=func_under_test ) - assert actual_schema_mldev == expected_schema_mldev - assert actual_schema_vertex == expected_schema_vertex + assert actual_schema_mldev == expected_schema + assert actual_schema_vertex == expected_schema def test_typing_union_return_type(): @@ -2300,22 +2241,21 @@ def func_under_test() -> ( """test typing union return type.""" pass - expected_schema_mldev = types.FunctionDeclaration( + expected_schema = types.FunctionDeclaration( name='func_under_test', description='test typing union return type.', - ) - expected_schema_vertex = copy.deepcopy(expected_schema_mldev) - expected_schema_vertex.response_json_schema = types.Schema( - type='OBJECT', - any_of=[ - types.Schema(type='INTEGER'), - types.Schema(type='STRING'), - types.Schema(type='NUMBER'), - types.Schema(type='BOOLEAN'), - types.Schema(type='ARRAY'), - types.Schema(type='OBJECT'), - ], - nullable=True, + response_json_schema={ + 'type': 'object', + 'anyOf': [ + {'type': 'integer'}, + {'type': 'string'}, + {'type': 'number'}, + {'type': 'boolean'}, + {'type': 'array', 'items': {}}, + {'type': 'object', 'additionalProperties': True}, + {'type': 'null'}, + ], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -2325,8 +2265,8 @@ def func_under_test() -> ( client=vertex_client, callable=func_under_test ) - assert actual_schema_mldev == expected_schema_mldev - assert actual_schema_vertex == expected_schema_vertex + assert actual_schema_mldev == expected_schema + assert actual_schema_vertex == expected_schema def test_return_type_optional(): @@ -2334,14 +2274,16 @@ def func_under_test() -> typing.Optional[int]: """test return type optional.""" pass - expected_schema_mldev = types.FunctionDeclaration( + expected_schema = types.FunctionDeclaration( name='func_under_test', description='test return type optional.', - ) - expected_schema_vertex = copy.deepcopy(expected_schema_mldev) - expected_schema_vertex.response = types.Schema( - type='INTEGER', - nullable=True, + response_json_schema={ + 'type': 'object', + 'anyOf': [ + {'type': 'integer'}, + {'type': 'null'}, + ], + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -2351,8 +2293,8 @@ def func_under_test() -> typing.Optional[int]: client=vertex_client, callable=func_under_test ) - assert actual_schema_mldev == expected_schema_mldev - assert actual_schema_vertex == expected_schema_vertex + assert actual_schema_mldev == expected_schema + assert actual_schema_vertex == expected_schema def test_return_type_pydantic_model(): @@ -2368,35 +2310,41 @@ def func_under_test() -> MyComplexPydanticModel: """test return type pydantic model.""" pass - expected_schema_mldev = types.FunctionDeclaration( + expected_schema = types.FunctionDeclaration( name='func_under_test', description='test return type pydantic model.', - ) - expected_schema_vertex = copy.deepcopy(expected_schema_mldev) - expected_schema_vertex.response = types.Schema( - type='OBJECT', - properties={ - 'a_complex': types.Schema( - type='OBJECT', - properties={ - 'a_simple': types.Schema(type='INTEGER'), - 'b_simple': types.Schema(type='STRING'), - }, - required=['a_simple', 'b_simple'], - ), - 'b_complex': types.Schema( - type='ARRAY', - items=types.Schema( - type='OBJECT', - properties={ - 'a_simple': types.Schema(type='INTEGER'), - 'b_simple': types.Schema(type='STRING'), + response_json_schema={ + 'title': 'MyComplexPydanticModel', + 'type': 'object', + '$defs': { + 'MySimplePydanticModel': { + 'properties': { + 'a_simple': { + 'title': 'A Simple', + 'type': 'integer', + }, + 'b_simple': { + 'title': 'B Simple', + 'type': 'string', + }, }, - required=['a_simple', 'b_simple'], - ), - ), + 'required': ['a_simple', 'b_simple'], + 'title': 'MySimplePydanticModel', + 'type': 'object', + }, + }, + 'properties': { + 'a_complex': { + '$ref': '#/$defs/MySimplePydanticModel', + }, + 'b_complex': { + 'items': {'$ref': '#/$defs/MySimplePydanticModel'}, + 'title': 'B Complex', + 'type': 'array', + }, + }, + 'required': ['a_complex', 'b_complex'], }, - required=['a_complex', 'b_complex'], ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -2406,86 +2354,15 @@ def func_under_test() -> MyComplexPydanticModel: client=vertex_client, callable=func_under_test ) - assert actual_schema_mldev == expected_schema_mldev - assert actual_schema_vertex == expected_schema_vertex - - -def test_function_with_return_type(): - def func_under_test1() -> set: - pass - - def func_under_test2() -> frozenset[int]: - pass - - def func_under_test3() -> typing.Set[int]: - pass - - def func_under_test4() -> typing.FrozenSet[int]: - pass - - def func_under_test5() -> typing.Iterable[int]: - pass - - def func_under_test6() -> bytes: - pass - - def func_under_test7() -> typing.OrderedDict[str, int]: - pass - - def func_under_test8() -> typing.MutableMapping[str, int]: - pass - - def func_under_test9() -> typing.MutableSequence[int]: - pass - - def func_under_test10() -> typing.MutableSet[int]: - pass - - def func_under_test11() -> typing.Counter[int]: - pass - - all_func_under_test = [ - func_under_test1, - func_under_test2, - func_under_test3, - func_under_test4, - func_under_test5, - func_under_test6, - func_under_test7, - func_under_test8, - func_under_test9, - func_under_test10, - func_under_test11, - ] - for i, func_under_test in enumerate(all_func_under_test): - - expected_schema_mldev = types.FunctionDeclaration( - name=f'func_under_test{i+1}', - description=None, - ) - actual_schema_mldev = types.FunctionDeclaration.from_callable( - client=mldev_client, callable=func_under_test - ) - assert actual_schema_mldev == expected_schema_mldev - - types.FunctionDeclaration.from_callable( - client=vertex_client, callable=func_under_test - ) + assert actual_schema_mldev == expected_schema + assert actual_schema_vertex == expected_schema def test_function_with_tuple_return_type(): def func_under_test() -> tuple[int, str, str]: pass - expected_schema_mldev = types.FunctionDeclaration( - name=f'func_under_test', - description=None, - ) - actual_schema_mldev = types.FunctionDeclaration.from_callable( - client=mldev_client, callable=func_under_test - ) - - expected_schema_vertex = types.FunctionDeclaration( + expected_schema = types.FunctionDeclaration( name=f'func_under_test', description=None, response_json_schema={ @@ -2503,8 +2380,12 @@ def func_under_test() -> tuple[int, str, str]: actual_schema_vertex = types.FunctionDeclaration.from_callable( client=vertex_client, callable=func_under_test ) - assert actual_schema_mldev == expected_schema_mldev - assert actual_schema_vertex == expected_schema_vertex + actual_schema_mldev = types.FunctionDeclaration.from_callable( + client=mldev_client, callable=func_under_test + ) + + assert actual_schema_mldev == expected_schema + assert actual_schema_vertex == expected_schema def test_function_with_return_type_not_supported(): @@ -2532,33 +2413,41 @@ def func_under_test4() -> MyClass: ] for i, func_under_test in enumerate(all_func_under_test): - expected_schema_mldev = types.FunctionDeclaration( - name=f'func_under_test{i+1}', - description=None, - ) - actual_schema_mldev = types.FunctionDeclaration.from_callable( - client=mldev_client, callable=func_under_test - ) - assert actual_schema_mldev == expected_schema_mldev with pytest.raises(ValueError): types.FunctionDeclaration.from_callable( client=vertex_client, callable=func_under_test ) + with pytest.raises(ValueError): + types.FunctionDeclaration.from_callable( + client=mldev_client, callable=func_under_test + ) + def test_function_with_tuple_contains_unevaluated_items(): def func_under_test(a: tuple[int, int]) -> str: """test return type.""" return '' - expected_parameters_json_schema = { - 'a': { - 'maxItems': 2, - 'minItems': 2, - 'prefixItems': [{'type': 'integer'}, {'type': 'integer'}], - 'type': 'array', - 'unevaluatedItems': False, - } - } + expected_schema = types.FunctionDeclaration( + name='func_under_test', + description='test return type.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': { + 'maxItems': 2, + 'minItems': 2, + 'prefixItems': [{'type': 'integer'}, {'type': 'integer'}], + 'type': 'array', + 'unevaluatedItems': False, + } + }, + 'required': ['a'], + }, + response_json_schema={ + 'type': 'string', + }, + ) actual_schema_mldev = types.FunctionDeclaration.from_callable( client=mldev_client, callable=func_under_test @@ -2567,8 +2456,8 @@ def func_under_test(a: tuple[int, int]) -> str: client=vertex_client, callable=func_under_test ) - assert actual_schema_mldev.parameters_json_schema == expected_parameters_json_schema - assert actual_schema_vertex.parameters_json_schema == expected_parameters_json_schema + assert actual_schema_mldev == expected_schema + assert actual_schema_vertex == expected_schema def test_function_gemini_api(monkeypatch): @@ -2581,14 +2470,17 @@ def func_under_test(a: int) -> str: expected_schema_mldev = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema(type='INTEGER'), - }, - required=['a'], - ), description='test return type.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': {'type': 'integer'}, + }, + 'required': ['a'], + }, + response_json_schema={ + 'type': 'string', + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable( @@ -2606,14 +2498,17 @@ def func_under_test(a: int) -> str: expected_schema_mldev = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema(type='INTEGER'), - }, - required=['a'], - ), description='test return type.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': {'type': 'integer'}, + }, + 'required': ['a'], + }, + response_json_schema={ + 'type': 'string', + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable_with_api_option( @@ -2631,14 +2526,17 @@ def func_under_test(a: int) -> str: expected_schema_mldev = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema(type='INTEGER'), - }, - required=['a'], - ), description='test return type.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': {'type': 'integer'}, + }, + 'required': ['a'], + }, + response_json_schema={ + 'type': 'string', + }, ) actual_schema_mldev = types.FunctionDeclaration.from_callable_with_api_option( @@ -2668,23 +2566,24 @@ def func_under_test(a: int) -> str: expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema(type='INTEGER'), - }, - ), description='test return type.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': {'type': 'integer'}, + }, + 'required': ['a'], + }, + response_json_schema={ + 'type': 'string', + }, ) - expected_schema_vertex = copy.deepcopy(expected_schema) - expected_schema_vertex.response = types.Schema(type='STRING') - expected_schema_vertex.parameters.required = ['a'] actual_schema_vertex = types.FunctionDeclaration.from_callable( client=vertex_client, callable=func_under_test ) - assert actual_schema_vertex == expected_schema_vertex + assert actual_schema_vertex == expected_schema def test_function_with_option_vertex(monkeypatch): @@ -2695,49 +2594,26 @@ def func_under_test(a: int) -> str: expected_schema = types.FunctionDeclaration( name='func_under_test', - parameters=types.Schema( - type='OBJECT', - properties={ - 'a': types.Schema(type='INTEGER'), - }, - ), description='test return type.', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'a': {'type': 'integer'}, + }, + 'required': ['a'], + }, + response_json_schema={ + 'type': 'string', + }, ) - expected_schema_vertex = copy.deepcopy(expected_schema) - expected_schema_vertex.response = types.Schema(type='STRING') - expected_schema_vertex.parameters.required = ['a'] actual_schema_vertex = ( types.FunctionDeclaration.from_callable_with_api_option( - callable=func_under_test, api_option='VERTEX_AI' + callable=func_under_test, api_option='ENTERPRISE' ) ) - assert actual_schema_vertex == expected_schema_vertex - - -def test_convert_json_schema_with_cycle(): - json_schema_dict = { - 'type': 'object', - 'properties': { - 'foo': {'$ref': '#/$defs/Foo'} - }, - '$defs': { - 'Foo': { - 'type': 'object', - 'properties': { - 'foo': {'$ref': '#/$defs/Foo'} - } - } - } - } - - json_schema = types.JSONSchema(**json_schema_dict) - schema = types.Schema.from_json_schema(json_schema=json_schema) - - assert schema.type == types.Type.OBJECT - assert schema.properties['foo'].type == types.Type.OBJECT - assert schema.properties['foo'].properties['foo'] == types.Schema() + assert actual_schema_vertex == expected_schema def test_case_insensitive_enum(): @@ -2969,4 +2845,3 @@ def test_computer_use_types(): assert c.enable_prompt_injection_detection is True assert len(c.disabled_safety_policies) == 2 assert types.SafetyPolicy.FINANCIAL_TRANSACTIONS in c.disabled_safety_policies - diff --git a/google/genai/types.py b/google/genai/types.py index 359c00d45..b1268efcd 100644 --- a/google/genai/types.py +++ b/google/genai/types.py @@ -1040,22 +1040,6 @@ class ResourceScope(_common.CaseInSensitiveEnum): "https://aiplatform.googleapis.com/publishers/google/models/gemini-3-pro-preview""" -class JSONSchemaType(Enum): - """The type of the data supported by JSON Schema. - - The values of the enums are lower case strings, while the values of the enums - for the Type class are upper case strings. - """ - - NULL = 'null' - BOOLEAN = 'boolean' - OBJECT = 'object' - ARRAY = 'array' - NUMBER = 'number' - INTEGER = 'integer' - STRING = 'string' - - class FeatureSelectionPreference(_common.CaseInSensitiveEnum): """Options for feature selection preference.""" @@ -2674,173 +2658,6 @@ class HttpOptionsDict(TypedDict, total=False): HttpOptionsOrDict = Union[HttpOptions, HttpOptionsDict] -class JSONSchema(_common.BaseModel): - """A subset of JSON Schema according to 2020-12 JSON Schema draft. - - Represents a subset of a JSON Schema object that is used by the Gemini model. - The difference between this class and the Schema class is that this class is - compatible with OpenAPI 3.1 schema objects. And the Schema class is used to - make API call to Gemini model. - """ - - type: Optional[Union[JSONSchemaType, list[JSONSchemaType]]] = Field( - default=None, - description="""Validation succeeds if the type of the instance matches the type represented by the given type, or matches at least one of the given types.""", - ) - format: Optional[str] = Field( - default=None, - description='Define semantic information about a string instance.', - ) - title: Optional[str] = Field( - default=None, - description=( - 'A preferably short description about the purpose of the instance' - ' described by the schema.' - ), - ) - description: Optional[str] = Field( - default=None, - description=( - 'An explanation about the purpose of the instance described by the' - ' schema.' - ), - ) - default: Optional[Any] = Field( - default=None, - description=( - 'This keyword can be used to supply a default JSON value associated' - ' with a particular schema.' - ), - ) - items: Optional['JSONSchema'] = Field( - default=None, - description=( - 'Validation succeeds if each element of the instance not covered by' - ' prefixItems validates against this schema.' - ), - ) - min_items: Optional[int] = Field( - default=None, - description=( - 'An array instance is valid if its size is greater than, or equal to,' - ' the value of this keyword.' - ), - ) - max_items: Optional[int] = Field( - default=None, - description=( - 'An array instance is valid if its size is less than, or equal to,' - ' the value of this keyword.' - ), - ) - enum: Optional[list[Any]] = Field( - default=None, - description=( - 'Validation succeeds if the instance is equal to one of the elements' - ' in this keyword’s array value.' - ), - ) - properties: Optional[dict[str, 'JSONSchema']] = Field( - default=None, - description=( - 'Validation succeeds if, for each name that appears in both the' - ' instance and as a name within this keyword’s value, the child' - ' instance for that name successfully validates against the' - ' corresponding schema.' - ), - ) - required: Optional[list[str]] = Field( - default=None, - description=( - 'An object instance is valid against this keyword if every item in' - ' the array is the name of a property in the instance.' - ), - ) - min_properties: Optional[int] = Field( - default=None, - description=( - 'An object instance is valid if its number of properties is greater' - ' than, or equal to, the value of this keyword.' - ), - ) - max_properties: Optional[int] = Field( - default=None, - description=( - 'An object instance is valid if its number of properties is less' - ' than, or equal to, the value of this keyword.' - ), - ) - minimum: Optional[float] = Field( - default=None, - description=( - 'Validation succeeds if the numeric instance is greater than or equal' - ' to the given number.' - ), - ) - maximum: Optional[float] = Field( - default=None, - description=( - 'Validation succeeds if the numeric instance is less than or equal to' - ' the given number.' - ), - ) - min_length: Optional[int] = Field( - default=None, - description=( - 'A string instance is valid against this keyword if its length is' - ' greater than, or equal to, the value of this keyword.' - ), - ) - max_length: Optional[int] = Field( - default=None, - description=( - 'A string instance is valid against this keyword if its length is' - ' less than, or equal to, the value of this keyword.' - ), - ) - pattern: Optional[str] = Field( - default=None, - description=( - 'A string instance is considered valid if the regular expression' - ' matches the instance successfully.' - ), - ) - additional_properties: Optional[Any] = Field( - default=None, - description="""Can either be a boolean or an object; controls the presence of additional properties.""", - ) - any_of: Optional[list['JSONSchema']] = Field( - default=None, - description=( - 'An instance validates successfully against this keyword if it' - ' validates successfully against at least one schema defined by this' - ' keyword’s value.' - ), - ) - unique_items: Optional[bool] = Field( - default=None, - description="""Boolean value that indicates whether the items in an array are unique.""", - ) - ref: Optional[str] = Field( - default=None, - alias='$ref', - description="""Allows indirect references between schema nodes.""", - ) - defs: Optional[dict[str, 'JSONSchema']] = Field( - default=None, - alias='$defs', - description="""Schema definitions to be used with $ref.""", - ) - one_of: Optional[list['JSONSchema']] = Field( - default=None, - description=( - 'An instance validates successfully against this keyword if it' - ' validates successfully against exactly one schema defined by this' - " keyword's value." - ), - ) - - class Schema(_common.BaseModel): """Schema is used to define the format of input/output data. @@ -2948,482 +2765,6 @@ class Schema(_common.BaseModel): default=None, description="""Optional. Data type of the schema field.""" ) - @property - def json_schema(self) -> 'JSONSchema': - """Converts the Schema object to a JSONSchema object, that is compatible with 2020-12 JSON Schema draft. - - Note: Conversion of fields that are not included in the JSONSchema class - are ignored. - Json Schema is now supported natively by both Gemini Enterprise Agent - Platform and Gemini API. Users - are recommended to pass/receive Json Schema directly to/from the API. For - example: - 1. the counter part of GenerateContentConfig.response_schema is - GenerateContentConfig.response_json_schema, which accepts [JSON - Schema](https://json-schema.org/) - 2. the counter part of FunctionDeclaration.parameters is - FunctionDeclaration.parameters_json_schema, which accepts [JSON - Schema](https://json-schema.org/) - 3. the counter part of FunctionDeclaration.response is - FunctionDeclaration.response_json_schema, which accepts [JSON - Schema](https://json-schema.org/) - """ - - global _json_schema_warning_logged - if not _json_schema_warning_logged: - info_message = """ -Note: Conversion of fields that are not included in the JSONSchema class are -ignored. -Json Schema is now supported natively by both Gemini Enterprise Agent Platform and Gemini API. Users -are recommended to pass/receive Json Schema directly to/from the API. For example: -1. the counter part of GenerateContentConfig.response_schema is - GenerateContentConfig.response_json_schema, which accepts [JSON - Schema](https://json-schema.org/) -2. the counter part of FunctionDeclaration.parameters is - FunctionDeclaration.parameters_json_schema, which accepts [JSON - Schema](https://json-schema.org/) -3. the counter part of FunctionDeclaration.response is - FunctionDeclaration.response_json_schema, which accepts [JSON - Schema](https://json-schema.org/) -""" - logger.info(info_message) - _json_schema_warning_logged = True - - json_schema_field_names: set[str] = set(JSONSchema.model_fields.keys()) - schema_field_names: tuple[str] = ( - 'items', - ) # 'additional_properties' to come - list_schema_field_names: tuple[str] = ( - 'any_of', # 'one_of', 'all_of', 'not' to come - ) - dict_schema_field_names: tuple[str] = ('properties',) # 'defs' to come - - def convert_schema(schema: Union['Schema', dict[str, Any]]) -> 'JSONSchema': - if isinstance(schema, pydantic.BaseModel): - schema_dict = schema.model_dump(exclude_none=True) - else: - schema_dict = schema - json_schema = JSONSchema() - for field_name, field_value in schema_dict.items(): - if field_value is None: - continue - elif field_name == 'nullable': - if json_schema.type is None: - json_schema.type = JSONSchemaType.NULL - elif isinstance(json_schema.type, JSONSchemaType): - current_type: JSONSchemaType = json_schema.type - json_schema.type = [current_type, JSONSchemaType.NULL] - elif isinstance(json_schema.type, list): - json_schema.type.append(JSONSchemaType.NULL) - elif field_name not in json_schema_field_names: - continue - elif field_name == 'type': - if field_value == Type.TYPE_UNSPECIFIED: - continue - json_schema_type = JSONSchemaType(field_value.lower()) - if json_schema.type is None: - json_schema.type = json_schema_type - elif isinstance(json_schema.type, JSONSchemaType): - existing_type: JSONSchemaType = json_schema.type - json_schema.type = [existing_type, json_schema_type] - elif isinstance(json_schema.type, list): - json_schema.type.append(json_schema_type) - elif field_name in schema_field_names: - schema_field_value: 'JSONSchema' = convert_schema(field_value) - setattr(json_schema, field_name, schema_field_value) - elif field_name in list_schema_field_names: - list_schema_field_value: list['JSONSchema'] = [ - convert_schema(this_field_value) - for this_field_value in field_value - ] - setattr(json_schema, field_name, list_schema_field_value) - elif field_name in dict_schema_field_names: - dict_schema_field_value: dict[str, 'JSONSchema'] = { - key: convert_schema(value) for key, value in field_value.items() - } - setattr(json_schema, field_name, dict_schema_field_value) - else: - setattr(json_schema, field_name, field_value) - - return json_schema - - return convert_schema(self) - - @classmethod - def from_json_schema( - cls, - *, - json_schema: 'JSONSchema', - api_option: Literal['VERTEX_AI', 'GEMINI_API'] = 'GEMINI_API', - raise_error_on_unsupported_field: bool = False, - ) -> 'Schema': - """Converts a JSONSchema object to a Schema object. - - Note: Conversion of fields that are not included in the JSONSchema class - are ignored. - Json Schema is now supported natively by both Gemini Enterprise Agent - Platform and Gemini API. Users - are recommended to pass/receive Json Schema directly to/from the API. For - example: - 1. the counter part of GenerateContentConfig.response_schema is - GenerateContentConfig.response_json_schema, which accepts [JSON - Schema](https://json-schema.org/) - 2. the counter part of FunctionDeclaration.parameters is - FunctionDeclaration.parameters_json_schema, which accepts [JSON - Schema](https://json-schema.org/) - 3. the counter part of FunctionDeclaration.response is - FunctionDeclaration.response_json_schema, which accepts [JSON - Schema](https://json-schema.org/) - The JSONSchema is compatible with 2020-12 JSON Schema draft, specified by - OpenAPI 3.1. - - Args: - json_schema: JSONSchema object to be converted. - api_option: API option to be used. If set to 'VERTEX_AI', the - JSONSchema will be converted to a Schema object that is compatible - with Gemini Enterprise Agent Platform API. If set to 'GEMINI_API', - the JSONSchema will be converted to a Schema object that is - compatible with Gemini API. Default is 'GEMINI_API'. - raise_error_on_unsupported_field: If set to True, an error will be - raised if the JSONSchema contains any unsupported fields. Default is - False. - - Returns: - Schema object that is compatible with the specified API option. - Raises: - ValueError: If the JSONSchema contains any unsupported fields and - raise_error_on_unsupported_field is set to True. Or if the JSONSchema - is not compatible with the specified API option. - """ - global _from_json_schema_warning_logged - if not _from_json_schema_warning_logged: - info_message = """ -Note: Conversion of fields that are not included in the JSONSchema class are ignored. -Json Schema is now supported natively by both Gemini Enterprise Agent Platform and Gemini API. Users -are recommended to pass/receive Json Schema directly to/from the API. For example: -1. the counter part of GenerateContentConfig.response_schema is - GenerateContentConfig.response_json_schema, which accepts [JSON - Schema](https://json-schema.org/) -2. the counter part of FunctionDeclaration.parameters is - FunctionDeclaration.parameters_json_schema, which accepts [JSON - Schema](https://json-schema.org/) -3. the counter part of FunctionDeclaration.response is - FunctionDeclaration.response_json_schema, which accepts [JSON - Schema](https://json-schema.org/) -""" - logger.info(info_message) - _from_json_schema_warning_logged = True - - google_schema_field_names: set[str] = set(cls.model_fields.keys()) - schema_field_names: tuple[str, ...] = ( - 'items', - ) # 'additional_properties' to come - list_schema_field_names: tuple[str, ...] = ( - 'any_of', # 'one_of', 'all_of', 'not' to come - ) - dict_schema_field_names: tuple[str, ...] = ('properties',) - - related_field_names_by_type: dict[str, tuple[str, ...]] = { - JSONSchemaType.NUMBER.value: ( - 'description', - 'enum', - 'format', - 'maximum', - 'minimum', - 'title', - ), - JSONSchemaType.STRING.value: ( - 'description', - 'enum', - 'format', - 'max_length', - 'min_length', - 'pattern', - 'title', - ), - JSONSchemaType.OBJECT.value: ( - 'any_of', - 'description', - 'max_properties', - 'min_properties', - 'properties', - 'required', - 'title', - ), - JSONSchemaType.ARRAY.value: ( - 'description', - 'items', - 'max_items', - 'min_items', - 'title', - ), - JSONSchemaType.BOOLEAN.value: ( - 'description', - 'title', - ), - } - # Treat `INTEGER` like `NUMBER`. - related_field_names_by_type[JSONSchemaType.INTEGER.value] = ( - related_field_names_by_type[JSONSchemaType.NUMBER.value] - ) - - # placeholder for potential gemini api unsupported fields - gemini_api_unsupported_field_names: tuple[str, ...] = () - - def _resolve_ref( - ref_path: str, root_schema_dict: dict[str, Any] - ) -> dict[str, Any]: - """Helper to resolve a $ref path.""" - current = root_schema_dict - for part in ref_path.lstrip('#/').split('/'): - if part == '$defs': - part = 'defs' - current = current[part] - current.pop('title', None) - if 'properties' in current and current['properties'] is not None: - for prop_schema in current['properties'].values(): - if isinstance(prop_schema, dict): - prop_schema.pop('title', None) - - return current - - def normalize_json_schema_type( - json_schema_type: Optional[ - Union[JSONSchemaType, Sequence[JSONSchemaType], str, Sequence[str]] - ], - ) -> tuple[list[str], bool]: - """Returns (non_null_types, nullable)""" - if json_schema_type is None: - return [], False - type_sequence: Sequence[Union[JSONSchemaType, str]] - if isinstance(json_schema_type, str) or not isinstance( - json_schema_type, Sequence - ): - type_sequence = [json_schema_type] - else: - type_sequence = json_schema_type - non_null_types = [] - nullable = False - for type_value in type_sequence: - if isinstance(type_value, JSONSchemaType): - type_value = type_value.value - if type_value == JSONSchemaType.NULL.value: - nullable = True - else: - non_null_types.append(type_value) - return non_null_types, nullable - - def raise_error_if_cannot_convert( - json_schema_dict: dict[str, Any], - api_option: Literal['VERTEX_AI', 'GEMINI_API'], - raise_error_on_unsupported_field: bool, - ) -> None: - """Raises an error if the JSONSchema cannot be converted to the specified Schema object.""" - if not raise_error_on_unsupported_field: - return - for field_name, field_value in json_schema_dict.items(): - if field_value is None: - continue - if field_name not in google_schema_field_names and field_name not in [ - 'ref', - 'defs', - ]: - raise ValueError( - f'JSONSchema field "{field_name}" is not supported by the Schema' - ' object. And the "raise_error_on_unsupported_field" argument is' - ' set to True. If you still want to convert it into the Schema' - f' object, please either remove the field "{field_name}" from the' - ' JSONSchema object, leave the' - ' "raise_error_on_unsupported_field" unset, or try using' - ' response_json_schema instead.' - ) - if ( - field_name in gemini_api_unsupported_field_names - and api_option == 'GEMINI_API' - ): - raise ValueError( - f'The "{field_name}" field is not supported by the Schema ' - 'object for GEMINI_API.' - ) - - def copy_schema_fields( - json_schema_dict: dict[str, Any], - related_fields_to_copy: tuple[str, ...], - sub_schema_in_any_of: dict[str, Any], - ) -> None: - """Copies the fields from json_schema_dict to sub_schema_in_any_of.""" - for field_name in related_fields_to_copy: - sub_schema_in_any_of[field_name] = json_schema_dict.get( - field_name, None - ) - - def convert_json_schema( - current_json_schema: 'JSONSchema', - root_json_schema_dict: dict[str, Any], - api_option: Literal['VERTEX_AI', 'GEMINI_API'], - raise_error_on_unsupported_field: bool, - visited_refs: Optional[set[str]] = None, - ) -> 'Schema': - if visited_refs is None: - visited_refs = set() - - schema = Schema() - json_schema_dict = current_json_schema.model_dump() - - ref = json_schema_dict.get('ref') - if ref: - if ref in visited_refs: - return Schema() - visited_refs.add(ref) - json_schema_dict = _resolve_ref(ref, root_json_schema_dict) - - raise_error_if_cannot_convert( - json_schema_dict=json_schema_dict, - api_option=api_option, - raise_error_on_unsupported_field=raise_error_on_unsupported_field, - ) - - # At the highest level of the logic, there are two passes: - # Pass 1: the JSONSchema.type is union-like, - # e.g. ['null', 'string', 'array']. - # for this case, we need to split the JSONSchema into multiple - # sub-schemas, and copy them into the any_of field of the Schema. - # And when we copy the non-type fields into any_of field, - # we only copy the fields related to the specific type. - # Detailed logic is commented below with `Pass 1` keyword tag. - # Pass 2: the JSONSchema.type is not union-like, - # e.g. 'string', ['string'], ['null', 'string']. - # for this case, no splitting is needed. Detailed - # logic is commented below with `Pass 2` keyword tag. - # - # - # Pass 1: the JSONSchema.type is union-like - # e.g. ['null', 'string', 'array']. - non_null_types, nullable = normalize_json_schema_type( - json_schema_dict.get('type', None) - ) - is_union_like_type = len(non_null_types) > 1 - if len(non_null_types) > 1: - logger.warning( - 'JSONSchema type is union-like, e.g. ["null", "string", "array"]. ' - 'Converting it into multiple sub-schemas, and copying them into ' - 'the any_of field of the Schema. The value of `default` field is ' - 'ignored because it is ambiguous to tell which sub-schema it ' - 'belongs to.' - ) - reformed_json_schema = JSONSchema() - # start splitting the JSONSchema into multiple sub-schemas - any_of = [] - if nullable: - schema.nullable = True - for normalized_type in non_null_types: - sub_schema_in_any_of = {'type': normalized_type} - related_field_names = related_field_names_by_type.get(normalized_type) - if related_field_names is not None: - copy_schema_fields( - json_schema_dict=json_schema_dict, - related_fields_to_copy=related_field_names, - sub_schema_in_any_of=sub_schema_in_any_of, - ) - any_of.append(JSONSchema(**sub_schema_in_any_of)) - reformed_json_schema.any_of = any_of - json_schema_dict = reformed_json_schema.model_dump() - - # Pass 2: the JSONSchema.type is not union-like, - # e.g. 'string', ['string'], ['null', 'string']. - for field_name, field_value in json_schema_dict.items(): - if field_value is None or field_name == 'defs': - continue - if field_name in schema_field_names: - if field_name == 'items' and not field_value: - continue - schema_field_value: 'Schema' = convert_json_schema( - current_json_schema=JSONSchema(**field_value), - root_json_schema_dict=root_json_schema_dict, - api_option=api_option, - raise_error_on_unsupported_field=raise_error_on_unsupported_field, - visited_refs=visited_refs, - ) - setattr(schema, field_name, schema_field_value) - elif field_name in list_schema_field_names: - list_schema_field_value: list['Schema'] = [ - convert_json_schema( - current_json_schema=JSONSchema(**this_field_value), - root_json_schema_dict=root_json_schema_dict, - api_option=api_option, - raise_error_on_unsupported_field=raise_error_on_unsupported_field, - visited_refs=visited_refs, - ) - for this_field_value in field_value - ] - setattr(schema, field_name, list_schema_field_value) - if not schema.type and not is_union_like_type and not schema.any_of: - schema.type = Type('OBJECT') - elif field_name in dict_schema_field_names: - dict_schema_field_value: dict[str, 'Schema'] = { - key: convert_json_schema( - current_json_schema=JSONSchema(**value), - root_json_schema_dict=root_json_schema_dict, - api_option=api_option, - raise_error_on_unsupported_field=raise_error_on_unsupported_field, - visited_refs=visited_refs, - ) - for key, value in field_value.items() - } - setattr(schema, field_name, dict_schema_field_value) - elif field_name == 'type': - non_null_types, nullable = normalize_json_schema_type(field_value) - if nullable: - schema.nullable = True - if non_null_types: - schema.type = Type(non_null_types[0]) - else: - if ( - hasattr(schema, field_name) - and field_name != 'additional_properties' - ): - setattr(schema, field_name, field_value) - - if ( - schema.type == 'ARRAY' - and schema.items - and not schema.items.model_dump(exclude_unset=True) - ): - schema.items = None - - if schema.any_of and len(schema.any_of) == 2: - nullable_part = None - type_part = None - for part in schema.any_of: - # A schema representing `None` will either be of type NULL or just be nullable. - part_dict = part.model_dump(exclude_unset=True) - if part_dict == {'nullable': True} or part_dict == {'type': 'NULL'}: - nullable_part = part - else: - type_part = part - - # If we found both parts, unwrap them into a single schema. - if nullable_part and type_part: - default_value = schema.default - schema = type_part - schema.nullable = True - # Carry the default value over to the unwrapped schema - if default_value is not None: - schema.default = default_value - - if ref: - visited_refs.remove(ref) - return schema - - # This is the initial call to the recursive function. - root_schema_dict = json_schema.model_dump() - return convert_json_schema( - current_json_schema=json_schema, - root_json_schema_dict=root_schema_dict, - api_option=api_option, - raise_error_on_unsupported_field=raise_error_on_unsupported_field, - ) - class SchemaDict(TypedDict, total=False): """Schema is used to define the format of input/output data. @@ -4768,54 +4109,44 @@ def from_callable_with_api_option( cls, *, callable: Callable[..., Any], - api_option: Literal['VERTEX_AI', 'GEMINI_API'] = 'GEMINI_API', + api_option: Literal['ENTERPRISE', 'GEMINI_API'] = 'GEMINI_API', behavior: Optional[Behavior] = None, ) -> 'FunctionDeclaration': """Converts a Callable to a FunctionDeclaration based on the API option. - Supported API option is 'VERTEX_AI' or 'GEMINI_API'. If api_option is unset, + Supported API option is 'ENTERPRISE' or 'GEMINI_API'. If api_option is + unset, it will default to 'GEMINI_API'. If unsupported api_option is provided, it will raise ValueError. """ - supported_api_options = ['VERTEX_AI', 'GEMINI_API'] + supported_api_options = ['ENTERPRISE', 'GEMINI_API'] if api_option not in supported_api_options: raise ValueError( f'Unsupported api_option value: {api_option}. Supported api_option' f' value is one of: {supported_api_options}.' ) + from . import _automatic_function_calling_util + from . import _extra_utils - parameters_properties = {} - parameters_json_schema = {} annotation_under_future = typing.get_type_hints(callable) - try: - for name, param in inspect.signature(callable).parameters.items(): - if param.kind in ( - inspect.Parameter.POSITIONAL_OR_KEYWORD, - inspect.Parameter.KEYWORD_ONLY, - inspect.Parameter.POSITIONAL_ONLY, - ): + parameters_properties_json_schema = {} + root_defs = {} + + for name, param in inspect.signature(callable).parameters.items(): + if param.kind in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + inspect.Parameter.POSITIONAL_ONLY, + ): + try: param = _automatic_function_calling_util._handle_params_as_deferred_annotations( param, annotation_under_future, name ) - schema = ( - _automatic_function_calling_util._parse_schema_from_parameter( - api_option, param, callable.__name__ - ) - ) - parameters_properties[name] = schema - except ValueError: - parameters_properties = {} - for name, param in inspect.signature(callable).parameters.items(): - if param.kind in ( - inspect.Parameter.POSITIONAL_OR_KEYWORD, - inspect.Parameter.KEYWORD_ONLY, - inspect.Parameter.POSITIONAL_ONLY, - ): - try: - param = _automatic_function_calling_util._handle_params_as_deferred_annotations( - param, annotation_under_future, name - ) + json_schema_dict = {} + if _extra_utils.is_annotation_pydantic_model(param.annotation): + json_schema_dict = param.annotation.model_json_schema() + else: param_schema_adapter = pydantic.TypeAdapter( param.annotation, config=pydantic.ConfigDict(arbitrary_types_allowed=True), @@ -4824,36 +4155,23 @@ def from_callable_with_api_option( json_schema_dict = _automatic_function_calling_util._add_unevaluated_items_to_fixed_len_tuple_schema( json_schema_dict ) - if 'prefixItems' in json_schema_dict: - parameters_json_schema[name] = json_schema_dict - continue - - union_args = typing.get_args(param.annotation) - has_primitive = any( - _automatic_function_calling_util._is_builtin_primitive_or_compound( - arg - ) - for arg in union_args - ) - if ( - '$ref' in json_schema_dict or '$defs' in json_schema_dict - ) and has_primitive: - # This is a complex schema with a primitive (e.g., str | MyModel) - # that is better represented by raw JSON schema. - parameters_json_schema[name] = json_schema_dict - continue - - schema = Schema.from_json_schema( - json_schema=JSONSchema(**json_schema_dict), - api_option=api_option, - ) - if param.default is not inspect.Parameter.empty: - schema.default = param.default - parameters_properties[name] = schema - except Exception as e: - _automatic_function_calling_util._raise_for_unsupported_param( - param, callable.__name__, e - ) + + # Extract parameter-level $defs and promote to top-level root_defs + if '$defs' in json_schema_dict: + root_defs.update(json_schema_dict.pop('$defs')) + if 'definitions' in json_schema_dict: + root_defs.update(json_schema_dict.pop('definitions')) + # pydantic doesn't assign the `type` field when the schema has 'anyOf'. + # but Vertex requires it. + if not 'type' in json_schema_dict and 'anyOf' in json_schema_dict: + json_schema_dict['type'] = 'object' + if param.default is not inspect._empty: + json_schema_dict['default'] = param.default + parameters_properties_json_schema[name] = json_schema_dict + except Exception as e: + _automatic_function_calling_util._raise_for_unsupported_param( + param, callable.__name__, e + ) declaration = FunctionDeclaration( name=callable.__name__, @@ -4862,21 +4180,18 @@ def from_callable_with_api_option( else callable.__doc__, behavior=behavior, ) - if parameters_properties: - declaration.parameters = Schema( - type='OBJECT', - properties=parameters_properties, - ) - declaration.parameters.required = ( + if parameters_properties_json_schema: + declaration.parameters_json_schema = { + 'type': 'object', + 'properties': parameters_properties_json_schema, + } + if root_defs: + declaration.parameters_json_schema['$defs'] = root_defs + declaration.parameters_json_schema['required'] = ( _automatic_function_calling_util._get_required_fields( - declaration.parameters + declaration.parameters_json_schema ) ) - elif parameters_json_schema: - declaration.parameters_json_schema = parameters_json_schema - # TODO: b/421991354 - Remove this check once the bug is fixed. - if api_option == 'GEMINI_API': - return declaration return_annotation = inspect.signature(callable).return_annotation if return_annotation is inspect._empty: @@ -4893,39 +4208,29 @@ def from_callable_with_api_option( return_value = return_value.replace( annotation=annotation_under_future['return'] ) - response_schema: Optional[Schema] = None - response_json_schema: Optional[Union[dict[str, Any], Schema]] = {} + response_json_schema: dict[str, Any] = {} try: - response_schema = ( - _automatic_function_calling_util._parse_schema_from_parameter( - api_option, - return_value, - callable.__name__, - ) - ) - if response_schema.any_of is not None: - # To handle any_of, we need to use responseJsonSchema - response_json_schema = response_schema - response_schema = None - except ValueError: - try: + if _extra_utils.is_annotation_pydantic_model(return_value.annotation): + response_json_schema = return_value.annotation.model_json_schema() + else: return_value_schema_adapter = pydantic.TypeAdapter( return_value.annotation, config=pydantic.ConfigDict(arbitrary_types_allowed=True), ) response_json_schema = return_value_schema_adapter.json_schema() - response_json_schema = _automatic_function_calling_util._add_unevaluated_items_to_fixed_len_tuple_schema( - response_json_schema - ) - except Exception as e: - _automatic_function_calling_util._raise_for_unsupported_param( - return_value, callable.__name__, e - ) + response_json_schema = _automatic_function_calling_util._add_unevaluated_items_to_fixed_len_tuple_schema( + response_json_schema + ) + # pydantic doesn't assign the `type` field when the schema has 'anyOf'. + # but Vertex requires it. + if not 'type' in response_json_schema and 'anyOf' in response_json_schema: + response_json_schema['type'] = 'object' + except Exception as e: + _automatic_function_calling_util._raise_for_unsupported_param( + return_value, callable.__name__, e + ) - if response_schema: - declaration.response = response_schema - elif response_json_schema: - declaration.response_json_schema = response_json_schema + declaration.response_json_schema = response_json_schema return declaration @classmethod @@ -4948,7 +4253,7 @@ def from_callable( """ if client.vertexai: return cls.from_callable_with_api_option( - callable=callable, api_option='VERTEX_AI', behavior=behavior + callable=callable, api_option='ENTERPRISE', behavior=behavior ) else: return cls.from_callable_with_api_option( @@ -5565,65 +4870,6 @@ class SpeechConfigDict(TypedDict, total=False): SpeechConfigOrDict = Union[SpeechConfig, SpeechConfigDict] -class AutomaticFunctionCallingConfig(_common.BaseModel): - """The configuration for automatic function calling.""" - - disable: Optional[bool] = Field( - default=None, - description="""Whether to disable automatic function calling. - If not set or set to False, will enable automatic function calling. - If set to True, will disable automatic function calling. - """, - ) - maximum_remote_calls: Optional[int] = Field( - default=10, - description="""If automatic function calling is enabled, - maximum number of remote calls for automatic function calling. - This number should be a positive integer. - If not set, SDK will set maximum number of remote calls to 10. - """, - ) - ignore_call_history: Optional[bool] = Field( - default=None, - description="""If automatic function calling is enabled, - whether to ignore call history to the response. - If not set, SDK will set ignore_call_history to false, - and will append the call history to - GenerateContentResponse.automatic_function_calling_history. - """, - ) - - -class AutomaticFunctionCallingConfigDict(TypedDict, total=False): - """The configuration for automatic function calling.""" - - disable: Optional[bool] - """Whether to disable automatic function calling. - If not set or set to False, will enable automatic function calling. - If set to True, will disable automatic function calling. - """ - - maximum_remote_calls: Optional[int] - """If automatic function calling is enabled, - maximum number of remote calls for automatic function calling. - This number should be a positive integer. - If not set, SDK will set maximum number of remote calls to 10. - """ - - ignore_call_history: Optional[bool] - """If automatic function calling is enabled, - whether to ignore call history to the response. - If not set, SDK will set ignore_call_history to false, - and will append the call history to - GenerateContentResponse.automatic_function_calling_history. - """ - - -AutomaticFunctionCallingConfigOrDict = Union[ - AutomaticFunctionCallingConfig, AutomaticFunctionCallingConfigDict -] - - class ThinkingConfig(_common.BaseModel): """The thinking features configuration.""" @@ -6534,11 +5780,6 @@ class GenerateContentConfig(_common.BaseModel): model. """, ) - automatic_function_calling: Optional[AutomaticFunctionCallingConfig] = Field( - default=None, - description="""The configuration for automatic function calling. - """, - ) thinking_config: Optional[ThinkingConfig] = Field( default=None, description="""The thinking features configuration. @@ -6762,10 +6003,6 @@ class GenerateContentConfigDict(TypedDict, total=False): model. """ - automatic_function_calling: Optional[AutomaticFunctionCallingConfigDict] - """The configuration for automatic function calling. - """ - thinking_config: Optional[ThinkingConfigDict] """The thinking features configuration. """ @@ -22534,6 +21771,64 @@ class EmbedContentParametersDict(TypedDict, total=False): ] +class AutomaticFunctionCallingConfig(_common.BaseModel): + """The configuration for automatic function calling.""" + + maximum_remote_calls: Optional[int] = Field( + default=10, + description="""If automatic function calling is enabled, + maximum number of remote calls for automatic function calling. + This number should be a positive integer. + If not set, SDK will set maximum number of remote calls to 10. + """, + ) + enable: Optional[bool] = Field( + default=None, + description="""Whether to enable automatic function calling. + If not set or set to False, will not enable automatic function calling. + If set to True, will enable automatic function calling. + """, + ) + + +class AutomaticFunctionCallingConfigDict(TypedDict, total=False): + """The configuration for automatic function calling.""" + + maximum_remote_calls: Optional[int] + """If automatic function calling is enabled, + maximum number of remote calls for automatic function calling. + This number should be a positive integer. + If not set, SDK will set maximum number of remote calls to 10. + """ + + enable: Optional[bool] + """Whether to enable automatic function calling. + If not set or set to False, will not enable automatic function calling. + If set to True, will enable automatic function calling. + """ + + +AutomaticFunctionCallingConfigOrDict = Union[ + AutomaticFunctionCallingConfig, AutomaticFunctionCallingConfigDict +] + + +class ChatConfig(GenerateContentConfig): + """Configuration for chat. + + This is a sub class of `GenerateContentConfig`, it supports all the + configurations in `GenerateContentConfig`. + """ + + automatic_function_calling_config: Optional[ + AutomaticFunctionCallingConfig + ] = Field( + default=None, + description="""The configuration for automatic function calling. + """, + ) + + class UserContent(Content): """UserContent facilitates the creation of a Content object with a user role.