722 lines
26 KiB
Python
722 lines
26 KiB
Python
# 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.
|
|
#
|
|
|
|
"""Extra utils depending on types that are shared between sync and async modules."""
|
|
|
|
import asyncio
|
|
import inspect
|
|
import io
|
|
import logging
|
|
import sys
|
|
import typing
|
|
from typing import Any, Callable, Dict, Optional, Union, get_args, get_origin
|
|
import mimetypes
|
|
import os
|
|
import pydantic
|
|
|
|
from . import _common
|
|
from . import _mcp_utils
|
|
from . import _transformers as t
|
|
from . import errors
|
|
from . import types
|
|
from ._adapters import McpToGenAiToolAdapter
|
|
|
|
|
|
if sys.version_info >= (3, 10):
|
|
from types import UnionType
|
|
else:
|
|
UnionType = typing._UnionGenericAlias # type: ignore[attr-defined]
|
|
|
|
if typing.TYPE_CHECKING:
|
|
from mcp import ClientSession as McpClientSession
|
|
from mcp.types import Tool as McpTool
|
|
else:
|
|
McpClientSession: typing.Type = Any
|
|
McpTool: typing.Type = Any
|
|
try:
|
|
from mcp import ClientSession as McpClientSession
|
|
from mcp.types import Tool as McpTool
|
|
except ImportError:
|
|
McpClientSession = None
|
|
McpTool = None
|
|
|
|
_DEFAULT_MAX_REMOTE_CALLS_AFC = 10
|
|
|
|
logger = logging.getLogger('google_genai.models')
|
|
|
|
|
|
def _create_generate_content_config_model(
|
|
config: types.GenerateContentConfigOrDict,
|
|
) -> types.GenerateContentConfig:
|
|
if isinstance(config, dict):
|
|
return types.GenerateContentConfig(**config)
|
|
else:
|
|
return config
|
|
|
|
|
|
def _get_gcs_uri(
|
|
src: Union[str, types.BatchJobSourceOrDict]
|
|
) -> Optional[str]:
|
|
"""Extracts the first GCS URI from the source, if available."""
|
|
if isinstance(src, str) and src.startswith('gs://'):
|
|
return src
|
|
elif isinstance(src, dict) and src.get('gcs_uri'):
|
|
return src['gcs_uri'][0] if src['gcs_uri'] else None
|
|
elif isinstance(src, types.BatchJobSource) and src.gcs_uri:
|
|
return src.gcs_uri[0] if src.gcs_uri else None
|
|
return None
|
|
|
|
|
|
def _get_bigquery_uri(
|
|
src: Union[str, types.BatchJobSourceOrDict]
|
|
) -> Optional[str]:
|
|
"""Extracts the BigQuery URI from the source, if available."""
|
|
if isinstance(src, str) and src.startswith('bq://'):
|
|
return src
|
|
elif isinstance(src, dict) and src.get('bigquery_uri'):
|
|
return src['bigquery_uri']
|
|
elif isinstance(src, types.BatchJobSource) and src.bigquery_uri:
|
|
return src.bigquery_uri
|
|
return None
|
|
|
|
|
|
def format_destination(
|
|
src: Union[str, types.BatchJobSource],
|
|
config: Optional[types.CreateBatchJobConfig] = None,
|
|
) -> types.CreateBatchJobConfig:
|
|
"""Formats the destination uri based on the source uri for Vertex AI."""
|
|
if config is None:
|
|
config = types.CreateBatchJobConfig()
|
|
|
|
unique_name = None
|
|
if not config.display_name:
|
|
unique_name = _common.timestamped_unique_name()
|
|
config.display_name = f'genai_batch_job_{unique_name}'
|
|
|
|
if not config.dest:
|
|
gcs_source_uri = _get_gcs_uri(src)
|
|
bigquery_source_uri = _get_bigquery_uri(src)
|
|
|
|
if gcs_source_uri and gcs_source_uri.endswith('.jsonl'):
|
|
config.dest = f'{gcs_source_uri[:-6]}/dest'
|
|
elif bigquery_source_uri:
|
|
unique_name = unique_name or _common.timestamped_unique_name()
|
|
config.dest = f'{bigquery_source_uri}_dest_{unique_name}'
|
|
|
|
return config
|
|
|
|
|
|
def find_afc_incompatible_tool_indexes(
|
|
config: Optional[types.GenerateContentConfigOrDict] = None,
|
|
is_agent_platform: bool = False,
|
|
) -> list[int]:
|
|
"""Checks if the config contains any AFC incompatible tools."""
|
|
if not config:
|
|
return []
|
|
config_model = _create_generate_content_config_model(config)
|
|
incompatible_tools_indexes: list[int] = []
|
|
|
|
if not config_model or not config_model.tools:
|
|
return incompatible_tools_indexes
|
|
|
|
for index, tool in enumerate(config_model.tools):
|
|
if not isinstance(tool, types.Tool):
|
|
continue
|
|
if getattr(tool, 'function_declarations', None):
|
|
incompatible_tools_indexes.append(index)
|
|
|
|
# Only mark it incompatible if it's MLDev, not Agent Platform.
|
|
if getattr(tool, 'mcp_servers', None) and not is_agent_platform:
|
|
incompatible_tools_indexes.append(index)
|
|
return incompatible_tools_indexes
|
|
|
|
|
|
def get_function_map(
|
|
config: Optional[types.GenerateContentConfigOrDict] = None,
|
|
mcp_to_genai_tool_adapters: Optional[
|
|
dict[str, McpToGenAiToolAdapter]
|
|
] = None,
|
|
is_caller_method_async: bool = False,
|
|
) -> dict[str, Union[Callable[..., Any], McpToGenAiToolAdapter]]:
|
|
"""Returns a function map from the config."""
|
|
function_map: dict[str, Union[Callable[..., Any], McpToGenAiToolAdapter]] = {}
|
|
if not config:
|
|
return function_map
|
|
config_model = _create_generate_content_config_model(config)
|
|
if config_model.tools:
|
|
for tool in config_model.tools:
|
|
if callable(tool):
|
|
if inspect.iscoroutinefunction(tool) and not is_caller_method_async:
|
|
raise errors.UnsupportedFunctionError(
|
|
f'Function {tool.__name__} is a coroutine function, which is not'
|
|
' supported for automatic function calling. Please manually'
|
|
f' invoke {tool.__name__} to get the function response.'
|
|
)
|
|
function_map[tool.__name__] = tool
|
|
if mcp_to_genai_tool_adapters:
|
|
if not is_caller_method_async:
|
|
raise errors.UnsupportedFunctionError(
|
|
'MCP tools are not supported in synchronous methods.'
|
|
)
|
|
for tool_name, _ in mcp_to_genai_tool_adapters.items():
|
|
if function_map.get(tool_name):
|
|
raise ValueError(
|
|
f'Tool {tool_name} is already defined for the request.'
|
|
)
|
|
function_map.update(mcp_to_genai_tool_adapters)
|
|
return function_map
|
|
|
|
|
|
def convert_number_values_for_dict_function_call_args(
|
|
args: _common.StringDict,
|
|
) -> _common.StringDict:
|
|
"""Converts float values in dict with no decimal to integers."""
|
|
return {
|
|
key: convert_number_values_for_function_call_args(value)
|
|
for key, value in args.items()
|
|
}
|
|
|
|
|
|
def convert_number_values_for_function_call_args(
|
|
args: Union[dict[str, object], list[object], object],
|
|
) -> Union[dict[str, object], list[object], object]:
|
|
"""Converts float values with no decimal to integers."""
|
|
if isinstance(args, float) and args.is_integer():
|
|
return int(args)
|
|
if isinstance(args, dict):
|
|
return {
|
|
key: convert_number_values_for_function_call_args(value)
|
|
for key, value in args.items()
|
|
}
|
|
if isinstance(args, list):
|
|
return [
|
|
convert_number_values_for_function_call_args(value) for value in args
|
|
]
|
|
return args
|
|
|
|
|
|
def is_annotation_pydantic_model(annotation: Any) -> bool:
|
|
try:
|
|
return inspect.isclass(annotation) and issubclass(
|
|
annotation, pydantic.BaseModel
|
|
)
|
|
# for python 3.10 and below, inspect.isclass(annotation) has inconsistent
|
|
# results with versions above. for example, inspect.isclass(dict[str, int]) is
|
|
# True in 3.10 and below but False in 3.11 and above.
|
|
except TypeError:
|
|
return False
|
|
|
|
|
|
def convert_if_exist_pydantic_model(
|
|
value: Any, annotation: Any, param_name: str, func_name: str
|
|
) -> Any:
|
|
if isinstance(value, dict) and is_annotation_pydantic_model(annotation):
|
|
try:
|
|
return annotation(**value)
|
|
except pydantic.ValidationError as e:
|
|
raise errors.UnknownFunctionCallArgumentError(
|
|
f'Failed to parse parameter {param_name} for function'
|
|
f' {func_name} from function call part because function call argument'
|
|
f' value {value} is not compatible with parameter annotation'
|
|
f' {annotation}, due to error {e}'
|
|
)
|
|
if isinstance(value, list) and get_origin(annotation) == list:
|
|
item_type = get_args(annotation)[0]
|
|
return [
|
|
convert_if_exist_pydantic_model(item, item_type, param_name, func_name)
|
|
for item in value
|
|
]
|
|
if isinstance(value, dict) and get_origin(annotation) == dict:
|
|
_, value_type = get_args(annotation)
|
|
return {
|
|
k: convert_if_exist_pydantic_model(v, value_type, param_name, func_name)
|
|
for k, v in value.items()
|
|
}
|
|
# example 1: typing.Union[int, float]
|
|
# example 2: int | float equivalent to UnionType[int, float]
|
|
if get_origin(annotation) in (Union, UnionType):
|
|
for arg in get_args(annotation):
|
|
if (
|
|
(get_args(arg) and get_origin(arg) is list)
|
|
or isinstance(value, arg)
|
|
or (isinstance(value, dict) and is_annotation_pydantic_model(arg))
|
|
):
|
|
try:
|
|
return convert_if_exist_pydantic_model(
|
|
value, arg, param_name, func_name
|
|
)
|
|
# do not raise here because there could be multiple pydantic model types
|
|
# in the union type.
|
|
except pydantic.ValidationError:
|
|
continue
|
|
# if none of the union type is matched, raise error
|
|
raise errors.UnknownFunctionCallArgumentError(
|
|
f'Failed to parse parameter {param_name} for function'
|
|
f' {func_name} from function call part because function call argument'
|
|
f' value {value} cannot be converted to parameter annotation'
|
|
f' {annotation}.'
|
|
)
|
|
# the only exception for value and annotation type to be different is int and
|
|
# float. see convert_number_values_for_function_call_args function for context
|
|
if isinstance(value, int) and annotation is float:
|
|
return value
|
|
if not isinstance(value, annotation):
|
|
raise errors.UnknownFunctionCallArgumentError(
|
|
f'Failed to parse parameter {param_name} for function {func_name} from'
|
|
f' function call part because function call argument value {value} is'
|
|
f' not compatible with parameter annotation {annotation}.'
|
|
)
|
|
return value
|
|
|
|
|
|
def convert_argument_from_function(
|
|
args: _common.StringDict, function: Callable[..., Any]
|
|
) -> _common.StringDict:
|
|
signature = inspect.signature(function)
|
|
func_name = function.__name__
|
|
converted_args = {}
|
|
for param_name, param in signature.parameters.items():
|
|
if param_name in args:
|
|
converted_args[param_name] = convert_if_exist_pydantic_model(
|
|
args[param_name],
|
|
param.annotation,
|
|
param_name,
|
|
func_name,
|
|
)
|
|
return converted_args
|
|
|
|
|
|
def invoke_function_from_dict_args(
|
|
args: _common.StringDict, function_to_invoke: Callable[..., Any]
|
|
) -> Any:
|
|
converted_args = convert_argument_from_function(args, function_to_invoke)
|
|
try:
|
|
return function_to_invoke(**converted_args)
|
|
except Exception as e:
|
|
raise errors.FunctionInvocationError(
|
|
f'Failed to invoke function {function_to_invoke.__name__} with'
|
|
f' converted arguments {converted_args} from model returned function'
|
|
f' call argument {args} because of error {e}'
|
|
)
|
|
|
|
|
|
async def invoke_function_from_dict_args_async(
|
|
args: _common.StringDict, function_to_invoke: Callable[..., Any]
|
|
) -> Any:
|
|
converted_args = convert_argument_from_function(args, function_to_invoke)
|
|
try:
|
|
return await function_to_invoke(**converted_args)
|
|
except Exception as e:
|
|
raise errors.FunctionInvocationError(
|
|
f'Failed to invoke function {function_to_invoke.__name__} with'
|
|
f' converted arguments {converted_args} from model returned function'
|
|
f' call argument {args} because of error {e}'
|
|
)
|
|
|
|
|
|
def get_function_response_parts(
|
|
response: types.GenerateContentResponse,
|
|
function_map: dict[str, Union[Callable[..., Any], McpToGenAiToolAdapter]],
|
|
) -> list[types.Part]:
|
|
"""Returns the function response parts from the response."""
|
|
func_response_parts = []
|
|
if (
|
|
response.candidates is not None
|
|
and isinstance(response.candidates[0].content, types.Content)
|
|
and response.candidates[0].content.parts is not None
|
|
):
|
|
for part in response.candidates[0].content.parts:
|
|
if not part.function_call:
|
|
continue
|
|
func_name = part.function_call.name
|
|
if func_name is not None and part.function_call.args is not None:
|
|
func = function_map[func_name]
|
|
args = convert_number_values_for_dict_function_call_args(
|
|
part.function_call.args
|
|
)
|
|
func_response: _common.StringDict
|
|
try:
|
|
if not isinstance(func, McpToGenAiToolAdapter):
|
|
func_response = {
|
|
'result': invoke_function_from_dict_args(args, func)
|
|
}
|
|
except Exception as e: # pylint: disable=broad-except
|
|
func_response = {'error': str(e)}
|
|
func_response_part = types.Part.from_function_response(
|
|
name=func_name, response=func_response
|
|
)
|
|
func_response_parts.append(func_response_part)
|
|
return func_response_parts
|
|
|
|
|
|
async def get_function_response_parts_async(
|
|
response: types.GenerateContentResponse,
|
|
function_map: dict[str, Union[Callable[..., Any], McpToGenAiToolAdapter]],
|
|
) -> list[types.Part]:
|
|
"""Returns the function response parts from the response."""
|
|
func_response_parts = []
|
|
if (
|
|
response.candidates is not None
|
|
and isinstance(response.candidates[0].content, types.Content)
|
|
and response.candidates[0].content.parts is not None
|
|
):
|
|
for part in response.candidates[0].content.parts:
|
|
if not part.function_call:
|
|
continue
|
|
func_name = part.function_call.name
|
|
if func_name is not None:
|
|
func = function_map[func_name]
|
|
# Treat None as an empty dictionary for execution
|
|
raw_args = (
|
|
part.function_call.args
|
|
if part.function_call.args is not None
|
|
else {}
|
|
)
|
|
args = convert_number_values_for_dict_function_call_args(raw_args)
|
|
try:
|
|
if isinstance(func, McpToGenAiToolAdapter):
|
|
mcp_tool_response = await func.call_tool(
|
|
types.FunctionCall(name=func_name, args=args)
|
|
)
|
|
if mcp_tool_response.isError:
|
|
func_response = {'error': mcp_tool_response}
|
|
else:
|
|
func_response = {'result': mcp_tool_response}
|
|
elif inspect.iscoroutinefunction(func):
|
|
func_response = {
|
|
'result': await invoke_function_from_dict_args_async(args, func)
|
|
}
|
|
else:
|
|
func_response = {
|
|
'result': await asyncio.to_thread(
|
|
invoke_function_from_dict_args, args, func
|
|
)
|
|
}
|
|
except Exception as e: # pylint: disable=broad-except
|
|
func_response = {'error': str(e)} # type: ignore[dict-item]
|
|
func_response_part = types.Part.from_function_response(
|
|
name=func_name, response=func_response
|
|
)
|
|
func_response_parts.append(func_response_part)
|
|
return func_response_parts
|
|
|
|
|
|
def should_disable_afc(
|
|
config: Optional[types.GenerateContentConfigOrDict] = None,
|
|
) -> bool:
|
|
"""Returns whether automatic function calling is enabled."""
|
|
if not config:
|
|
return False
|
|
config_model = _create_generate_content_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
|
|
is not None
|
|
and int(config_model.automatic_function_calling.maximum_remote_calls) <= 0
|
|
):
|
|
logger.warning(
|
|
'max_remote_calls in automatic_function_calling_config'
|
|
f' {config_model.automatic_function_calling.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
|
|
|
|
# Default to enable AFC if not specified.
|
|
if (
|
|
not config_model.automatic_function_calling
|
|
or config_model.automatic_function_calling.disable is None
|
|
):
|
|
return False
|
|
|
|
if (
|
|
config_model.automatic_function_calling.disable
|
|
and config_model.automatic_function_calling.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
|
|
):
|
|
logger.warning(
|
|
'`automatic_function_calling.disable` is set to `True`. And'
|
|
' `automatic_function_calling.maximum_remote_calls` is a'
|
|
' positive number'
|
|
f' {config_model.automatic_function_calling.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.maximum_remote_calls` unset.'
|
|
)
|
|
|
|
return config_model.automatic_function_calling.disable
|
|
|
|
|
|
def get_max_remote_calls_afc(
|
|
config: Optional[types.GenerateContentConfigOrDict] = 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):
|
|
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
|
|
):
|
|
return _DEFAULT_MAX_REMOTE_CALLS_AFC
|
|
return int(config_model.automatic_function_calling.maximum_remote_calls)
|
|
|
|
|
|
def raise_error_for_afc_incompatible_config(config: Optional[types.GenerateContentConfig]
|
|
) -> None:
|
|
"""Raises an error if the config is not compatible with AFC."""
|
|
if (
|
|
not config
|
|
or not config.tool_config
|
|
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
|
|
stream_function_call = (
|
|
config.tool_config.function_calling_config.stream_function_call_arguments
|
|
)
|
|
|
|
if stream_function_call and not disable_afc_config:
|
|
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.'
|
|
)
|
|
|
|
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:
|
|
return True
|
|
return not config_model.automatic_function_calling.ignore_call_history
|
|
|
|
|
|
def parse_config_for_mcp_usage(
|
|
config: Optional[types.GenerateContentConfigOrDict] = None,
|
|
) -> Optional[types.GenerateContentConfig]:
|
|
"""Returns a parsed config with an appended MCP header if MCP tools or sessions are used."""
|
|
if not config:
|
|
return None
|
|
config_model = _create_generate_content_config_model(config)
|
|
# Create a copy of the config model with the tools field cleared since some
|
|
# tools may not be pickleable.
|
|
config_model_copy = config_model.model_copy(update={'tools': None})
|
|
config_model_copy.tools = config_model.tools
|
|
if config_model.tools and _mcp_utils.has_mcp_tool_usage(config_model.tools):
|
|
if config_model_copy.http_options is None:
|
|
config_model_copy.http_options = types.HttpOptions(headers={})
|
|
if config_model_copy.http_options.headers is None:
|
|
config_model_copy.http_options.headers = {}
|
|
_mcp_utils.set_mcp_usage_header(config_model_copy.http_options.headers)
|
|
|
|
return config_model_copy
|
|
|
|
|
|
async def parse_config_for_mcp_sessions(
|
|
config: Optional[types.GenerateContentConfigOrDict] = None,
|
|
is_agent_platform: bool = False,
|
|
) -> tuple[
|
|
Optional[types.GenerateContentConfig],
|
|
dict[str, McpToGenAiToolAdapter],
|
|
]:
|
|
"""Returns a parsed config with MCP sessions converted to GenAI tools.
|
|
|
|
Also returns a map of MCP tools to GenAI tool adapters to be used for AFC.
|
|
"""
|
|
mcp_to_genai_tool_adapters: dict[str, McpToGenAiToolAdapter] = {}
|
|
parsed_config = parse_config_for_mcp_usage(config)
|
|
if not parsed_config:
|
|
return None, mcp_to_genai_tool_adapters
|
|
# Create a copy of the config model with the tools field cleared as they will
|
|
# be replaced with the MCP tools converted to GenAI tools.
|
|
parsed_config_copy = parsed_config.model_copy(update={'tools': None})
|
|
if parsed_config.tools:
|
|
parsed_config_copy.tools = []
|
|
for tool in parsed_config.tools:
|
|
if McpClientSession is not None and isinstance(tool, McpClientSession):
|
|
mcp_to_genai_tool_adapter = McpToGenAiToolAdapter(
|
|
tool, await tool.list_tools(), is_agent_platform=is_agent_platform
|
|
)
|
|
# Extend the config with the MCP session tools converted to GenAI tools.
|
|
parsed_config_copy.tools.extend(mcp_to_genai_tool_adapter.tools)
|
|
for genai_tool in mcp_to_genai_tool_adapter.tools:
|
|
if genai_tool.function_declarations:
|
|
for function_declaration in genai_tool.function_declarations:
|
|
if function_declaration.name:
|
|
if mcp_to_genai_tool_adapters.get(function_declaration.name):
|
|
raise ValueError(
|
|
f'Tool {function_declaration.name} is already defined for'
|
|
' the request.'
|
|
)
|
|
mcp_to_genai_tool_adapters[function_declaration.name] = (
|
|
mcp_to_genai_tool_adapter
|
|
)
|
|
else:
|
|
parsed_config_copy.tools.append(tool)
|
|
|
|
return parsed_config_copy, mcp_to_genai_tool_adapters
|
|
|
|
|
|
def append_chunk_contents(
|
|
contents: Union[types.ContentListUnion, types.ContentListUnionDict],
|
|
chunk: types.GenerateContentResponse,
|
|
) -> Union[types.ContentListUnion, types.ContentListUnionDict]:
|
|
"""Appends the contents of the chunk to the contents list and returns it."""
|
|
if chunk is not None and chunk.candidates is not None:
|
|
chunk_content = chunk.candidates[0].content
|
|
contents = t.t_contents(contents) # type: ignore[assignment]
|
|
if isinstance(contents, list) and chunk_content is not None:
|
|
contents.append(chunk_content) # type: ignore[arg-type]
|
|
return contents
|
|
|
|
|
|
def prepare_resumable_upload(
|
|
file: Union[str, os.PathLike[str], io.IOBase],
|
|
user_http_options: Optional[types.HttpOptionsOrDict] = None,
|
|
user_mime_type: Optional[str] = None,
|
|
) -> tuple[
|
|
types.HttpOptions,
|
|
int,
|
|
str,
|
|
]:
|
|
"""Prepares the HTTP options, file bytes size and mime type for a resumable upload.
|
|
|
|
This function inspects a file (from a path or an in-memory object) to
|
|
determine its size and MIME type. It then constructs the necessary HTTP
|
|
headers and options required to initiate a resumable upload session.
|
|
"""
|
|
size_bytes = None
|
|
mime_type = user_mime_type
|
|
if isinstance(file, io.IOBase):
|
|
if mime_type is None:
|
|
raise ValueError(
|
|
'Unknown mime type: Could not determine the mimetype for your'
|
|
' file\n please set the `mime_type` argument'
|
|
)
|
|
if hasattr(file, 'mode'):
|
|
if 'b' not in file.mode:
|
|
raise ValueError('The file must be opened in binary mode.')
|
|
offset = file.tell()
|
|
file.seek(0, os.SEEK_END)
|
|
size_bytes = file.tell() - offset
|
|
file.seek(offset, os.SEEK_SET)
|
|
else:
|
|
fs_path = os.fspath(file)
|
|
if not fs_path or not os.path.isfile(fs_path):
|
|
raise FileNotFoundError(f'{file} is not a valid file path.')
|
|
size_bytes = os.path.getsize(fs_path)
|
|
if mime_type is None:
|
|
mime_type, _ = mimetypes.guess_type(fs_path)
|
|
if mime_type is None:
|
|
raise ValueError(
|
|
'Unknown mime type: Could not determine the mimetype for your'
|
|
' file\n please set the `mime_type` argument'
|
|
)
|
|
http_options: types.HttpOptions
|
|
if user_http_options:
|
|
if isinstance(user_http_options, dict):
|
|
user_http_options = types.HttpOptions(**user_http_options)
|
|
http_options = user_http_options
|
|
http_options.api_version = ''
|
|
http_options.headers = {
|
|
'Content-Type': 'application/json',
|
|
'X-Goog-Upload-Protocol': 'resumable',
|
|
'X-Goog-Upload-Command': 'start',
|
|
'X-Goog-Upload-Header-Content-Length': f'{size_bytes}',
|
|
'X-Goog-Upload-Header-Content-Type': f'{mime_type}',
|
|
}
|
|
else:
|
|
http_options = types.HttpOptions(
|
|
api_version='',
|
|
headers={
|
|
'Content-Type': 'application/json',
|
|
'X-Goog-Upload-Protocol': 'resumable',
|
|
'X-Goog-Upload-Command': 'start',
|
|
'X-Goog-Upload-Header-Content-Length': f'{size_bytes}',
|
|
'X-Goog-Upload-Header-Content-Type': f'{mime_type}',
|
|
},
|
|
)
|
|
if isinstance(file, (str, os.PathLike)):
|
|
if http_options.headers is None:
|
|
http_options.headers = {}
|
|
http_options.headers['X-Goog-Upload-File-Name'] = os.path.basename(file)
|
|
return http_options, size_bytes, mime_type
|
|
|
|
|
|
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
|
|
return False
|
|
|
|
|
|
def get_usage_header(
|
|
config: Optional[types.GenerateContentConfigOrDict] = None,
|
|
usage: str = 'afc',
|
|
) -> types.GenerateContentConfig:
|
|
"""Sets the afc version label."""
|
|
usage_header = f'google-genai-sdk/{usage}'
|
|
if not config:
|
|
config_model = types.GenerateContentConfig()
|
|
elif isinstance(config, dict):
|
|
config_model = types.GenerateContentConfig(**config)
|
|
else:
|
|
config_model = config
|
|
|
|
if not config_model.http_options:
|
|
config_model.http_options = types.HttpOptions()
|
|
existing_headers = config_model.http_options.headers or {}
|
|
if 'user-agent' in existing_headers:
|
|
if usage_header not in existing_headers['user-agent']:
|
|
existing_headers['user-agent'] += f' {usage_header}'
|
|
else:
|
|
existing_headers['user-agent'] = usage_header
|
|
if 'x-goog-api-client' in existing_headers:
|
|
if usage_header not in existing_headers['x-goog-api-client']:
|
|
existing_headers['x-goog-api-client'] += f' {usage_header}'
|
|
else:
|
|
existing_headers['x-goog-api-client'] = usage_header
|
|
config_model.http_options.headers = existing_headers
|
|
return config_model |