feat(phase1): add image storage utilities
- Create backend/services/image_storage.py with 4 core functions:
- sanitize_filename(): remove unsafe chars, limit to 255 chars, convert to lowercase
- get_unique_filename(): handle collisions with UUID suffix (format: {name}_{uuid8}_{variant}.jpg)
- ensure_image_directories(): create /images/ root and category subdirs on startup
- save_image(): save bytes to /images/{category}/{filename}, returns relative path
- Create comprehensive test suite (22 tests) covering all functionality
- Integrate ensure_image_directories() into FastAPI startup event
- Directory structure: /images/{category}/{filename}
- Collision handling: auto-suffix with UUID if filename exists
- All tests passing, pathlib.Path for safe operations
This commit is contained in:
@@ -0,0 +1,21 @@
|
||||
# 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."""
|
||||
|
||||
|
||||
import pytest
|
||||
pytest.register_assert_rewrite('genai.replay_api_client')
|
||||
@@ -0,0 +1,21 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""AFC helper function unit tests.
|
||||
|
||||
For AFC end to end test, please write test cases in
|
||||
test_generate_content_tools.py module
|
||||
"""
|
||||
@@ -0,0 +1,309 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""Test convert_if_exist_pydantic_model."""
|
||||
|
||||
import inspect
|
||||
from typing import Optional, Union
|
||||
import pydantic
|
||||
import pytest
|
||||
import sys
|
||||
from ... import errors
|
||||
from ..._extra_utils import convert_if_exist_pydantic_model
|
||||
|
||||
|
||||
def test_builtin_types():
|
||||
assert convert_if_exist_pydantic_model(1, int, 'param_name', 'func_name') == 1
|
||||
assert (
|
||||
convert_if_exist_pydantic_model(1.0, float, 'param_name', 'func_name')
|
||||
== 1.0
|
||||
)
|
||||
assert (
|
||||
convert_if_exist_pydantic_model('1.0', str, 'param_name', 'func_name')
|
||||
== '1.0'
|
||||
)
|
||||
assert (
|
||||
convert_if_exist_pydantic_model(True, bool, 'param_name', 'func_name')
|
||||
is True
|
||||
)
|
||||
assert convert_if_exist_pydantic_model(
|
||||
[1], list, 'param_name', 'func_name'
|
||||
) == [1]
|
||||
assert convert_if_exist_pydantic_model(
|
||||
{'key1': 1}, dict, 'param_name', 'func_name'
|
||||
) == {'key1': 1}
|
||||
assert convert_if_exist_pydantic_model(
|
||||
{'key1': 1, 'key2': 2}, dict[str, int], 'param_name', 'func_name'
|
||||
) == {'key1': 1, 'key2': 2}
|
||||
|
||||
|
||||
def test_value_int_annotation_float():
|
||||
assert (
|
||||
convert_if_exist_pydantic_model(1, float, 'param_name', 'func_name')
|
||||
== 1.0
|
||||
)
|
||||
|
||||
|
||||
def test_union_types():
|
||||
assert (
|
||||
convert_if_exist_pydantic_model(
|
||||
1, Union[int, float], 'param_name', 'func_name'
|
||||
)
|
||||
== 1
|
||||
)
|
||||
assert (
|
||||
convert_if_exist_pydantic_model(
|
||||
1.0, Union[int, float], 'param_name', 'func_name'
|
||||
)
|
||||
== 1.0
|
||||
)
|
||||
assert (
|
||||
convert_if_exist_pydantic_model(
|
||||
'1.0', Union[str, float], 'param_name', 'func_name'
|
||||
)
|
||||
== '1.0'
|
||||
)
|
||||
assert (
|
||||
convert_if_exist_pydantic_model(
|
||||
True, Union[bool, str], 'param_name', 'func_name'
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
# | in 3.10+
|
||||
if sys.version_info >= (3, 10):
|
||||
assert (
|
||||
convert_if_exist_pydantic_model(1, int | float, 'param_name', 'func_name')
|
||||
== 1
|
||||
)
|
||||
|
||||
|
||||
def test_nested_pydantic_model():
|
||||
class SimpleModel(pydantic.BaseModel):
|
||||
key1_simple: int
|
||||
key2_simple: float
|
||||
|
||||
class ComplexModel(pydantic.BaseModel):
|
||||
key1_complex: SimpleModel
|
||||
key2_complex: list[SimpleModel]
|
||||
key3_complex: dict[str, SimpleModel]
|
||||
|
||||
def foo(_: ComplexModel):
|
||||
pass
|
||||
|
||||
annotation = inspect.signature(foo).parameters['_'].annotation
|
||||
original_args = {
|
||||
'key1_complex': {'key1_simple': 1, 'key2_simple': 1.0},
|
||||
'key2_complex': [
|
||||
{'key1_simple': 2, 'key2_simple': 2.0},
|
||||
{'key1_simple': 3, 'key2_simple': 3.0},
|
||||
],
|
||||
'key3_complex': {
|
||||
'key1_simple': {'key1_simple': 4, 'key2_simple': 4.0},
|
||||
'key2_simple': {'key1_simple': 5, 'key2_simple': 5.0},
|
||||
},
|
||||
}
|
||||
|
||||
converted_args = convert_if_exist_pydantic_model(
|
||||
original_args, annotation, 'param_name', 'func_name'
|
||||
)
|
||||
assert isinstance(converted_args, ComplexModel)
|
||||
assert isinstance(converted_args.key1_complex, SimpleModel)
|
||||
assert isinstance(converted_args.key2_complex[0], SimpleModel)
|
||||
assert isinstance(converted_args.key2_complex[1], SimpleModel)
|
||||
assert isinstance(converted_args.key3_complex['key1_simple'], SimpleModel)
|
||||
assert isinstance(converted_args.key3_complex['key2_simple'], SimpleModel)
|
||||
assert converted_args.model_dump() == original_args
|
||||
|
||||
|
||||
def test_pydantic_model_in_list_union_type():
|
||||
class SimpleModel(pydantic.BaseModel):
|
||||
key1_simple: int
|
||||
key2_simple: float
|
||||
|
||||
def foo(_: list[Union[int, SimpleModel]]):
|
||||
pass
|
||||
|
||||
annotation = inspect.signature(foo).parameters['_'].annotation
|
||||
original_args = [1, {'key1_simple': 1, 'key2_simple': 1.0}]
|
||||
|
||||
converted_args = convert_if_exist_pydantic_model(
|
||||
original_args, annotation, 'param_name', 'func_name'
|
||||
)
|
||||
|
||||
assert isinstance(converted_args, list)
|
||||
assert len(converted_args) == 2
|
||||
assert isinstance(converted_args[0], int)
|
||||
assert isinstance(converted_args[1], SimpleModel)
|
||||
assert converted_args[0] == original_args[0]
|
||||
assert converted_args[1].model_dump() == original_args[1]
|
||||
|
||||
|
||||
def test_generic_list_type():
|
||||
def foo(_: Optional[list[Union[int, float]]]):
|
||||
pass
|
||||
|
||||
annotation = inspect.signature(foo).parameters['_'].annotation
|
||||
original_args = [1, 1.5]
|
||||
|
||||
converted_args = convert_if_exist_pydantic_model(
|
||||
original_args, annotation, 'param_name', 'func_name'
|
||||
)
|
||||
|
||||
assert isinstance(converted_args, list)
|
||||
assert len(converted_args) == 2
|
||||
assert isinstance(converted_args[0], int)
|
||||
assert isinstance(converted_args[1], float)
|
||||
assert converted_args[0] == original_args[0]
|
||||
|
||||
|
||||
@pytest.mark.skipif(sys.version_info < (3, 10), reason='need python 3.10+')
|
||||
def test_pydantic_model_with_union_and_generic_type():
|
||||
class Model1(pydantic.BaseModel):
|
||||
arg1: Union[int, float]
|
||||
arg2: list[str]
|
||||
arg3: int | float # python 3.9+
|
||||
arg4: Union[list[str], str]
|
||||
|
||||
def foo(_: Optional[list[Union[int, Model1]]]):
|
||||
pass
|
||||
|
||||
annotation = inspect.signature(foo).parameters['_'].annotation
|
||||
original_args = [1, {'arg1': 1, 'arg2': ['a'], 'arg3': 1.0, 'arg4': '1.0'}]
|
||||
|
||||
converted_args = convert_if_exist_pydantic_model(
|
||||
original_args, annotation, 'param_name', 'func_name'
|
||||
)
|
||||
|
||||
assert isinstance(converted_args, list)
|
||||
assert len(converted_args) == 2
|
||||
assert isinstance(converted_args[0], int)
|
||||
assert isinstance(converted_args[1], Model1)
|
||||
assert converted_args[0] == original_args[0]
|
||||
|
||||
|
||||
def test_unknown_pydantic_model_argument_error():
|
||||
class SimpleModel(pydantic.BaseModel):
|
||||
key1_simple: int
|
||||
key2_simple: float
|
||||
|
||||
def foo(_: SimpleModel):
|
||||
pass
|
||||
|
||||
annotation = inspect.signature(foo).parameters['_'].annotation
|
||||
original_args = {'key3_simple': 1, 'key2_simple': 1.0}
|
||||
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
convert_if_exist_pydantic_model(
|
||||
original_args, annotation, 'param_name', 'func_name'
|
||||
)
|
||||
|
||||
|
||||
def test_unknown_pydantic_model_argument_error_with_union_type():
|
||||
class SimpleModel1(pydantic.BaseModel):
|
||||
key1_simple: int
|
||||
key2_simple: float
|
||||
|
||||
class SimpleModel2(pydantic.BaseModel):
|
||||
key3_simple: str
|
||||
key4_simple: float
|
||||
|
||||
def foo(_: Union[SimpleModel1, SimpleModel2]):
|
||||
pass
|
||||
|
||||
annotation = inspect.signature(foo).parameters['_'].annotation
|
||||
original_args = {'key5_simple': 1, 'key4_simple': 1.0}
|
||||
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
convert_if_exist_pydantic_model(
|
||||
original_args, annotation, 'param_name', 'func_name'
|
||||
)
|
||||
|
||||
|
||||
def test_unknown_pydantic_model_argument_error_with_union_type_and_builtin_type():
|
||||
class SimpleModel1(pydantic.BaseModel):
|
||||
key1_simple: int
|
||||
key2_simple: float
|
||||
|
||||
def foo(_: Union[SimpleModel1, int]):
|
||||
pass
|
||||
|
||||
annotation = inspect.signature(foo).parameters['_'].annotation
|
||||
original_args = {'key5_simple': 1, 'key4_simple': 1.0}
|
||||
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
convert_if_exist_pydantic_model(
|
||||
original_args, annotation, 'param_name', 'func_name'
|
||||
)
|
||||
|
||||
|
||||
def test_incompatible_value_and_annotation():
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
convert_if_exist_pydantic_model(
|
||||
{'key1_simple': 1, 'key2_simple': 1.0},
|
||||
int,
|
||||
'param_name',
|
||||
'func_name',
|
||||
)
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
convert_if_exist_pydantic_model(
|
||||
{'key1_simple': 1, 'key2_simple': 1.0},
|
||||
float,
|
||||
'param_name',
|
||||
'func_name',
|
||||
)
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
convert_if_exist_pydantic_model(
|
||||
1,
|
||||
str,
|
||||
'param_name',
|
||||
'func_name',
|
||||
)
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
convert_if_exist_pydantic_model(
|
||||
1.0,
|
||||
str,
|
||||
'param_name',
|
||||
'func_name',
|
||||
)
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
convert_if_exist_pydantic_model(
|
||||
True,
|
||||
str,
|
||||
'param_name',
|
||||
'func_name',
|
||||
)
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
convert_if_exist_pydantic_model(
|
||||
[1],
|
||||
str,
|
||||
'param_name',
|
||||
'func_name',
|
||||
)
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
convert_if_exist_pydantic_model(
|
||||
{'key1': 1},
|
||||
str,
|
||||
'param_name',
|
||||
'func_name',
|
||||
)
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
convert_if_exist_pydantic_model(
|
||||
{'key1': 1, 'key2': 2},
|
||||
str,
|
||||
'param_name',
|
||||
'func_name',
|
||||
)
|
||||
@@ -0,0 +1,63 @@
|
||||
# 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 convert_number_values_for_function_call_args."""
|
||||
|
||||
from ..._extra_utils import convert_number_values_for_function_call_args
|
||||
|
||||
|
||||
def test_integer_value():
|
||||
assert convert_number_values_for_function_call_args(1) == 1
|
||||
|
||||
|
||||
def test_float_value():
|
||||
assert convert_number_values_for_function_call_args(1.0) == 1
|
||||
|
||||
|
||||
def test_string_value():
|
||||
assert convert_number_values_for_function_call_args('1.0') == '1.0'
|
||||
|
||||
|
||||
def test_boolean_value():
|
||||
assert convert_number_values_for_function_call_args(True) is True
|
||||
|
||||
|
||||
def test_none_value():
|
||||
assert convert_number_values_for_function_call_args(None) is None
|
||||
|
||||
|
||||
def test_float_value_with_decimal():
|
||||
assert convert_number_values_for_function_call_args(1.1) == 1.1
|
||||
|
||||
|
||||
def test_dict_value():
|
||||
assert convert_number_values_for_function_call_args(
|
||||
{'key1': 1.0, 'key2': 1.1}
|
||||
) == {'key1': 1, 'key2': 1.1}
|
||||
|
||||
|
||||
def test_list_value():
|
||||
assert convert_number_values_for_function_call_args([1.0, 1.1, 1.2]) == [
|
||||
1,
|
||||
1.1,
|
||||
1.2,
|
||||
]
|
||||
|
||||
|
||||
def test_nested_value():
|
||||
assert convert_number_values_for_function_call_args(
|
||||
{'key1': 1.0, 'key2': {'key3': 1.0, 'key4': [1.2, 2.0]}}
|
||||
) == {'key1': 1, 'key2': {'key3': 1, 'key4': [1.2, 2]}}
|
||||
@@ -0,0 +1,268 @@
|
||||
# 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 find_afc_incompatible_tool_indexes."""
|
||||
|
||||
from typing import Any
|
||||
import pytest
|
||||
from ... import types
|
||||
from ..._extra_utils import find_afc_incompatible_tool_indexes
|
||||
|
||||
try:
|
||||
from mcp import types as mcp_types
|
||||
from mcp import ClientSession as McpClientSession
|
||||
from ..._adapters import McpToGenAiToolAdapter
|
||||
except ImportError as e:
|
||||
import sys
|
||||
|
||||
if sys.version_info < (3, 10):
|
||||
raise ImportError(
|
||||
'MCP Tool requires Python 3.10 or above. Please upgrade your Python'
|
||||
' version.'
|
||||
) from e
|
||||
else:
|
||||
raise e
|
||||
|
||||
|
||||
def get_weather_tool(city: str) -> str:
|
||||
"""Get the weather in a city."""
|
||||
return f'The weather in {city} is sunny and 100 degrees.'
|
||||
|
||||
|
||||
class MockMcpClientSession(McpClientSession):
|
||||
|
||||
def __init__(self):
|
||||
self._read_stream = None
|
||||
self._write_stream = None
|
||||
|
||||
async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
|
||||
return mcp_types.CallToolResult(
|
||||
content=[mcp_types.TextContent(type='text', text='1.01')]
|
||||
)
|
||||
|
||||
mcp_to_genai_tool_adapter = McpToGenAiToolAdapter(
|
||||
session=MockMcpClientSession(),
|
||||
list_tools_result=mcp_types.ListToolsResult(tools=[]),
|
||||
)
|
||||
|
||||
|
||||
def test_no_config_returns_empty_list():
|
||||
"""Verifies that an empty list is returned if the input config is None.
|
||||
"""
|
||||
result = find_afc_incompatible_tool_indexes(config=None)
|
||||
assert result == []
|
||||
|
||||
def test_config_with_no_tools_returns_empty_list():
|
||||
"""Verifies an empty list is returned if the config has no 'tools' attribute.
|
||||
"""
|
||||
result = find_afc_incompatible_tool_indexes(
|
||||
config=types.GenerateContentConfig()
|
||||
)
|
||||
assert result == []
|
||||
|
||||
def test_empty_tools_list_returns_empty_list():
|
||||
"""Verifies that an empty list is returned if the 'tools' list is empty."""
|
||||
result = find_afc_incompatible_tool_indexes(
|
||||
config=types.GenerateContentConfig(tools=[])
|
||||
)
|
||||
|
||||
assert result == []
|
||||
|
||||
def test_all_compatible_tools_returns_empty_list_with_empty_fd():
|
||||
"""Verifies that an empty list is returned when all tools are compatible.
|
||||
|
||||
A tool is compatible if it's not a `types.Tool` or if its
|
||||
`function_declarations` attribute is empty or None from config.
|
||||
"""
|
||||
result = find_afc_incompatible_tool_indexes(
|
||||
config=types.GenerateContentConfig(
|
||||
tools=[
|
||||
types.Tool(
|
||||
google_search_retrieval=types.GoogleSearchRetrieval()
|
||||
),
|
||||
types.Tool(retrieval=types.Retrieval()),
|
||||
types.Tool(google_search=types.GoogleSearch()),
|
||||
types.Tool(code_execution=types.ToolCodeExecution()),
|
||||
types.Tool(google_maps=types.GoogleMaps()),
|
||||
types.Tool(url_context=types.UrlContext()),
|
||||
types.Tool(computer_use=types.ComputerUse()),
|
||||
types.Tool(code_execution=types.ToolCodeExecution()),
|
||||
types.Tool(function_declarations=[]),
|
||||
mcp_types.Tool(
|
||||
name='get_weather',
|
||||
description='Get the weather in a city.',
|
||||
inputSchema={
|
||||
'type': 'object',
|
||||
'properties': {'location': {'type': 'string'}},
|
||||
},
|
||||
),
|
||||
get_weather_tool,
|
||||
mcp_to_genai_tool_adapter,
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
assert result == []
|
||||
|
||||
def test_all_compatible_tools_returns_empty_list_with_none_fd():
|
||||
"""Verifies that an empty list is returned when all tools are compatible.
|
||||
|
||||
A tool is compatible if it's not a `types.Tool` or if its
|
||||
`function_declarations` attribute is empty or None from config.
|
||||
"""
|
||||
result = find_afc_incompatible_tool_indexes(
|
||||
config=types.GenerateContentConfig(
|
||||
tools=[
|
||||
types.Tool(
|
||||
google_search_retrieval=types.GoogleSearchRetrieval()
|
||||
),
|
||||
types.Tool(retrieval=types.Retrieval()),
|
||||
types.Tool(google_search=types.GoogleSearch()),
|
||||
types.Tool(code_execution=types.ToolCodeExecution()),
|
||||
types.Tool(google_maps=types.GoogleMaps()),
|
||||
types.Tool(url_context=types.UrlContext()),
|
||||
types.Tool(computer_use=types.ComputerUse()),
|
||||
types.Tool(code_execution=types.ToolCodeExecution()),
|
||||
types.Tool(function_declarations=None),
|
||||
mcp_types.Tool(
|
||||
name='get_weather',
|
||||
description='Get the weather in a city.',
|
||||
inputSchema={
|
||||
'type': 'object',
|
||||
'properties': {'location': {'type': 'string'}},
|
||||
},
|
||||
),
|
||||
get_weather_tool,
|
||||
mcp_to_genai_tool_adapter,
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
assert result == []
|
||||
|
||||
def test_all_compatible_tools_returns_empty_list():
|
||||
"""Verifies that an empty list is returned when all tools are compatible.
|
||||
|
||||
A tool is compatible if it's not a `types.Tool` or if its
|
||||
`function_declarations` attribute is empty or None from config.
|
||||
"""
|
||||
result = find_afc_incompatible_tool_indexes(
|
||||
config=types.GenerateContentConfig(
|
||||
tools=[
|
||||
types.Tool(
|
||||
google_search_retrieval=types.GoogleSearchRetrieval()
|
||||
),
|
||||
types.Tool(retrieval=types.Retrieval()),
|
||||
types.Tool(google_search=types.GoogleSearch()),
|
||||
types.Tool(code_execution=types.ToolCodeExecution()),
|
||||
types.Tool(google_maps=types.GoogleMaps()),
|
||||
types.Tool(url_context=types.UrlContext()),
|
||||
types.Tool(computer_use=types.ComputerUse()),
|
||||
types.Tool(code_execution=types.ToolCodeExecution()),
|
||||
types.Tool(function_declarations=[]),
|
||||
mcp_types.Tool(
|
||||
name='get_weather',
|
||||
description='Get the weather in a city.',
|
||||
inputSchema={
|
||||
'type': 'object',
|
||||
'properties': {'location': {'type': 'string'}},
|
||||
},
|
||||
),
|
||||
get_weather_tool,
|
||||
mcp_to_genai_tool_adapter,
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
assert result == []
|
||||
|
||||
def test_single_incompatible_tool():
|
||||
"""Verifies that the correct index is returned for a single incompatible
|
||||
tool.
|
||||
"""
|
||||
result = find_afc_incompatible_tool_indexes(
|
||||
config=types.GenerateContentConfig(
|
||||
tools=[
|
||||
types.Tool(
|
||||
google_search_retrieval=types.GoogleSearchRetrieval()
|
||||
),
|
||||
types.Tool(retrieval=types.Retrieval()),
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(name='test_function')
|
||||
]
|
||||
),
|
||||
get_weather_tool,
|
||||
mcp_to_genai_tool_adapter,
|
||||
]
|
||||
)
|
||||
)
|
||||
assert result == [2]
|
||||
|
||||
|
||||
def test_multiple_incompatible_tools():
|
||||
"""Verifies correct indexes are returned for multiple incompatible tools."""
|
||||
result = find_afc_incompatible_tool_indexes(
|
||||
config=types.GenerateContentConfig(
|
||||
tools=[
|
||||
types.Tool(
|
||||
google_search_retrieval=types.GoogleSearchRetrieval()
|
||||
),
|
||||
types.Tool(retrieval=types.Retrieval()),
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(name='test_function')
|
||||
]
|
||||
),
|
||||
types.Tool(computer_use=types.ComputerUse()),
|
||||
types.Tool(code_execution=types.ToolCodeExecution()),
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(name='test_function_2')
|
||||
]
|
||||
),
|
||||
get_weather_tool,
|
||||
mcp_to_genai_tool_adapter,
|
||||
]
|
||||
)
|
||||
)
|
||||
assert result == [2, 5]
|
||||
|
||||
def test_mcp_tool_incompatible():
|
||||
"""Verifies correct indexes are returned for multiple incompatible tools."""
|
||||
result = find_afc_incompatible_tool_indexes(
|
||||
config=types.GenerateContentConfig(
|
||||
tools=[
|
||||
types.Tool(
|
||||
google_search_retrieval=types.GoogleSearchRetrieval()
|
||||
),
|
||||
types.Tool(retrieval=types.Retrieval()),
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(name='test_function')
|
||||
]
|
||||
),
|
||||
types.Tool(code_execution=types.ToolCodeExecution()),
|
||||
|
||||
get_weather_tool,
|
||||
mcp_to_genai_tool_adapter,
|
||||
types.Tool(
|
||||
mcp_servers=[types.McpServer(name='test_mcp_server')]
|
||||
),
|
||||
]
|
||||
)
|
||||
)
|
||||
assert result == [2, 6]
|
||||
@@ -0,0 +1,530 @@
|
||||
# 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 _api_client
|
||||
from ... import _extra_utils
|
||||
from ... import client
|
||||
from ... import models
|
||||
from ... import types
|
||||
|
||||
|
||||
TEST_NO_AFC_PART = types.Part(
|
||||
text=(
|
||||
'Okay, here is the weather in San Francisco'
|
||||
' as of approximately 8:10 pm PST on May'
|
||||
' 10, 2023. Please note that the weather'
|
||||
' can change rapidly.'
|
||||
)
|
||||
)
|
||||
|
||||
TEST_FUNCTION_CALL_PART = types.Part(
|
||||
function_call=types.FunctionCall(
|
||||
name='get_current_weather',
|
||||
args={'location': 'San Francisco'},
|
||||
)
|
||||
)
|
||||
|
||||
TEST_FUNCTION_RESPONSE_PART = types.Part(
|
||||
function_response=types.FunctionResponse(
|
||||
name='get_current_weather',
|
||||
response={'result': 'sunny'},
|
||||
)
|
||||
)
|
||||
|
||||
TEST_AFC_TEXT_PART = types.Part(text='San Francisco weather is sunny.')
|
||||
|
||||
TEST_NO_AFC_CONTENT = types.Content(
|
||||
parts=[TEST_NO_AFC_PART],
|
||||
role='model',
|
||||
)
|
||||
|
||||
TEST_FUNCTION_CALL_CONTENT = types.Content(
|
||||
parts=[TEST_FUNCTION_CALL_PART],
|
||||
role='model',
|
||||
)
|
||||
|
||||
TEST_FUNCTION_RESPONSE_CONTENT = types.Content(
|
||||
parts=[TEST_FUNCTION_RESPONSE_PART],
|
||||
role='user',
|
||||
)
|
||||
|
||||
TEST_AFC_TEXT_CONTENT = types.Content(
|
||||
parts=[TEST_AFC_TEXT_PART],
|
||||
role='model',
|
||||
)
|
||||
|
||||
TEST_AFC_HISTORY = [
|
||||
types.Content(
|
||||
parts=[types.Part(text='what is the weather in San Francisco?')],
|
||||
role='user',
|
||||
),
|
||||
TEST_FUNCTION_CALL_CONTENT,
|
||||
TEST_FUNCTION_RESPONSE_CONTENT,
|
||||
]
|
||||
|
||||
|
||||
def get_current_weather(location: str) -> str:
|
||||
"""Returns the current weather.
|
||||
|
||||
Args:
|
||||
location: The location of a city and state, e.g. "San Francisco, CA".
|
||||
"""
|
||||
return 'windy'
|
||||
|
||||
|
||||
async def get_current_weather_async(location: str) -> str:
|
||||
"""Returns the current weather.
|
||||
|
||||
Args:
|
||||
location: The location of a city and state, e.g. "San Francisco, CA".
|
||||
"""
|
||||
return 'windy'
|
||||
|
||||
|
||||
def get_aqi_from_city(location: str) -> str:
|
||||
"""Returns the aqi index of a city.
|
||||
|
||||
Args:
|
||||
location: The city and State, e.g. San Francisco, CA.
|
||||
"""
|
||||
return None
|
||||
|
||||
|
||||
@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_get_function_response_parts_none():
|
||||
with mock.patch.object(
|
||||
_extra_utils,
|
||||
'get_function_response_parts',
|
||||
) as mock_get_function_response_parts_none:
|
||||
mock_get_function_response_parts_none.return_value = None
|
||||
yield mock_get_function_response_parts_none
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_get_function_response_parts_none_async():
|
||||
with mock.patch.object(
|
||||
_extra_utils,
|
||||
'get_function_response_parts_async',
|
||||
) as mock_get_function_response_parts_none_async:
|
||||
mock_get_function_response_parts_none_async.return_value = None
|
||||
yield mock_get_function_response_parts_none_async
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_get_function_response_parts() -> list[types.Part]:
|
||||
with mock.patch.object(
|
||||
_extra_utils, 'get_function_response_parts'
|
||||
) as mock_get_function_response_parts:
|
||||
mock_get_function_response_parts.side_effect = [
|
||||
[TEST_FUNCTION_RESPONSE_PART],
|
||||
[], # Breaks when the function response is not returned.
|
||||
]
|
||||
yield mock_get_function_response_parts
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_get_function_response_parts_async() -> list[types.Part]:
|
||||
with mock.patch.object(
|
||||
_extra_utils, 'get_function_response_parts_async'
|
||||
) as mock_get_function_response_parts_async:
|
||||
mock_get_function_response_parts_async.side_effect = [
|
||||
[TEST_FUNCTION_RESPONSE_PART],
|
||||
[], # Breaks when the function response is not returned.
|
||||
]
|
||||
yield mock_get_function_response_parts_async
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_generate_content_stream_no_afc():
|
||||
with mock.patch.object(
|
||||
models.Models, '_generate_content_stream'
|
||||
) as mock_stream_no_afc:
|
||||
mock_stream_no_afc.return_value = [
|
||||
types.GenerateContentResponse(
|
||||
candidates=[types.Candidate(content=TEST_NO_AFC_CONTENT)]
|
||||
)
|
||||
]
|
||||
yield mock_stream_no_afc
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_generate_content_stream_with_afc():
|
||||
with mock.patch.object(
|
||||
models.Models, '_generate_content_stream'
|
||||
) as mock_stream_with_afc:
|
||||
mock_stream_with_afc.side_effect = [
|
||||
[
|
||||
types.GenerateContentResponse(
|
||||
candidates=[types.Candidate(content=TEST_FUNCTION_CALL_CONTENT)]
|
||||
)
|
||||
],
|
||||
[
|
||||
types.GenerateContentResponse(
|
||||
candidates=[types.Candidate(content=TEST_AFC_TEXT_CONTENT)]
|
||||
)
|
||||
],
|
||||
]
|
||||
yield mock_stream_with_afc
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_generate_content_stream_no_afc_async():
|
||||
with mock.patch.object(
|
||||
models.AsyncModels, '_generate_content_stream'
|
||||
) as mock_stream_no_afc:
|
||||
|
||||
async def async_generator():
|
||||
yield types.GenerateContentResponse(
|
||||
candidates=[types.Candidate(content=TEST_NO_AFC_CONTENT)]
|
||||
)
|
||||
|
||||
mock_stream_no_afc.return_value = async_generator()
|
||||
yield mock_stream_no_afc
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_generate_content_stream_with_afc_async():
|
||||
with mock.patch.object(
|
||||
models.AsyncModels, '_generate_content_stream'
|
||||
) as mock_stream_with_afc:
|
||||
|
||||
async def async_generator_1():
|
||||
yield types.GenerateContentResponse(
|
||||
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
|
||||
|
||||
|
||||
def test_generate_content_stream_no_function_map(
|
||||
mock_generate_content_stream_no_afc,
|
||||
mock_get_function_response_parts_none,
|
||||
):
|
||||
"""Test when no function tools are provided.
|
||||
|
||||
Expected to answer past weather.
|
||||
"""
|
||||
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?'
|
||||
)
|
||||
for chunk in stream:
|
||||
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 == 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.
|
||||
|
||||
Expected to answer past weather.
|
||||
"""
|
||||
models_instance = models.Models(api_client_=mock_api_client)
|
||||
config = types.GenerateContentConfig(tools=[get_aqi_from_city])
|
||||
stream = models_instance.generate_content_stream(
|
||||
model='test_model',
|
||||
contents='what is the weather in San Francisco?',
|
||||
config=config,
|
||||
)
|
||||
for chunk in stream:
|
||||
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,
|
||||
):
|
||||
"""Test when function tools are provided and thought summaries are enabled.
|
||||
|
||||
Expected to answer weather based on function response.
|
||||
"""
|
||||
models_instance = models.Models(api_client_=mock_api_client)
|
||||
config = types.GenerateContentConfig(
|
||||
tools=[get_current_weather],
|
||||
thinking_config=types.ThinkingConfig(include_thoughts=True),
|
||||
)
|
||||
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)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_content_stream_no_function_map_async(
|
||||
mock_generate_content_stream_no_afc,
|
||||
mock_get_function_response_parts_none,
|
||||
):
|
||||
"""Test when no function tools are provided.
|
||||
|
||||
Expected to answer past weather.
|
||||
"""
|
||||
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?'
|
||||
)
|
||||
for chunk in stream:
|
||||
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 == 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.
|
||||
|
||||
Expected to answer past weather.
|
||||
"""
|
||||
models_instance = models.AsyncModels(api_client_=mock_api_client)
|
||||
config = types.GenerateContentConfig(tools=[get_aqi_from_city])
|
||||
stream = await models_instance.generate_content_stream(
|
||||
model='test_model',
|
||||
contents='what is the weather in San Francisco?',
|
||||
config=config,
|
||||
)
|
||||
async for chunk in stream:
|
||||
assert chunk.text == TEST_NO_AFC_PART.text
|
||||
|
||||
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,
|
||||
):
|
||||
"""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])
|
||||
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 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 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_thought_summaries_async(
|
||||
mock_generate_content_stream_with_afc_async,
|
||||
mock_get_function_response_parts_async,
|
||||
):
|
||||
"""Test when function tools are provided and thought summaries are enabled.
|
||||
|
||||
Expected to answer weather based on function response.
|
||||
"""
|
||||
models_instance = models.AsyncModels(api_client_=mock_api_client)
|
||||
config = types.GenerateContentConfig(
|
||||
tools=[get_current_weather],
|
||||
thinking_config=types.ThinkingConfig(include_thoughts=True),
|
||||
)
|
||||
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 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)
|
||||
@@ -0,0 +1,77 @@
|
||||
# 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 pytest
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
|
||||
def get_current_weather(location: str) -> str:
|
||||
"""Returns the current weather.
|
||||
|
||||
Args:
|
||||
location: The location of a city and state, e.g. "San Francisco, CA".
|
||||
"""
|
||||
return 'windy'
|
||||
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='models.generate_content',
|
||||
)
|
||||
pytest_plugins = ('pytest_asyncio',)
|
||||
|
||||
|
||||
def test_generate_content_stream_with_function_and_thought_summaries(client):
|
||||
"""Test when function tools are provided and thought summaries are enabled.
|
||||
|
||||
Expected to answer weather based on function response.
|
||||
"""
|
||||
config = types.GenerateContentConfig(
|
||||
tools=[get_current_weather],
|
||||
thinking_config=types.ThinkingConfig(include_thoughts=True),
|
||||
)
|
||||
stream = client.models.generate_content_stream(
|
||||
model='gemini-2.5-flash',
|
||||
contents='what is the weather in San Francisco, CA?',
|
||||
config=config,
|
||||
)
|
||||
|
||||
chunk = None
|
||||
for chunk in stream:
|
||||
assert chunk is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_content_stream_with_function_and_thought_summaries_async(
|
||||
client,
|
||||
):
|
||||
"""Test when function tools are provided and thought summaries are enabled.
|
||||
|
||||
Expected to answer weather based on function response.
|
||||
"""
|
||||
config = types.GenerateContentConfig(
|
||||
tools=[get_current_weather],
|
||||
thinking_config=types.ThinkingConfig(include_thoughts=True),
|
||||
)
|
||||
stream = await client.aio.models.generate_content_stream(
|
||||
model='gemini-2.5-flash',
|
||||
contents='what is the weather in San Francisco, CA?',
|
||||
config=config,
|
||||
)
|
||||
|
||||
chunk = None
|
||||
async for chunk in stream:
|
||||
assert chunk is not None
|
||||
@@ -0,0 +1,176 @@
|
||||
# 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 get_function_map."""
|
||||
|
||||
import typing
|
||||
from typing import Any
|
||||
import pytest
|
||||
from ..._extra_utils import get_function_map
|
||||
from ...errors import UnsupportedFunctionError
|
||||
from ...types import GenerateContentConfig
|
||||
|
||||
_is_mcp_imported = False
|
||||
if typing.TYPE_CHECKING:
|
||||
from mcp import types as mcp_types
|
||||
from mcp import ClientSession as McpClientSession
|
||||
from ..._adapters import McpToGenAiToolAdapter
|
||||
|
||||
_is_mcp_imported = True
|
||||
else:
|
||||
McpClientSession: typing.Type = Any
|
||||
McpToGenAiToolAdapter: typing.Type = Any
|
||||
try:
|
||||
from mcp import types as mcp_types
|
||||
from mcp import ClientSession as McpClientSession
|
||||
from ..._adapters import McpToGenAiToolAdapter
|
||||
|
||||
_is_mcp_imported = True
|
||||
except ImportError:
|
||||
McpClientSession = None
|
||||
McpToGenAiToolAdapter = None
|
||||
|
||||
|
||||
def test_coroutine_function():
|
||||
async def func_under_test():
|
||||
pass
|
||||
|
||||
config = GenerateContentConfig(tools=[func_under_test])
|
||||
|
||||
with pytest.raises(UnsupportedFunctionError):
|
||||
get_function_map(config)
|
||||
|
||||
|
||||
def test_empty_config():
|
||||
config = {}
|
||||
|
||||
assert get_function_map(config) == {}
|
||||
|
||||
|
||||
def test_empty_tools():
|
||||
config = GenerateContentConfig(top_p=0.5)
|
||||
|
||||
assert get_function_map(config) == {}
|
||||
|
||||
|
||||
def test_valid_function():
|
||||
def func_under_test():
|
||||
pass
|
||||
|
||||
config = GenerateContentConfig(tools=[func_under_test])
|
||||
|
||||
assert get_function_map(config) == {'func_under_test': func_under_test}
|
||||
|
||||
|
||||
def test_mcp_tool_raises_error():
|
||||
if not _is_mcp_imported:
|
||||
return
|
||||
|
||||
session = McpClientSession(read_stream=None, write_stream=None)
|
||||
config = GenerateContentConfig(tools=[session])
|
||||
mcp_to_genai_tool_adapters = {'tool': McpToGenAiToolAdapter(session, [])}
|
||||
with pytest.raises(UnsupportedFunctionError):
|
||||
get_function_map(
|
||||
config, mcp_to_genai_tool_adapters, is_caller_method_async=False
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_tool():
|
||||
if not _is_mcp_imported:
|
||||
return
|
||||
|
||||
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='tool',
|
||||
description='tool-description',
|
||||
inputSchema={
|
||||
'type': 'OBJECT',
|
||||
'properties': {
|
||||
'key1': {
|
||||
'type': 'STRING',
|
||||
},
|
||||
'key2': {
|
||||
'type': 'NUMBER',
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
session = MockMcpClientSession()
|
||||
config = GenerateContentConfig(tools=[session])
|
||||
mcp_to_genai_tool_adapters = {
|
||||
'tool': McpToGenAiToolAdapter(session, [await session.list_tools()]),
|
||||
}
|
||||
result = get_function_map(
|
||||
config, mcp_to_genai_tool_adapters, is_caller_method_async=True
|
||||
)
|
||||
assert isinstance(result['tool'], McpToGenAiToolAdapter)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplicate_mcp_tool_raises_error():
|
||||
if not _is_mcp_imported:
|
||||
return
|
||||
|
||||
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='tool',
|
||||
description='tool-description',
|
||||
inputSchema={
|
||||
'type': 'OBJECT',
|
||||
'properties': {
|
||||
'key1': {
|
||||
'type': 'STRING',
|
||||
},
|
||||
'key2': {
|
||||
'type': 'NUMBER',
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
def tool():
|
||||
pass
|
||||
|
||||
session = MockMcpClientSession()
|
||||
config = GenerateContentConfig(tools=[tool, session])
|
||||
mcp_to_genai_tool_adapters = {
|
||||
'tool': McpToGenAiToolAdapter(session, [await session.list_tools()]),
|
||||
}
|
||||
with pytest.raises(ValueError):
|
||||
get_function_map(
|
||||
config, mcp_to_genai_tool_adapters, is_caller_method_async=True
|
||||
)
|
||||
@@ -0,0 +1,277 @@
|
||||
# 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 get_function_response_parts."""
|
||||
|
||||
import typing
|
||||
from typing import Any
|
||||
import pytest
|
||||
from ..._extra_utils import get_function_response_parts, get_function_response_parts_async
|
||||
from ...errors import UnsupportedFunctionError
|
||||
from ...types import Candidate
|
||||
from ...types import Content
|
||||
from ...types import FunctionCall
|
||||
from ...types import FunctionResponse
|
||||
from ...types import GenerateContentResponse
|
||||
from ...types import Part
|
||||
|
||||
_is_mcp_imported = False
|
||||
if typing.TYPE_CHECKING:
|
||||
from mcp import types as mcp_types
|
||||
from mcp import ClientSession as McpClientSession
|
||||
from ..._adapters import McpToGenAiToolAdapter
|
||||
|
||||
_is_mcp_imported = True
|
||||
else:
|
||||
McpClientSession: typing.Type = Any
|
||||
McpToGenAiToolAdapter: typing.Type = Any
|
||||
try:
|
||||
from mcp import types as mcp_types
|
||||
from mcp import ClientSession as McpClientSession
|
||||
from ..._adapters import McpToGenAiToolAdapter
|
||||
|
||||
_is_mcp_imported = True
|
||||
except ImportError:
|
||||
McpClientSession = None
|
||||
McpToGenAiToolAdapter = None
|
||||
|
||||
|
||||
def test_integer_value():
|
||||
def func_under_test(a: int) -> int:
|
||||
return a + 1
|
||||
|
||||
response = GenerateContentResponse(
|
||||
candidates=[
|
||||
Candidate(
|
||||
content=Content(
|
||||
parts=[
|
||||
Part(
|
||||
function_call=FunctionCall(
|
||||
name='func_under_test',
|
||||
args={'a': 1},
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
function_map = {'func_under_test': func_under_test}
|
||||
expected_parts = [
|
||||
Part(
|
||||
function_response=FunctionResponse(
|
||||
name='func_under_test',
|
||||
response={'result': 2},
|
||||
)
|
||||
)
|
||||
]
|
||||
actual_parts = get_function_response_parts(response, function_map)
|
||||
|
||||
for actual_part, expected_part in zip(actual_parts, expected_parts):
|
||||
assert actual_part.model_dump_json(
|
||||
exclude_none=True
|
||||
) == expected_part.model_dump_json(exclude_none=True)
|
||||
|
||||
|
||||
def test_float_value():
|
||||
def func_under_test(a: float) -> float:
|
||||
return a + 1.0
|
||||
|
||||
response = GenerateContentResponse(
|
||||
candidates=[
|
||||
Candidate(
|
||||
content=Content(
|
||||
parts=[
|
||||
Part(
|
||||
function_call=FunctionCall(
|
||||
name='func_under_test',
|
||||
args={'a': 1.0},
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
function_map = {'func_under_test': func_under_test}
|
||||
expected_parts = [
|
||||
Part(
|
||||
function_response=FunctionResponse(
|
||||
name='func_under_test',
|
||||
response={'result': 2.0},
|
||||
)
|
||||
)
|
||||
]
|
||||
actual_parts = get_function_response_parts(response, function_map)
|
||||
|
||||
for actual_part, expected_part in zip(actual_parts, expected_parts):
|
||||
assert actual_part.model_dump_json(
|
||||
exclude_none=True
|
||||
) == expected_part.model_dump_json(exclude_none=True)
|
||||
|
||||
|
||||
def test_string_value():
|
||||
def func_under_test(a: str) -> str:
|
||||
return a + '1'
|
||||
|
||||
response = GenerateContentResponse(
|
||||
candidates=[
|
||||
Candidate(
|
||||
content=Content(
|
||||
parts=[
|
||||
Part(
|
||||
function_call=FunctionCall(
|
||||
name='func_under_test',
|
||||
args={'a': '1.0'},
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
function_map = {'func_under_test': func_under_test}
|
||||
expected_parts = [
|
||||
Part(
|
||||
function_response=FunctionResponse(
|
||||
name='func_under_test',
|
||||
response={'result': '1.01'},
|
||||
)
|
||||
)
|
||||
]
|
||||
actual_parts = get_function_response_parts(response, function_map)
|
||||
|
||||
for actual_part, expected_part in zip(actual_parts, expected_parts):
|
||||
assert actual_part.model_dump_json(
|
||||
exclude_none=True
|
||||
) == expected_part.model_dump_json(exclude_none=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_tool():
|
||||
if not _is_mcp_imported:
|
||||
return
|
||||
|
||||
class MockMcpClientSession(McpClientSession):
|
||||
|
||||
def __init__(self):
|
||||
self._read_stream = None
|
||||
self._write_stream = None
|
||||
|
||||
async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
|
||||
return mcp_types.CallToolResult(
|
||||
content=[mcp_types.TextContent(type='text', text='1.01')]
|
||||
)
|
||||
|
||||
mcp_to_genai_tool_adapter = McpToGenAiToolAdapter(
|
||||
session=MockMcpClientSession(),
|
||||
list_tools_result=mcp_types.ListToolsResult(tools=[]),
|
||||
)
|
||||
response = GenerateContentResponse(
|
||||
candidates=[
|
||||
Candidate(
|
||||
content=Content(
|
||||
parts=[
|
||||
Part(
|
||||
function_call=FunctionCall(
|
||||
name='tool',
|
||||
args={'key1': 'value1', 'key2': 1},
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
function_map = {'tool': mcp_to_genai_tool_adapter}
|
||||
expected_parts = [
|
||||
Part(
|
||||
function_response=FunctionResponse(
|
||||
name='tool',
|
||||
response={
|
||||
'result': {
|
||||
'content': [{'type': 'text', 'text': '1.01'}],
|
||||
'isError': False,
|
||||
}
|
||||
},
|
||||
)
|
||||
)
|
||||
]
|
||||
actual_parts = await get_function_response_parts_async(response, function_map)
|
||||
|
||||
for actual_part, expected_part in zip(actual_parts, expected_parts):
|
||||
assert actual_part.model_dump_json(
|
||||
exclude_none=True
|
||||
) == expected_part.model_dump_json(exclude_none=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_tool_error():
|
||||
if not _is_mcp_imported:
|
||||
return
|
||||
|
||||
class MockMcpClientSession(McpClientSession):
|
||||
|
||||
def __init__(self):
|
||||
self._read_stream = None
|
||||
self._write_stream = None
|
||||
|
||||
async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
|
||||
return mcp_types.CallToolResult(
|
||||
content=[mcp_types.TextContent(type='text', text='Internal error')],
|
||||
isError=True,
|
||||
)
|
||||
|
||||
mcp_to_genai_tool_adapter = McpToGenAiToolAdapter(
|
||||
session=MockMcpClientSession(),
|
||||
list_tools_result=mcp_types.ListToolsResult(tools=[]),
|
||||
)
|
||||
response = GenerateContentResponse(
|
||||
candidates=[
|
||||
Candidate(
|
||||
content=Content(
|
||||
parts=[
|
||||
Part(
|
||||
function_call=FunctionCall(
|
||||
name='tool',
|
||||
args={'key1': 'value1', 'key2': 1},
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
function_map = {'tool': mcp_to_genai_tool_adapter}
|
||||
expected_parts = [
|
||||
Part(
|
||||
function_response=FunctionResponse(
|
||||
name='tool',
|
||||
response={
|
||||
'error': {
|
||||
'content': [{'type': 'text', 'text': 'Internal error'}],
|
||||
'isError': True,
|
||||
}
|
||||
},
|
||||
)
|
||||
)
|
||||
]
|
||||
actual_parts = await get_function_response_parts_async(response, function_map)
|
||||
|
||||
for actual_part, expected_part in zip(actual_parts, expected_parts):
|
||||
assert actual_part.model_dump_json(
|
||||
exclude_none=True
|
||||
) == expected_part.model_dump_json(exclude_none=True)
|
||||
@@ -0,0 +1,130 @@
|
||||
# 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 get_max_remote_calls_for_afc."""
|
||||
|
||||
from ... import types
|
||||
from ..._extra_utils import get_max_remote_calls_afc
|
||||
import pytest
|
||||
|
||||
|
||||
def test_config_is_none():
|
||||
assert get_max_remote_calls_afc(None) == 10
|
||||
|
||||
|
||||
def test_afc_unset_max_unset():
|
||||
assert get_max_remote_calls_afc(types.GenerateContentConfig()) == 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,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_afc_disabled_max_set():
|
||||
with pytest.raises(ValueError):
|
||||
get_max_remote_calls_afc(
|
||||
types.GenerateContentConfig(
|
||||
automatic_function_calling=types.AutomaticFunctionCallingConfig(
|
||||
disable=True,
|
||||
maximum_remote_calls=20,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_afc_d_max_unset():
|
||||
assert (
|
||||
get_max_remote_calls_afc(
|
||||
types.GenerateContentConfig(
|
||||
automatic_function_calling=types.AutomaticFunctionCallingConfig(
|
||||
disable=False,
|
||||
),
|
||||
)
|
||||
)
|
||||
== 10
|
||||
)
|
||||
|
||||
|
||||
def test_afc_d_max_set():
|
||||
assert (
|
||||
get_max_remote_calls_afc(
|
||||
types.GenerateContentConfig(
|
||||
automatic_function_calling=types.AutomaticFunctionCallingConfig(
|
||||
disable=False,
|
||||
maximum_remote_calls=5,
|
||||
),
|
||||
)
|
||||
)
|
||||
== 5
|
||||
)
|
||||
|
||||
|
||||
def test_afc_enabled_max_set_to_zero():
|
||||
with pytest.raises(ValueError):
|
||||
get_max_remote_calls_afc(
|
||||
types.GenerateContentConfig(
|
||||
automatic_function_calling=types.AutomaticFunctionCallingConfig(
|
||||
disable=False,
|
||||
maximum_remote_calls=0,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_afc_enabled_max_set_to_negative():
|
||||
with pytest.raises(ValueError):
|
||||
get_max_remote_calls_afc(
|
||||
types.GenerateContentConfig(
|
||||
automatic_function_calling=types.AutomaticFunctionCallingConfig(
|
||||
disable=False,
|
||||
maximum_remote_calls=-1,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_afc_enabled_max_set_to_float():
|
||||
assert (
|
||||
get_max_remote_calls_afc(
|
||||
types.GenerateContentConfig(
|
||||
automatic_function_calling=types.AutomaticFunctionCallingConfig(
|
||||
disable=False,
|
||||
maximum_remote_calls=5.0,
|
||||
),
|
||||
)
|
||||
)
|
||||
== 5
|
||||
)
|
||||
@@ -0,0 +1,241 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""Test invoke_function_from_dict_args."""
|
||||
|
||||
from typing import Union
|
||||
import pydantic
|
||||
import pytest
|
||||
import sys
|
||||
from ... import errors
|
||||
from ..._extra_utils import invoke_function_from_dict_args
|
||||
|
||||
|
||||
def test_builtin_primitive_types():
|
||||
def func_under_test(x: int, y: float, z: str, w: bool):
|
||||
return {
|
||||
'x': x + 1,
|
||||
'y': y + 1.0,
|
||||
'z': z + '1.0',
|
||||
'w': not w,
|
||||
}
|
||||
|
||||
original_args = {
|
||||
'x': 1,
|
||||
'y': 1.0,
|
||||
'z': '1.0',
|
||||
'w': True,
|
||||
}
|
||||
expected_response = {
|
||||
'x': 2,
|
||||
'y': 2.0,
|
||||
'z': '1.01.0',
|
||||
'w': False,
|
||||
}
|
||||
|
||||
actual_response = invoke_function_from_dict_args(
|
||||
original_args, func_under_test
|
||||
)
|
||||
|
||||
assert actual_response == expected_response
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.version_info < (3, 10),
|
||||
reason='| is only supported in Python 3.9 and above.',
|
||||
)
|
||||
def test_builtin_compound_types():
|
||||
def func_under_test(x: list[int], y: dict[str, float]):
|
||||
return {
|
||||
'new_x': x + [3],
|
||||
'new_y': y | {'key3': 3.0},
|
||||
}
|
||||
|
||||
original_args = {
|
||||
'x': [1, 2],
|
||||
'y': {'key1': 1.0, 'key2': 2.0},
|
||||
}
|
||||
expected_response = {
|
||||
'new_x': [1, 2, 3],
|
||||
'new_y': {'key1': 1.0, 'key2': 2.0, 'key3': 3.0},
|
||||
}
|
||||
|
||||
actual_response = invoke_function_from_dict_args(
|
||||
original_args, func_under_test
|
||||
)
|
||||
|
||||
assert actual_response == expected_response
|
||||
|
||||
|
||||
def test_nested_pydantic_model():
|
||||
class SimpleModel(pydantic.BaseModel):
|
||||
key1_simple: int
|
||||
key2_simple: float
|
||||
|
||||
class ComplexModel(pydantic.BaseModel):
|
||||
key1_complex: SimpleModel
|
||||
key2_complex: list[SimpleModel]
|
||||
key3_complex: dict[str, SimpleModel]
|
||||
|
||||
def func_under_test(x: ComplexModel):
|
||||
y = x.model_copy()
|
||||
y.key1_complex.key1_simple += 1
|
||||
y.key1_complex.key2_simple += 1.0
|
||||
for simple_model in y.key2_complex:
|
||||
simple_model.key1_simple += 1
|
||||
simple_model.key2_simple += 1.0
|
||||
for simple_model in y.key3_complex.values():
|
||||
simple_model.key1_simple += 1
|
||||
simple_model.key2_simple += 1.0
|
||||
return y
|
||||
|
||||
original_args = {
|
||||
'x': {
|
||||
'key1_complex': {'key1_simple': 1, 'key2_simple': 1.0},
|
||||
'key2_complex': [
|
||||
{'key1_simple': 2, 'key2_simple': 2.0},
|
||||
{'key1_simple': 3, 'key2_simple': 3.0},
|
||||
],
|
||||
'key3_complex': {
|
||||
'key1_simple': {'key1_simple': 4, 'key2_simple': 4.0},
|
||||
'key2_simple': {'key1_simple': 5, 'key2_simple': 5.0},
|
||||
},
|
||||
}
|
||||
}
|
||||
expected_response = {
|
||||
'key1_complex': {'key1_simple': 2, 'key2_simple': 2.0},
|
||||
'key2_complex': [
|
||||
{'key1_simple': 3, 'key2_simple': 3.0},
|
||||
{'key1_simple': 4, 'key2_simple': 4.0},
|
||||
],
|
||||
'key3_complex': {
|
||||
'key1_simple': {'key1_simple': 5, 'key2_simple': 5.0},
|
||||
'key2_simple': {'key1_simple': 6, 'key2_simple': 6.0},
|
||||
},
|
||||
}
|
||||
|
||||
actual_response = invoke_function_from_dict_args(
|
||||
original_args, func_under_test
|
||||
)
|
||||
|
||||
assert isinstance(actual_response, ComplexModel)
|
||||
assert isinstance(actual_response.key1_complex, SimpleModel)
|
||||
assert isinstance(actual_response.key2_complex[0], SimpleModel)
|
||||
assert isinstance(actual_response.key2_complex[1], SimpleModel)
|
||||
assert isinstance(actual_response.key3_complex['key1_simple'], SimpleModel)
|
||||
assert isinstance(actual_response.key3_complex['key2_simple'], SimpleModel)
|
||||
assert actual_response.model_dump() == expected_response
|
||||
|
||||
|
||||
def test_pydantic_model_in_list_union_type():
|
||||
class SimpleModel(pydantic.BaseModel):
|
||||
key1_simple: int
|
||||
key2_simple: float
|
||||
|
||||
def func_under_test(x: list[Union[int, SimpleModel]]):
|
||||
result = []
|
||||
for item in x:
|
||||
if isinstance(item, int):
|
||||
result.append(item + 1)
|
||||
elif isinstance(item, SimpleModel):
|
||||
result.append(item.model_copy())
|
||||
result[-1].key1_simple += 1
|
||||
result[-1].key2_simple += 1.0
|
||||
else:
|
||||
raise ValueError('Unsupported type: %s' % type(item))
|
||||
return result
|
||||
|
||||
expected_response = [2, {'key1_simple': 2, 'key2_simple': 2.0}]
|
||||
original_args = {'x': [1, {'key1_simple': 1, 'key2_simple': 1.0}]}
|
||||
actual_response = invoke_function_from_dict_args(
|
||||
original_args, func_under_test
|
||||
)
|
||||
assert isinstance(actual_response, list)
|
||||
assert len(actual_response) == 2
|
||||
assert isinstance(actual_response[0], int)
|
||||
assert isinstance(actual_response[1], SimpleModel)
|
||||
assert actual_response[0] == expected_response[0]
|
||||
assert actual_response[1].model_dump() == expected_response[1]
|
||||
|
||||
|
||||
def test_unknown_pydantic_model_argument():
|
||||
class SimpleModel(pydantic.BaseModel):
|
||||
key1_simple: int
|
||||
key2_simple: float
|
||||
|
||||
def func_under_test(x: SimpleModel):
|
||||
return x.model_copy()
|
||||
|
||||
original_args = {'x': {'key3_simple': 1, 'key2_simple': 1.0}}
|
||||
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
invoke_function_from_dict_args(original_args, func_under_test)
|
||||
|
||||
|
||||
def test_unknown_pydantic_model_argument_with_union_type():
|
||||
class SimpleModel1(pydantic.BaseModel):
|
||||
key1_simple: int
|
||||
key2_simple: float
|
||||
|
||||
class SimpleModel2(pydantic.BaseModel):
|
||||
key3_simple: str
|
||||
key4_simple: float
|
||||
|
||||
def func_under_test(x: Union[SimpleModel1, SimpleModel2]):
|
||||
return x.model_copy()
|
||||
|
||||
original_args = {'x': {'key5_simple': 1, 'key4_simple': 1.0}}
|
||||
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
invoke_function_from_dict_args(original_args, func_under_test)
|
||||
|
||||
|
||||
def test_unknown_pydantic_model_argument_with_union_type_and_builtin_type():
|
||||
class SimpleModel1(pydantic.BaseModel):
|
||||
key1_simple: int
|
||||
key2_simple: float
|
||||
|
||||
def func_under_test(x: Union[SimpleModel1, int]):
|
||||
return x.model_copy()
|
||||
|
||||
original_args = {'x': {'key5_simple': 1, 'key4_simple': 1.0}}
|
||||
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
invoke_function_from_dict_args(original_args, func_under_test)
|
||||
|
||||
|
||||
def test_incompatible_value_and_annotation():
|
||||
def func_under_test(x: int):
|
||||
return x + 1
|
||||
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
invoke_function_from_dict_args({'x': {'k': 'v'}}, func_under_test)
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
invoke_function_from_dict_args({'x': 'a'}, func_under_test)
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
invoke_function_from_dict_args({'x': []}, func_under_test)
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
invoke_function_from_dict_args({'x': 1.0}, func_under_test)
|
||||
with pytest.raises(errors.UnknownFunctionCallArgumentError):
|
||||
invoke_function_from_dict_args({'x': {}}, func_under_test)
|
||||
|
||||
|
||||
def test_function_invocation_error():
|
||||
def func_under_test(x: int):
|
||||
return x / 0
|
||||
|
||||
with pytest.raises(errors.FunctionInvocationError):
|
||||
invoke_function_from_dict_args({'x': 1}, func_under_test)
|
||||
@@ -0,0 +1,159 @@
|
||||
# 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 raise_error_for_afc_incompatible_config."""
|
||||
|
||||
import pytest
|
||||
from ... import types
|
||||
from ..._extra_utils import raise_error_for_afc_incompatible_config
|
||||
|
||||
|
||||
def test_config_is_none():
|
||||
assert raise_error_for_afc_incompatible_config(None) is None
|
||||
|
||||
|
||||
def test_tool_config_config_unset():
|
||||
assert (
|
||||
raise_error_for_afc_incompatible_config(types.GenerateContentConfig(
|
||||
automatic_function_calling=types.AutomaticFunctionCallingConfig(
|
||||
disable=False,
|
||||
maximum_remote_calls=1,
|
||||
),
|
||||
))
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
def test_function_calling_config_unset():
|
||||
assert (
|
||||
raise_error_for_afc_incompatible_config(types.GenerateContentConfig(
|
||||
automatic_function_calling=types.AutomaticFunctionCallingConfig(
|
||||
disable=False,
|
||||
maximum_remote_calls=1,
|
||||
),
|
||||
tool_config=types.ToolConfig(),
|
||||
))
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
|
||||
def test_compatible_config_afc_disabled():
|
||||
assert (
|
||||
raise_error_for_afc_incompatible_config(
|
||||
types.GenerateContentConfig(
|
||||
automatic_function_calling=types.AutomaticFunctionCallingConfig(
|
||||
disable=True,
|
||||
),
|
||||
tool_config=types.ToolConfig(
|
||||
function_calling_config=types.FunctionCallingConfig(
|
||||
stream_function_call_arguments=False,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
def test_compatible_config_stream_function_call_arguments_unset_afc_unset():
|
||||
assert (
|
||||
raise_error_for_afc_incompatible_config(
|
||||
types.GenerateContentConfig(
|
||||
tool_config=types.ToolConfig(
|
||||
function_calling_config=types.FunctionCallingConfig(
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
def test_compatible_config_stream_function_call_arguments_unset_no_disable_afc():
|
||||
assert (
|
||||
raise_error_for_afc_incompatible_config(
|
||||
types.GenerateContentConfig(
|
||||
automatic_function_calling=types.AutomaticFunctionCallingConfig(
|
||||
),
|
||||
tool_config=types.ToolConfig(
|
||||
function_calling_config=types.FunctionCallingConfig(
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
),
|
||||
tool_config=types.ToolConfig(
|
||||
function_calling_config=types.FunctionCallingConfig(
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
),
|
||||
tool_config=types.ToolConfig(
|
||||
function_calling_config=types.FunctionCallingConfig(
|
||||
stream_function_call_arguments=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_incompatible_config_stream_function_call_arguments_set_no_afc_config():
|
||||
with pytest.raises(ValueError):
|
||||
raise_error_for_afc_incompatible_config(
|
||||
types.GenerateContentConfig(
|
||||
tool_config=types.ToolConfig(
|
||||
function_calling_config=types.FunctionCallingConfig(
|
||||
stream_function_call_arguments=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,53 @@
|
||||
# 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
|
||||
@@ -0,0 +1,214 @@
|
||||
# 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 types
|
||||
from ..._extra_utils import should_disable_afc
|
||||
|
||||
|
||||
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
|
||||
)
|
||||
@@ -0,0 +1,17 @@
|
||||
# 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 batches module."""
|
||||
@@ -0,0 +1,77 @@
|
||||
# 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 batches.cancel()."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
|
||||
# Vertex AI batch job name.
|
||||
_BATCH_JOB_NAME = '6339625664542408704'
|
||||
_BATCH_JOB_FULL_RESOURCE_NAME = (
|
||||
'projects/964831358985/locations/us-central1/'
|
||||
f'batchPredictionJobs/{_BATCH_JOB_NAME}'
|
||||
)
|
||||
# MLDev batch operation name.
|
||||
_MLDEV_BATCH_OPERATION_NAME = 'batches/0yew7plxupyybd7appsrq5vw7w0lp3l79lab'
|
||||
_INVALID_BATCH_JOB_NAME = 'invalid_name'
|
||||
|
||||
|
||||
# All tests will be run for both Vertex and MLDev.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_cancel_batch_job',
|
||||
parameters=types._CancelBatchJobParameters(
|
||||
name=_BATCH_JOB_NAME,
|
||||
),
|
||||
exception_if_mldev='Invalid batch job name',
|
||||
skip_in_api_mode=('Cannot cancel a batch job multiple times in Vertex'),
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_cancel_batch_operation',
|
||||
parameters=types._CancelBatchJobParameters(
|
||||
name=_MLDEV_BATCH_OPERATION_NAME,
|
||||
),
|
||||
exception_if_vertex='Invalid batch job name',
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_cancel_batch_job_with_invalid_name',
|
||||
parameters=types._CancelBatchJobParameters(
|
||||
name=_INVALID_BATCH_JOB_NAME,
|
||||
),
|
||||
exception_if_mldev='Invalid batch job name',
|
||||
exception_if_vertex='Invalid batch job name',
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='batches.cancel',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_cancel(client):
|
||||
if client.vertexai:
|
||||
name = _BATCH_JOB_NAME
|
||||
else:
|
||||
name = _MLDEV_BATCH_OPERATION_NAME
|
||||
await client.aio.batches.cancel(name=name)
|
||||
@@ -0,0 +1,123 @@
|
||||
# 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 batches.create() with invalid source and destinations."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
|
||||
_GEMINI_MODEL = 'gemini-2.5-flash'
|
||||
_EMBEDDING_MODEL = 'gemini-embedding-001'
|
||||
_DISPLAY_NAME = 'test_batch'
|
||||
|
||||
_GENERATE_CONTENT_BQ_OUTPUT_PREFIX = (
|
||||
'bq://vertex-sdk-dev.unified_genai_tests_batches.generate_content_output'
|
||||
)
|
||||
|
||||
_EMBEDDING_BQ_INPUT_FILE = (
|
||||
'bq://vertex-sdk-dev.unified_genai_tests_batches.embedding_requests'
|
||||
)
|
||||
|
||||
|
||||
# All tests will be run for both Vertex and MLDev.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_union_with_invalid_src',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_GEMINI_MODEL,
|
||||
src='invalid_src',
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
'dest': _GENERATE_CONTENT_BQ_OUTPUT_PREFIX,
|
||||
},
|
||||
),
|
||||
exception_if_mldev='Unsupported source',
|
||||
exception_if_vertex='Unsupported source',
|
||||
has_union=True,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_union_with_invalid_dest',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_EMBEDDING_MODEL,
|
||||
src=_EMBEDDING_BQ_INPUT_FILE,
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
'dest': 'invalid_dest',
|
||||
},
|
||||
),
|
||||
exception_if_mldev='not supported in Gemini API',
|
||||
exception_if_vertex='Unsupported destination',
|
||||
has_union=True,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_create_with_webhook_config',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_GEMINI_MODEL,
|
||||
src={
|
||||
'inlined_requests': [
|
||||
{
|
||||
'contents': [
|
||||
{'parts': [{'text': 'say hello'}], 'role': 'user'}
|
||||
]
|
||||
},
|
||||
]
|
||||
},
|
||||
config=types.CreateBatchJobConfig(
|
||||
display_name=_DISPLAY_NAME,
|
||||
webhook_config=types.WebhookConfig(
|
||||
uris=['https://example.com/webhook'],
|
||||
user_metadata={'batch_id': '123'},
|
||||
),
|
||||
),
|
||||
),
|
||||
exception_if_vertex='not supported in Vertex AI',
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_create_with_webhook_config_dict',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_GEMINI_MODEL,
|
||||
src={
|
||||
'inlined_requests': [
|
||||
{
|
||||
'contents': [
|
||||
{'parts': [{'text': 'say hello'}], 'role': 'user'}
|
||||
]
|
||||
},
|
||||
]
|
||||
},
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
'webhook_config': {
|
||||
'uris': ['https://example.com/webhook'],
|
||||
},
|
||||
},
|
||||
),
|
||||
exception_if_vertex='not supported in Vertex AI',
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.usefixtures('mock_timestamped_unique_name'),
|
||||
pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='batches.create',
|
||||
test_table=test_table,
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1,113 @@
|
||||
# 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 batches.create() with BigQuery source."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
|
||||
_GEMINI_MODEL = 'gemini-2.5-flash'
|
||||
_GEMINI_MODEL_FULL_NAME = 'publishers/google/models/gemini-2.5-flash'
|
||||
_EMBEDDING_MODEL = 'gemini-embedding-001'
|
||||
_EMBEDDING_MODEL_FULL_NAME = 'publishers/google/models/gemini-embedding-001'
|
||||
_DISPLAY_NAME = 'test_batch'
|
||||
_GENERATE_CONTENT_BQ_INPUT_FILE = (
|
||||
'bq://vertex-sdk-dev.unified_genai_tests_batches.generate_content_requests'
|
||||
)
|
||||
_GENERATE_CONTENT_BQ_OUTPUT_PREFIX = (
|
||||
'bq://vertex-sdk-dev.unified_genai_tests_batches.generate_content_output'
|
||||
)
|
||||
|
||||
_EMBEDDING_BQ_INPUT_FILE = (
|
||||
'bq://vertex-sdk-dev.unified_genai_tests_batches.embedding_requests'
|
||||
)
|
||||
_EMBEDDING_BQ_OUTPUT_PREFIX = (
|
||||
'bq://vertex-sdk-dev.unified_genai_tests_batches.embedding_output'
|
||||
)
|
||||
|
||||
|
||||
# All tests will be run for both Vertex and MLDev.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_union_generate_content_with_bigquery',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_GEMINI_MODEL_FULL_NAME,
|
||||
src=_GENERATE_CONTENT_BQ_INPUT_FILE,
|
||||
),
|
||||
exception_if_mldev='not supported in Gemini API',
|
||||
has_union=True,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_union_embedding_with_bigquery',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model='publishers/google/models/text-embedding-004',
|
||||
src=_EMBEDDING_BQ_INPUT_FILE,
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
'dest': _EMBEDDING_BQ_OUTPUT_PREFIX,
|
||||
},
|
||||
),
|
||||
exception_if_mldev='not supported in Gemini API',
|
||||
has_union=True,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_generate_content_with_bigquery',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_GEMINI_MODEL_FULL_NAME,
|
||||
src={
|
||||
'bigquery_uri': _GENERATE_CONTENT_BQ_INPUT_FILE,
|
||||
'format': 'bigquery',
|
||||
},
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
'dest': {
|
||||
'bigquery_uri': _GENERATE_CONTENT_BQ_OUTPUT_PREFIX,
|
||||
'format': 'bigquery',
|
||||
},
|
||||
},
|
||||
),
|
||||
exception_if_mldev='one of',
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.usefixtures('mock_timestamped_unique_name'),
|
||||
pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='batches.create',
|
||||
test_table=test_table,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_create(client):
|
||||
with pytest_helper.exception_if_mldev(client, ValueError):
|
||||
batch_job = await client.aio.batches.create(
|
||||
model=_GEMINI_MODEL,
|
||||
src=_GENERATE_CONTENT_BQ_INPUT_FILE,
|
||||
)
|
||||
|
||||
assert batch_job.name.startswith('projects/')
|
||||
assert (
|
||||
batch_job.model == _GEMINI_MODEL_FULL_NAME
|
||||
) # Converted to Vertex full name.
|
||||
assert batch_job.src.bigquery_uri == _GENERATE_CONTENT_BQ_INPUT_FILE
|
||||
assert batch_job.src.format == 'bigquery'
|
||||
@@ -0,0 +1,82 @@
|
||||
# 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 batches.create() with file source."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
|
||||
_GEMINI_MODEL = 'gemini-2.5-flash'
|
||||
_DISPLAY_NAME = 'test_batch'
|
||||
|
||||
_MLDEV_GEMINI_MODEL = 'gemini-2.5-flash'
|
||||
_FILE_NAME = 'files/s0pa54alni6w'
|
||||
|
||||
# All tests will be run for both Vertex and MLDev.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_union_generate_content_with_file',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_MLDEV_GEMINI_MODEL,
|
||||
src={'file_name': _FILE_NAME},
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
},
|
||||
),
|
||||
exception_if_vertex='not supported',
|
||||
has_union=True,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_generate_content_with_file',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_MLDEV_GEMINI_MODEL,
|
||||
src={'file_name': _FILE_NAME},
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
},
|
||||
),
|
||||
exception_if_vertex='not supported',
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.usefixtures('mock_timestamped_unique_name'),
|
||||
pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='batches.create',
|
||||
test_table=test_table,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_create(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
batch_job = await client.aio.batches.create(
|
||||
model=_MLDEV_GEMINI_MODEL,
|
||||
src={'file_name': _FILE_NAME},
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
},
|
||||
)
|
||||
assert batch_job.name.startswith('batches/')
|
||||
assert (
|
||||
batch_job.model == 'models/' + _MLDEV_GEMINI_MODEL
|
||||
) # Converted to Gemini full name.
|
||||
@@ -0,0 +1,125 @@
|
||||
# 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 batches.create() with Google Cloud Storage source."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
|
||||
_GEMINI_MODEL = 'gemini-2.5-flash'
|
||||
_GEMINI_MODEL_FULL_NAME = 'publishers/google/models/gemini-2.5-flash'
|
||||
_EMBEDDING_MODEL = 'gemini-embedding-001'
|
||||
_EMBEDDING_MODEL_FULL_NAME = 'publishers/google/models/gemini-embedding-001'
|
||||
_DISPLAY_NAME = 'test_batch'
|
||||
|
||||
_GENERATE_CONTENT_GCS_INPUT_FILE = (
|
||||
'gs://unified-genai-tests/batches/input/generate_content_requests.jsonl'
|
||||
)
|
||||
_GENERATE_CONTENT_GCS_OUTPUT_PREFIX = 'gs://unified-genai-tests/batches/output'
|
||||
|
||||
_EMBEDDING_GCS_INPUT_FILE = (
|
||||
'gs://unified-genai-tests/batches/input/embedding_requests.jsonl'
|
||||
)
|
||||
_EMBEDDING_GCS_OUTPUT_PREFIX = 'gs://unified-genai-tests/batches/output'
|
||||
|
||||
|
||||
# All tests will be run for both Vertex and MLDev.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_union_generate_content_with_gcs',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_GEMINI_MODEL,
|
||||
src=_GENERATE_CONTENT_GCS_INPUT_FILE,
|
||||
),
|
||||
exception_if_mldev='not supported in Gemini API',
|
||||
has_union=True,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_union_embedding_with_gcs',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model='text-embedding-004',
|
||||
src=_EMBEDDING_GCS_INPUT_FILE,
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
'dest': _EMBEDDING_GCS_OUTPUT_PREFIX,
|
||||
},
|
||||
),
|
||||
exception_if_mldev='not supported in Gemini API',
|
||||
has_union=True,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_union_with_http_options',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_GEMINI_MODEL,
|
||||
src=_GENERATE_CONTENT_GCS_INPUT_FILE,
|
||||
config={
|
||||
'http_options': {
|
||||
'api_version': 'v1',
|
||||
'headers': {'test': 'headers'},
|
||||
},
|
||||
},
|
||||
),
|
||||
exception_if_mldev='not supported in Gemini API',
|
||||
has_union=True,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_generate_content_with_gcs',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_GEMINI_MODEL,
|
||||
src={
|
||||
'gcs_uri': [_GENERATE_CONTENT_GCS_INPUT_FILE],
|
||||
'format': 'jsonl',
|
||||
},
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
'dest': {
|
||||
'gcs_uri': _GENERATE_CONTENT_GCS_OUTPUT_PREFIX,
|
||||
'format': 'jsonl',
|
||||
},
|
||||
},
|
||||
),
|
||||
exception_if_mldev='one of',
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.usefixtures('mock_timestamped_unique_name'),
|
||||
pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='batches.create',
|
||||
test_table=test_table,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_create(client):
|
||||
with pytest_helper.exception_if_mldev(client, ValueError):
|
||||
batch_job = await client.aio.batches.create(
|
||||
model=_GEMINI_MODEL,
|
||||
src=_GENERATE_CONTENT_GCS_INPUT_FILE,
|
||||
)
|
||||
|
||||
assert batch_job.name.startswith('projects/')
|
||||
assert (
|
||||
batch_job.model == _GEMINI_MODEL_FULL_NAME
|
||||
) # Converted to Vertex full name.
|
||||
assert batch_job.src.gcs_uri == [_GENERATE_CONTENT_GCS_INPUT_FILE]
|
||||
assert batch_job.src.format == 'jsonl'
|
||||
@@ -0,0 +1,271 @@
|
||||
# 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 batches.create() with inlined requests."""
|
||||
import base64
|
||||
import copy
|
||||
import datetime
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import _transformers as t
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
_GEMINI_MODEL = 'gemini-2.5-flash'
|
||||
_DISPLAY_NAME = 'test_batch'
|
||||
|
||||
_MLDEV_GEMINI_MODEL = 'gemini-2.5-flash'
|
||||
|
||||
_SAFETY_SETTINGS = [
|
||||
{
|
||||
'category': 'HARM_CATEGORY_HATE_SPEECH',
|
||||
'threshold': 'BLOCK_ONLY_HIGH',
|
||||
},
|
||||
{
|
||||
'category': 'HARM_CATEGORY_DANGEROUS_CONTENT',
|
||||
'threshold': 'BLOCK_LOW_AND_ABOVE',
|
||||
},
|
||||
]
|
||||
|
||||
_INLINED_REQUESTS = [
|
||||
{
|
||||
'contents': [{
|
||||
'parts': [{
|
||||
'text': 'what is the number after 1? return just the number.',
|
||||
}],
|
||||
'role': 'user',
|
||||
}],
|
||||
'metadata': {
|
||||
'key': 'request-1',
|
||||
},
|
||||
'config': {
|
||||
'safety_settings': _SAFETY_SETTINGS,
|
||||
},
|
||||
},
|
||||
{
|
||||
'contents': [{
|
||||
'parts': [{
|
||||
'text': 'what is the number after 2? return just the number.',
|
||||
}],
|
||||
'role': 'user',
|
||||
}],
|
||||
'metadata': {
|
||||
'key': 'request-2',
|
||||
},
|
||||
'config': {
|
||||
'safety_settings': _SAFETY_SETTINGS,
|
||||
},
|
||||
},
|
||||
]
|
||||
_INLINED_TEXT_REQUEST_UNION = {
|
||||
'contents': [{
|
||||
'parts': [{
|
||||
'text': 'high',
|
||||
}],
|
||||
'role': 'user',
|
||||
}],
|
||||
'config': {
|
||||
'response_modalities': ['TEXT'],
|
||||
'system_instruction': 'I say high, you say low',
|
||||
'thinking_config': {
|
||||
'include_thoughts': True,
|
||||
'thinking_budget': 4000,
|
||||
},
|
||||
},
|
||||
}
|
||||
_INLINED_TEXT_REQUEST = copy.deepcopy(_INLINED_TEXT_REQUEST_UNION)
|
||||
_INLINED_TEXT_REQUEST['config']['system_instruction'] = t.t_content(
|
||||
'I say high, you say low'
|
||||
)
|
||||
_INLINED_TEXT_REQUEST['config']['tools'] = [{'google_search': {}}]
|
||||
_INLINED_TEXT_REQUEST['config']['tool_config'] = {
|
||||
'retrieval_config': {'lat_lng': {'latitude': 37.422, 'longitude': -122.084}}
|
||||
}
|
||||
|
||||
_INLINED_IMAGE_REQUEST = {
|
||||
'contents': [{
|
||||
'parts': [
|
||||
{'text': 'What is in this image?'},
|
||||
{
|
||||
'file_data': {
|
||||
'file_uri': (
|
||||
'https://generativelanguage.googleapis.com/v1beta/files/kje1wewvo85z'
|
||||
),
|
||||
'mime_type': 'image/jpeg',
|
||||
},
|
||||
},
|
||||
],
|
||||
'role': 'user',
|
||||
}],
|
||||
'config': {
|
||||
'temperature': 0.7,
|
||||
'top_p': 0.9,
|
||||
'top_k': 10,
|
||||
},
|
||||
}
|
||||
_INLINED_VIDEO_REQUEST = {
|
||||
'contents': [{
|
||||
'parts': [
|
||||
{
|
||||
'text': 'Summerize this video.',
|
||||
},
|
||||
{
|
||||
'file_data': {
|
||||
'file_uri': (
|
||||
'https://generativelanguage.googleapis.com/v1beta/files/tyvaih24jwje'
|
||||
),
|
||||
'mime_type': 'video/mp4',
|
||||
},
|
||||
'video_metadata': {
|
||||
'start_offset': '0s',
|
||||
'end_offset': '5s',
|
||||
'fps': 3,
|
||||
},
|
||||
},
|
||||
],
|
||||
'role': 'user',
|
||||
}],
|
||||
}
|
||||
_IMAGE_PNG_FILE_PATH = os.path.abspath(
|
||||
os.path.join(os.path.dirname(__file__), '../data/google.png')
|
||||
)
|
||||
with open(_IMAGE_PNG_FILE_PATH, 'rb') as image_file:
|
||||
image_bytes = image_file.read()
|
||||
image_string = base64.b64encode(image_bytes).decode('utf-8')
|
||||
_INLINED_IMAGE_BLOB_REQUEST = types.InlinedRequest(
|
||||
contents=[
|
||||
types.Content(
|
||||
parts=[
|
||||
types.Part(text='What is this image about?'),
|
||||
types.Part(
|
||||
inline_data=types.Blob(
|
||||
data=image_string,
|
||||
mime_type='image/png',
|
||||
),
|
||||
),
|
||||
],
|
||||
role='user',
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# All tests will be run for both Vertex and MLDev.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_union_with_inlined_request',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_MLDEV_GEMINI_MODEL,
|
||||
src=_INLINED_REQUESTS,
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
},
|
||||
),
|
||||
exception_if_vertex='not supported in Vertex',
|
||||
has_union=True,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_with_inlined_request',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_MLDEV_GEMINI_MODEL,
|
||||
src={'inlined_requests': _INLINED_REQUESTS},
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
},
|
||||
),
|
||||
exception_if_vertex='not supported',
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_with_inlined_request_config',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_MLDEV_GEMINI_MODEL,
|
||||
src={'inlined_requests': _INLINED_REQUESTS},
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
},
|
||||
),
|
||||
exception_if_vertex='not supported',
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_union_with_inlined_request_system_instruction',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_MLDEV_GEMINI_MODEL,
|
||||
src={'inlined_requests': [_INLINED_TEXT_REQUEST_UNION]},
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
},
|
||||
),
|
||||
has_union=True,
|
||||
exception_if_vertex='not supported',
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_with_image_file',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_MLDEV_GEMINI_MODEL,
|
||||
src={'inlined_requests': [_INLINED_IMAGE_REQUEST]},
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
},
|
||||
),
|
||||
exception_if_vertex='not supported',
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_with_image_blob',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_MLDEV_GEMINI_MODEL,
|
||||
src={'inlined_requests': [_INLINED_IMAGE_BLOB_REQUEST]},
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
},
|
||||
),
|
||||
exception_if_vertex='not supported',
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_with_video_file',
|
||||
parameters=types._CreateBatchJobParameters(
|
||||
model=_MLDEV_GEMINI_MODEL,
|
||||
src={'inlined_requests': [_INLINED_VIDEO_REQUEST]},
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
},
|
||||
),
|
||||
exception_if_vertex='not supported',
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.usefixtures('mock_timestamped_unique_name'),
|
||||
pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='batches.create',
|
||||
test_table=test_table,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_create(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
batch_job = await client.aio.batches.create(
|
||||
model=_GEMINI_MODEL,
|
||||
src=_INLINED_REQUESTS,
|
||||
)
|
||||
assert batch_job.name.startswith('batches/')
|
||||
assert (
|
||||
batch_job.model == 'models/' + _GEMINI_MODEL
|
||||
) # Converted to Gemini full name.
|
||||
@@ -0,0 +1,86 @@
|
||||
# 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 batches.delete()."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
|
||||
_BATCH_JOB_NAME = '7085929781874655232'
|
||||
_BATCH_JOB_FULL_RESOURCE_NAME = (
|
||||
'projects/964831358985/locations/us-central1/'
|
||||
f'batchPredictionJobs/{_BATCH_JOB_NAME}'
|
||||
)
|
||||
_INVALID_BATCH_JOB_NAME = 'invalid_name'
|
||||
_MLDEV_BATCH_OPERATION_NAME = 'batches/70h2jo0ic2t1zejyl0p4jgi8mk1gj0wvjusv'
|
||||
|
||||
|
||||
# All tests will be run for both Vertex and MLDev.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_delete_batch_job',
|
||||
parameters=types._DeleteBatchJobParameters(
|
||||
name=_BATCH_JOB_NAME,
|
||||
),
|
||||
exception_if_mldev='Invalid batch job name',
|
||||
skip_in_api_mode=('Cannot cancel a batch job multiple times in Vertex'),
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_delete_batch_job_full_resource_name',
|
||||
override_replay_id='test_delete_batch_job',
|
||||
parameters=types._DeleteBatchJobParameters(
|
||||
name=_BATCH_JOB_FULL_RESOURCE_NAME,
|
||||
),
|
||||
exception_if_mldev='Invalid batch job name',
|
||||
skip_in_api_mode=('Cannot cancel a batch job multiple times in Vertex'),
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_delete_batch_job_with_invalid_name',
|
||||
parameters=types._DeleteBatchJobParameters(
|
||||
name=_INVALID_BATCH_JOB_NAME,
|
||||
),
|
||||
exception_if_mldev='Invalid batch job name',
|
||||
exception_if_vertex='Invalid batch job name',
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_delete_batch_operation',
|
||||
parameters=types._DeleteBatchJobParameters(
|
||||
name=_MLDEV_BATCH_OPERATION_NAME,
|
||||
),
|
||||
exception_if_vertex='Invalid batch job name',
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='batches.delete',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_delete(client):
|
||||
if client.vertexai:
|
||||
name = _BATCH_JOB_NAME
|
||||
else:
|
||||
name = _MLDEV_BATCH_OPERATION_NAME
|
||||
|
||||
delete_job = await client.aio.batches.delete(name=name)
|
||||
assert delete_job
|
||||
@@ -0,0 +1,157 @@
|
||||
# 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 batches._create_embeddings()"""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
_MLDEV_EMBEDDING_BATCH_INLINE_OPERATION_NAME = (
|
||||
'batches/wdx71o8cgbzoa6gg3be1mg7g8ulrhapcjgo3'
|
||||
)
|
||||
_MLDEV_EMBEDDING_BATCH_FILE_OPERATION_NAME = (
|
||||
'batches/507oatd242het8ox60pwsmn7tcmtkrj8itff'
|
||||
)
|
||||
|
||||
_DISPLAY_NAME = 'test_batch'
|
||||
_MLDEV_EMBEDDING_MODEL = 'gemini-embedding-001'
|
||||
_EMBED_CONTENT_FILE_NAME = 'files/mq9e3mg3u2y5'
|
||||
_INLINED_EMBED_CONTENT_REQUESTS = {
|
||||
'config': {'output_dimensionality': 64},
|
||||
'contents': [
|
||||
{
|
||||
'parts': [{
|
||||
'text': '1',
|
||||
}],
|
||||
},
|
||||
{
|
||||
'parts': [{
|
||||
'text': '2',
|
||||
}],
|
||||
},
|
||||
{
|
||||
'parts': [{
|
||||
'text': '3',
|
||||
}],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_from_inlined',
|
||||
parameters=types._CreateEmbeddingsBatchJobParameters(
|
||||
model=_MLDEV_EMBEDDING_MODEL,
|
||||
src={'inlined_requests': _INLINED_EMBED_CONTENT_REQUESTS},
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
},
|
||||
),
|
||||
exception_if_vertex='Vertex AI does not support',
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.usefixtures('mock_timestamped_unique_name'),
|
||||
pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='batches.create_embeddings',
|
||||
test_table=test_table,
|
||||
http_options={
|
||||
'api_version': 'v1alpha',
|
||||
'base_url': (
|
||||
'https://autopush-generativelanguage.sandbox.googleapis.com'
|
||||
),
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_from_inline(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
batch_job = await client.aio.batches.create_embeddings(
|
||||
model=_MLDEV_EMBEDDING_MODEL,
|
||||
src={'inlined_requests': _INLINED_EMBED_CONTENT_REQUESTS},
|
||||
)
|
||||
assert batch_job.name.startswith('batches/')
|
||||
|
||||
|
||||
def test_from_file(client):
|
||||
"""Tests creating a batch job with an embedding file name."""
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
batch_job = client.batches.create_embeddings(
|
||||
model=_MLDEV_EMBEDDING_MODEL,
|
||||
src={'file_name': _EMBED_CONTENT_FILE_NAME},
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
},
|
||||
)
|
||||
assert batch_job.name.startswith('batches/')
|
||||
assert batch_job.model == 'models/' + _MLDEV_EMBEDDING_MODEL
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_from_file(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
batch_job = await client.aio.batches.create_embeddings(
|
||||
model=_MLDEV_EMBEDDING_MODEL,
|
||||
src={'file_name': _EMBED_CONTENT_FILE_NAME},
|
||||
config={
|
||||
'display_name': _DISPLAY_NAME,
|
||||
},
|
||||
)
|
||||
assert batch_job.name.startswith('batches/')
|
||||
assert (
|
||||
batch_job.model == 'models/' + _MLDEV_EMBEDDING_MODEL
|
||||
) # Converted to Gemini full name.
|
||||
|
||||
|
||||
def test_get_inline(client):
|
||||
"""Tests getting a batch job that used inline requests."""
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
name = _MLDEV_EMBEDDING_BATCH_INLINE_OPERATION_NAME
|
||||
batch_job = client.batches.get(name=name)
|
||||
assert batch_job.dest.inlined_embed_content_responses is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_get_inline(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
name = _MLDEV_EMBEDDING_BATCH_INLINE_OPERATION_NAME
|
||||
batch_job = await client.aio.batches.get(name=name)
|
||||
|
||||
assert batch_job.dest.inlined_embed_content_responses is not None
|
||||
|
||||
|
||||
def test_get_file(client):
|
||||
"""Tests getting a batch job that used a file source."""
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
name = _MLDEV_EMBEDDING_BATCH_FILE_OPERATION_NAME
|
||||
batch_job = client.batches.get(name=name)
|
||||
assert batch_job.dest.file_name is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_get_file(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
name = _MLDEV_EMBEDDING_BATCH_FILE_OPERATION_NAME
|
||||
batch_job = await client.aio.batches.get(name=name)
|
||||
|
||||
assert batch_job.dest.file_name is not None
|
||||
@@ -0,0 +1,78 @@
|
||||
# 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 batches.get()."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
|
||||
# Vertex AI batch job name.
|
||||
_BATCH_JOB_NAME = '5798522612028014592'
|
||||
_BATCH_JOB_FULL_RESOURCE_NAME = (
|
||||
'projects/964831358985/locations/us-central1/'
|
||||
f'batchPredictionJobs/{_BATCH_JOB_NAME}'
|
||||
)
|
||||
# MLDev batch operation name.
|
||||
_MLDEV_BATCH_OPERATION_NAME = 'batches/z2p8ksus4lyxt25rntl3fpd67p2niw4hfij5'
|
||||
_INVALID_BATCH_JOB_NAME = 'invalid_name'
|
||||
|
||||
|
||||
# All tests will be run for both Vertex and MLDev.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_get_batch_job',
|
||||
parameters=types._GetBatchJobParameters(
|
||||
name=_BATCH_JOB_NAME,
|
||||
),
|
||||
exception_if_mldev='Invalid batch job name',
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_get_batch_operation',
|
||||
parameters=types._GetBatchJobParameters(
|
||||
name=_MLDEV_BATCH_OPERATION_NAME,
|
||||
),
|
||||
exception_if_vertex='Invalid batch job name',
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_get_batch_job_with_invalid_name',
|
||||
parameters=types._GetBatchJobParameters(
|
||||
name=_INVALID_BATCH_JOB_NAME,
|
||||
),
|
||||
exception_if_mldev='Invalid batch job name',
|
||||
exception_if_vertex='Invalid batch job name',
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='batches.get',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_get(client):
|
||||
if client.vertexai:
|
||||
name = _BATCH_JOB_NAME
|
||||
else:
|
||||
name = _MLDEV_BATCH_OPERATION_NAME
|
||||
batch_job = await client.aio.batches.get(name=name)
|
||||
|
||||
assert batch_job
|
||||
@@ -0,0 +1,79 @@
|
||||
# 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 batches.list()."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import errors
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
|
||||
# All tests will be run for both Vertex and MLDev.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_list_batch_jobs',
|
||||
parameters=types._ListBatchJobsParameters(
|
||||
config=types.ListBatchJobsConfig(
|
||||
page_size=5,
|
||||
),
|
||||
),
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_list_batch_jobs_with_config',
|
||||
parameters=types._ListBatchJobsParameters(
|
||||
config=types.ListBatchJobsConfig(
|
||||
filter='display_name:"genai_*"',
|
||||
page_size=5,
|
||||
),
|
||||
),
|
||||
exception_if_mldev='filter parameter is not supported in Gemini API',
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='batches.list',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
def test_pager(client):
|
||||
batch_jobs = client.batches.list(config={'page_size': 10})
|
||||
assert 'content-type' in batch_jobs.sdk_http_response.headers
|
||||
assert batch_jobs.name == 'batch_jobs'
|
||||
assert batch_jobs.page_size == 10
|
||||
assert len(batch_jobs) <= 10
|
||||
|
||||
# Iterate through the first page. Otherwise, too many batches are returned.
|
||||
for _ in batch_jobs.page:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pager(client):
|
||||
batch_jobs = await client.aio.batches.list(config={'page_size': 10})
|
||||
|
||||
assert 'Content-Type' in batch_jobs.sdk_http_response.headers
|
||||
assert batch_jobs.name == 'batch_jobs'
|
||||
assert batch_jobs.page_size == 10
|
||||
assert len(batch_jobs) <= 10
|
||||
|
||||
# Iterate through the first page. Otherwise, too many batches are returned.
|
||||
for _ in batch_jobs.page:
|
||||
pass
|
||||
@@ -0,0 +1,17 @@
|
||||
# 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 caches module."""
|
||||
@@ -0,0 +1,29 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""Constants for caches tests."""
|
||||
|
||||
CACHED_CONTENT_NAME_MLDEV = 'cachedContents/o239k1gxzz0juy9wqstndhncr85krehehf551hqh'
|
||||
CACHED_CONTENT_NAME_VERTEX = 'cachedContents/2164089915711684608'
|
||||
|
||||
VERTEX_HTTP_OPTIONS = {
|
||||
'api_version': 'v1beta1',
|
||||
'base_url': 'https://us-central1-aiplatform.googleapis.com/',
|
||||
}
|
||||
MLDEV_HTTP_OPTIONS = {
|
||||
'api_version': 'v1beta',
|
||||
'base_url': 'https://generativelanguage.googleapis.com/',
|
||||
}
|
||||
@@ -0,0 +1,210 @@
|
||||
# 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 copy import deepcopy
|
||||
import datetime
|
||||
import pytest
|
||||
import sys
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
from ... import _transformers as t
|
||||
|
||||
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI = types._CreateCachedContentParameters(
|
||||
model='gemini-2.5-flash',
|
||||
config={
|
||||
'contents': [
|
||||
types.Content(
|
||||
role='user',
|
||||
parts=[
|
||||
types.Part(
|
||||
fileData=types.FileData(
|
||||
fileUri='gs://cloud-samples-data/generative-ai/pdf/2312.11805v3.pdf',
|
||||
mimeType='application/pdf',
|
||||
)
|
||||
),
|
||||
types.Part(
|
||||
fileData=types.FileData(
|
||||
fileUri='gs://cloud-samples-data/generative-ai/pdf/2403.05530.pdf',
|
||||
mimeType='application/pdf',
|
||||
)
|
||||
),
|
||||
],
|
||||
)
|
||||
],
|
||||
'system_instruction': t.t_content('What is the sum of the two pdfs?'),
|
||||
'display_name': 'test cache',
|
||||
'ttl': '86400s',
|
||||
},
|
||||
)
|
||||
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE = types._CreateCachedContentParameters(
|
||||
model='gemini-2.5-flash',
|
||||
config={
|
||||
'contents': [
|
||||
types.Content(
|
||||
role='user',
|
||||
parts=[
|
||||
types.Part(
|
||||
fileData=types.FileData(
|
||||
mimeType='application/pdf',
|
||||
fileUri='https://generativelanguage.googleapis.com/v1beta/files/v200dhvn15h7',
|
||||
)
|
||||
)
|
||||
],
|
||||
)
|
||||
],
|
||||
'system_instruction': t.t_content('What is the sum of the two pdfs?'),
|
||||
'display_name': 'test cache',
|
||||
'ttl': '86400s',
|
||||
},
|
||||
)
|
||||
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_PARTIAL_MODEL_1 = deepcopy(
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI
|
||||
)
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_PARTIAL_MODEL_1.model = (
|
||||
'models/gemini-2.5-flash'
|
||||
)
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_PARTIAL_MODEL_2 = deepcopy(
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI
|
||||
)
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_PARTIAL_MODEL_2.model = (
|
||||
'publishers/google/models/gemini-2.5-flash'
|
||||
)
|
||||
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE_PARTIAL_MODEL_1 = deepcopy(
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE
|
||||
)
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE_PARTIAL_MODEL_1.model = (
|
||||
'models/gemini-2.5-flash'
|
||||
)
|
||||
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_CMEK = deepcopy(
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI
|
||||
)
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_CMEK.config.kms_key_name = (
|
||||
'projects/test-project/locations/us-central1/keyRings/test-keyring/cryptoKeys/test-key'
|
||||
)
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
_EXPIRE_TIME = datetime.datetime.fromisoformat('2025-12-20T00:00:00Z')
|
||||
else:
|
||||
_EXPIRE_TIME = datetime.datetime.fromisoformat('2025-12-20T00:00:00+00:00')
|
||||
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_EXPIRE_TIME = deepcopy(
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI
|
||||
)
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_EXPIRE_TIME.config.ttl = None
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_EXPIRE_TIME.config.expire_time = (
|
||||
_EXPIRE_TIME
|
||||
)
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_EXPIRE_TIME.config.display_name = (
|
||||
'test cache'
|
||||
)
|
||||
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE_EXPIRE_TIME = deepcopy(
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE
|
||||
)
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE_EXPIRE_TIME.config.ttl = None
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE_EXPIRE_TIME.config.expire_time = (
|
||||
_EXPIRE_TIME
|
||||
)
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE_EXPIRE_TIME.config.display_name = (
|
||||
'test cache'
|
||||
)
|
||||
|
||||
|
||||
# Replay mode is not supported for caches tests due to the error message
|
||||
# inconsistency in api and replay mode.
|
||||
# To run api mode tests, use the following steps:
|
||||
# 1. First create the resource.
|
||||
# sh run_tests.sh pytest -s tests/caches/test_create.py --mode=api
|
||||
# 1.1 If mldev test_create fails, update the uploaded file using this colab
|
||||
# https://colab.sandbox.google.com/drive/1Fv6KGSs0cg6tlpcUHdsclHussXMEGOXk#scrollTo=RSKmFPx00MVL.
|
||||
# 2. Find the resource name in debugging print and change the resource name constants.py.
|
||||
# 3. Run and record get and update tests.
|
||||
# sh run_tests.sh pytest -s tests/caches/test_get.py --mode=api && sh run_tests.sh pytest -s tests/caches/test_update.py --mode=api && sh run_tests.sh pytest -s tests/caches/test_delete.py --mode=api
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_caches_create_with_gcs_uri',
|
||||
exception_if_mldev='INVALID_ARGUMENT',
|
||||
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_caches_create_with_gcs_uri_cmek',
|
||||
exception_if_mldev='not supported',
|
||||
exception_if_vertex='INVALID_ARGUMENT', # The key is invalid.
|
||||
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_CMEK,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_caches_create_with_gcs_uri_expire_time',
|
||||
exception_if_mldev='INVALID_ARGUMENT',
|
||||
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_EXPIRE_TIME,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_caches_create_model_partial_path_1',
|
||||
exception_if_mldev='INVALID_ARGUMENT',
|
||||
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_PARTIAL_MODEL_1,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_caches_create_model_partial_path_2',
|
||||
exception_if_mldev='404',
|
||||
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_PARTIAL_MODEL_2,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_caches_create_with_googleai_file',
|
||||
exception_if_vertex='Internal',
|
||||
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE,
|
||||
skip_in_api_mode='Create is not reproducible in the API mode.',
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_caches_create_with_googleai_file_expire_time',
|
||||
exception_if_vertex='Internal',
|
||||
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE_EXPIRE_TIME,
|
||||
skip_in_api_mode='Create is not reproducible in the API mode.',
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_caches_create_with_googleai_file_model_partial_path_1',
|
||||
exception_if_vertex='Internal',
|
||||
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE_PARTIAL_MODEL_1,
|
||||
skip_in_api_mode='Create is not reproducible in the API mode.',
|
||||
),
|
||||
]
|
||||
pytestmark = [
|
||||
pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='caches.create',
|
||||
test_table=test_table,
|
||||
),
|
||||
]
|
||||
pytest_plugins = ('pytest_asyncio',)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_googleai_file_create(client):
|
||||
if client._api_client.vertexai:
|
||||
with pytest.raises(Exception):
|
||||
await client.aio.caches.create(
|
||||
model=_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE.model,
|
||||
config=_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE.config,
|
||||
)
|
||||
else:
|
||||
await client.aio.caches.create(
|
||||
model=_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE.model,
|
||||
config=_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE.config,
|
||||
)
|
||||
@@ -0,0 +1,105 @@
|
||||
# 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 copy
|
||||
import pytest
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
from . import constants
|
||||
from ... import _transformers as t
|
||||
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI = types._CreateCachedContentParameters(
|
||||
model='gemini-2.5-flash',
|
||||
config={
|
||||
'contents': [
|
||||
types.Content(
|
||||
role='user',
|
||||
parts=[
|
||||
types.Part(
|
||||
fileData=types.FileData(
|
||||
fileUri='gs://cloud-samples-data/generative-ai/pdf/2312.11805v3.pdf',
|
||||
mimeType='application/pdf',
|
||||
)
|
||||
),
|
||||
types.Part(
|
||||
fileData=types.FileData(
|
||||
fileUri='gs://cloud-samples-data/generative-ai/pdf/2403.05530.pdf',
|
||||
mimeType='application/pdf',
|
||||
)
|
||||
),
|
||||
],
|
||||
)
|
||||
],
|
||||
'system_instruction': t.t_content('What is the sum of the two pdfs?'),
|
||||
'display_name': 'test cache',
|
||||
'ttl': '86400s',
|
||||
'http_options': constants.VERTEX_HTTP_OPTIONS,
|
||||
},
|
||||
)
|
||||
|
||||
_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE = types._CreateCachedContentParameters(
|
||||
model='gemini-2.5-flash',
|
||||
config={
|
||||
'contents': [
|
||||
types.Content(
|
||||
role='user',
|
||||
parts=[
|
||||
types.Part(
|
||||
fileData=types.FileData(
|
||||
mimeType='application/pdf',
|
||||
fileUri='https://generativelanguage.googleapis.com/v1beta/files/v200dhvn15h7',
|
||||
)
|
||||
)
|
||||
],
|
||||
)
|
||||
],
|
||||
'system_instruction': t.t_content('What is the sum of the two pdfs?'),
|
||||
'display_name': 'test cache',
|
||||
'ttl': '86400s',
|
||||
'http_options': constants.MLDEV_HTTP_OPTIONS,
|
||||
},
|
||||
)
|
||||
|
||||
# Replay mode is not supported for caches tests due to the error message
|
||||
# inconsistency in api and replay mode.
|
||||
# To run api mode tests, use the following steps:
|
||||
# 1. First create the resource.
|
||||
# sh run_tests.sh pytest -s tests/caches/test_create.py --mode=api
|
||||
# 1.1 If mldev test_create fails, update the uploaded file using this colab
|
||||
# https://colab.sandbox.google.com/drive/1Fv6KGSs0cg6tlpcUHdsclHussXMEGOXk#scrollTo=RSKmFPx00MVL.
|
||||
# 2. Find the resource name in debugging print and change the resource name constants.py.
|
||||
# 3. Run and record get and update tests.
|
||||
# sh run_tests.sh pytest -s tests/caches/test_get.py --mode=api && sh run_tests.sh pytest -s tests/caches/test_update.py --mode=api && sh run_tests.sh pytest -s tests/caches/test_delete.py --mode=api
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_caches_create_with_gcs_uri',
|
||||
exception_if_mldev='404',
|
||||
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_caches_create_with_googleai_file',
|
||||
exception_if_vertex='404',
|
||||
parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GOOGLEAI_FILE,
|
||||
skip_in_api_mode='Create is not reproducible in the API mode.',
|
||||
),
|
||||
]
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='caches.create',
|
||||
test_table=test_table,
|
||||
)
|
||||
pytest_plugins = ('pytest_asyncio',)
|
||||
@@ -0,0 +1,54 @@
|
||||
# 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 pytest
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
from . import constants
|
||||
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
skip_in_api_mode='Delete is not reproducible in the API mode.',
|
||||
name='test_caches_delete_with_vertex_cache_name',
|
||||
exception_if_mldev='PERMISSION_DENIED',
|
||||
parameters=types._DeleteCachedContentParameters(
|
||||
name=constants.CACHED_CONTENT_NAME_VERTEX,
|
||||
),
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
skip_in_api_mode='Delete is not reproducible in the API mode.',
|
||||
name='test_caches_delete_with_mldev_cache_name',
|
||||
exception_if_vertex='NOT_FOUND',
|
||||
parameters=types._DeleteCachedContentParameters(
|
||||
name=constants.CACHED_CONTENT_NAME_MLDEV,
|
||||
),
|
||||
),
|
||||
]
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='caches.delete',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_delete(client):
|
||||
if client._api_client.vertexai:
|
||||
await client.aio.caches.delete(name=constants.CACHED_CONTENT_NAME_VERTEX)
|
||||
else:
|
||||
await client.aio.caches.delete(name=constants.CACHED_CONTENT_NAME_MLDEV)
|
||||
@@ -0,0 +1,52 @@
|
||||
# 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 pytest
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
from . import constants
|
||||
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
skip_in_api_mode='Delete is not reproducible in the API mode.',
|
||||
name='test_caches_delete_with_vertex_cache_name',
|
||||
exception_if_mldev='404',
|
||||
parameters=types._DeleteCachedContentParameters(
|
||||
name=constants.CACHED_CONTENT_NAME_VERTEX,
|
||||
config={
|
||||
'http_options': constants.VERTEX_HTTP_OPTIONS,
|
||||
},
|
||||
),
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
skip_in_api_mode='Delete is not reproducible in the API mode.',
|
||||
name='test_caches_delete_with_mldev_cache_name',
|
||||
exception_if_vertex='404',
|
||||
parameters=types._DeleteCachedContentParameters(
|
||||
name=constants.CACHED_CONTENT_NAME_MLDEV,
|
||||
config={
|
||||
'http_options': constants.MLDEV_HTTP_OPTIONS,
|
||||
},
|
||||
),
|
||||
),
|
||||
]
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='caches.delete',
|
||||
test_table=test_table,
|
||||
)
|
||||
@@ -0,0 +1,94 @@
|
||||
# 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 pytest
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
from . import constants
|
||||
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
skip_in_api_mode='Get is not reproducible in the API mode.',
|
||||
name='test_caches_get_with_vertex_cache_name',
|
||||
exception_if_mldev='PERMISSION_DENIED',
|
||||
parameters=types._GetCachedContentParameters(
|
||||
name=constants.CACHED_CONTENT_NAME_VERTEX,
|
||||
),
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
skip_in_api_mode='Get is not reproducible in the API mode.',
|
||||
name='test_caches_get_with_mldev_cache_name',
|
||||
exception_if_vertex='NOT_FOUND',
|
||||
parameters=types._GetCachedContentParameters(
|
||||
name=constants.CACHED_CONTENT_NAME_MLDEV,
|
||||
),
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
skip_in_api_mode='Get is not reproducible in the API mode.',
|
||||
name='test_caches_get_with_mldev_cache_name_partial_1',
|
||||
exception_if_vertex='NOT_FOUND',
|
||||
parameters=types._GetCachedContentParameters(
|
||||
name='cachedContents/o239k1gxzz0juy9wqstndhncr85krehehf551hqh'
|
||||
),
|
||||
),
|
||||
]
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='caches.get',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_get(client):
|
||||
if client._api_client.vertexai:
|
||||
response = await client.aio.caches.get(
|
||||
name=constants.CACHED_CONTENT_NAME_VERTEX
|
||||
)
|
||||
assert response
|
||||
else:
|
||||
await client.aio.caches.get(name=constants.CACHED_CONTENT_NAME_MLDEV)
|
||||
|
||||
|
||||
def test_different_cache_name_formats(client):
|
||||
if client._api_client.vertexai:
|
||||
response1 = client.caches.get(
|
||||
name='projects/964831358985/locations/us-central1/cachedContents/2164089915711684608'
|
||||
)
|
||||
assert response1
|
||||
response2 = client.caches.get(
|
||||
name='locations/us-central1/cachedContents/2164089915711684608'
|
||||
)
|
||||
assert response2
|
||||
response3 = client.caches.get(
|
||||
name='cachedContents/2164089915711684608'
|
||||
)
|
||||
assert response3
|
||||
response4 = client.caches.get(
|
||||
name='2164089915711684608'
|
||||
)
|
||||
assert response4
|
||||
else:
|
||||
response1 = client.caches.get(
|
||||
name='cachedContents/o239k1gxzz0juy9wqstndhncr85krehehf551hqh'
|
||||
)
|
||||
assert response1
|
||||
response2 = client.caches.get(
|
||||
name='o239k1gxzz0juy9wqstndhncr85krehehf551hqh'
|
||||
)
|
||||
assert response2
|
||||
@@ -0,0 +1,52 @@
|
||||
# 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 pytest
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
from . import constants
|
||||
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
skip_in_api_mode='Get is not reproducible in the API mode.',
|
||||
name='test_caches_get_with_vertex_cache_name',
|
||||
exception_if_mldev='404',
|
||||
parameters=types._GetCachedContentParameters(
|
||||
name=constants.CACHED_CONTENT_NAME_VERTEX,
|
||||
config={
|
||||
'http_options': constants.VERTEX_HTTP_OPTIONS,
|
||||
},
|
||||
),
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
skip_in_api_mode='Get is not reproducible in the API mode.',
|
||||
name='test_caches_get_with_mldev_cache_name',
|
||||
exception_if_vertex='404',
|
||||
parameters=types._GetCachedContentParameters(
|
||||
name=constants.CACHED_CONTENT_NAME_MLDEV,
|
||||
config={
|
||||
'http_options': constants.MLDEV_HTTP_OPTIONS,
|
||||
},
|
||||
),
|
||||
),
|
||||
]
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='caches.get',
|
||||
test_table=test_table,
|
||||
)
|
||||
@@ -0,0 +1,68 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""Test caches list method."""
|
||||
|
||||
import pytest
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
skip_in_api_mode='List is not reproducible in the API mode.',
|
||||
name='test_caches_list',
|
||||
parameters=types._ListCachedContentsParameters(config={'page_size': 3}),
|
||||
),
|
||||
]
|
||||
pytestmark = [
|
||||
pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='caches.list',
|
||||
test_table=test_table,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def test_pager(client):
|
||||
cached_contents = client.caches.list(config={'page_size': 2})
|
||||
assert 'content-type' in cached_contents.sdk_http_response.headers
|
||||
assert cached_contents.name == 'cached_contents'
|
||||
assert cached_contents.page_size == 2
|
||||
assert len(cached_contents) <= 2
|
||||
|
||||
# Iterate through all the pages. Then next_page() should raise an exception.
|
||||
for _ in cached_contents:
|
||||
pass
|
||||
with pytest.raises(IndexError, match='No more pages to fetch.'):
|
||||
cached_contents.next_page()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pager(client):
|
||||
cached_contents = await client.aio.caches.list(config={'page_size': 2})
|
||||
|
||||
assert 'Content-Type' in cached_contents.sdk_http_response.headers
|
||||
assert cached_contents.name == 'cached_contents'
|
||||
assert cached_contents.page_size == 2
|
||||
assert len(cached_contents) <= 2
|
||||
|
||||
# Iterate through all the pages. Then next_page() should raise an exception.
|
||||
async for _ in cached_contents:
|
||||
pass
|
||||
with pytest.raises(IndexError, match='No more pages to fetch.'):
|
||||
await cached_contents.next_page()
|
||||
@@ -0,0 +1,70 @@
|
||||
# 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 pytest
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
from . import constants
|
||||
|
||||
_VERTEX_UPDATE_PARAMETERS = types._UpdateCachedContentParameters(
|
||||
name=constants.CACHED_CONTENT_NAME_VERTEX,
|
||||
config={
|
||||
'ttl': '7600s',
|
||||
},
|
||||
)
|
||||
_MLDEV_UPDATE_PARAMETERS = types._UpdateCachedContentParameters(
|
||||
name=constants.CACHED_CONTENT_NAME_MLDEV,
|
||||
config={
|
||||
'ttl': '7600s',
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
skip_in_api_mode='Update has permission issues in the API mode.',
|
||||
name='test_caches_update_with_vertex_cache_name',
|
||||
exception_if_mldev='PERMISSION_DENIED',
|
||||
parameters=_VERTEX_UPDATE_PARAMETERS,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
skip_in_api_mode='Update has permission issues in the API mode.',
|
||||
name='test_caches_update_with_mldev_cache_name',
|
||||
exception_if_vertex='NOT_FOUND',
|
||||
parameters=_MLDEV_UPDATE_PARAMETERS,
|
||||
),
|
||||
]
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='caches.update',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_update(client):
|
||||
if client._api_client.vertexai:
|
||||
response = await client.aio.caches.update(
|
||||
name=_VERTEX_UPDATE_PARAMETERS.name,
|
||||
config=_VERTEX_UPDATE_PARAMETERS.config,
|
||||
)
|
||||
assert response
|
||||
else:
|
||||
await client.aio.caches.update(
|
||||
name=_MLDEV_UPDATE_PARAMETERS.name,
|
||||
config=_MLDEV_UPDATE_PARAMETERS.config,
|
||||
)
|
||||
@@ -0,0 +1,58 @@
|
||||
# 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 pytest
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
from . import constants
|
||||
|
||||
_VERTEX_UPDATE_PARAMETERS = types._UpdateCachedContentParameters(
|
||||
name=constants.CACHED_CONTENT_NAME_VERTEX,
|
||||
config={
|
||||
'ttl': '7600s',
|
||||
'http_options': constants.VERTEX_HTTP_OPTIONS,
|
||||
},
|
||||
)
|
||||
_MLDEV_UPDATE_PARAMETERS = types._UpdateCachedContentParameters(
|
||||
name=constants.CACHED_CONTENT_NAME_MLDEV,
|
||||
config={
|
||||
'ttl': '7600s',
|
||||
'http_options': constants.MLDEV_HTTP_OPTIONS,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
skip_in_api_mode='Update has permission issues in the API mode.',
|
||||
name='test_caches_update_with_vertex_cache_name',
|
||||
exception_if_mldev='404',
|
||||
parameters=_VERTEX_UPDATE_PARAMETERS,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
skip_in_api_mode='Update has permission issues in the API mode.',
|
||||
name='test_caches_update_with_mldev_cache_name',
|
||||
exception_if_vertex='404',
|
||||
parameters=_MLDEV_UPDATE_PARAMETERS,
|
||||
),
|
||||
]
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='caches.update',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Tests for the Google GenAI SDK's chats module."""
|
||||
@@ -0,0 +1,598 @@
|
||||
# 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)
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
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
|
||||
@@ -0,0 +1,941 @@
|
||||
# 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 json
|
||||
import os
|
||||
import sys
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic import ValidationError
|
||||
import pytest
|
||||
|
||||
from .. import pytest_helper
|
||||
from ... import errors
|
||||
from ... import types
|
||||
|
||||
try:
|
||||
from mcp import types as mcp_types
|
||||
from mcp import ClientSession as McpClientSession
|
||||
except ImportError as e:
|
||||
import sys
|
||||
|
||||
if sys.version_info < (3, 10):
|
||||
raise ImportError(
|
||||
'MCP Tool requires Python 3.10 or above. Please upgrade your Python'
|
||||
' version.'
|
||||
) from e
|
||||
else:
|
||||
raise e
|
||||
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
)
|
||||
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 square_integer(given_integer: int) -> int:
|
||||
return given_integer*given_integer
|
||||
|
||||
|
||||
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.
|
||||
|
||||
Args:
|
||||
energetic: Whether the music is energetic or not.
|
||||
loud: Whether the music is loud or not.
|
||||
bpm: The beats per minute of the music.
|
||||
|
||||
Returns: The name of the song being played.
|
||||
"""
|
||||
print(f"Starting music! {energetic=} {loud=}, {bpm=}")
|
||||
return "Never gonna give you up."
|
||||
|
||||
def dim_lights(brightness: float) -> bool:
|
||||
"""Dim the lights.
|
||||
|
||||
Args:
|
||||
brightness: The brightness of the lights, 0.0 is off, 1.0 is full.
|
||||
"""
|
||||
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(
|
||||
'tell me a story in 100 words',
|
||||
)
|
||||
|
||||
|
||||
def test_part(client):
|
||||
chat = client.chats.create(model=MODEL_NAME)
|
||||
chat.send_message(
|
||||
types.Part.from_text(text='tell me a story in 100 words'),
|
||||
)
|
||||
|
||||
|
||||
def test_parts(client):
|
||||
chat = client.chats.create(model=MODEL_NAME)
|
||||
chat.send_message(
|
||||
[
|
||||
types.Part.from_text(text='tell me a US city'),
|
||||
types.Part.from_text(text='the city is in west coast'),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_image(client, image_jpeg):
|
||||
chat = client.chats.create(model=MODEL_NAME)
|
||||
chat.send_message(
|
||||
[
|
||||
'what is the image about?',
|
||||
image_jpeg,
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
},
|
||||
},
|
||||
)
|
||||
response1 = chat.send_message(
|
||||
'what is the sum of natural numbers from 1 to 100?',
|
||||
)
|
||||
has_thought1 = False
|
||||
if response1.candidates:
|
||||
for candidate in response1.candidates:
|
||||
for part in candidate.content.parts:
|
||||
if part.thought:
|
||||
has_thought1 = True
|
||||
break
|
||||
assert has_thought1
|
||||
|
||||
response2 = chat.send_message(
|
||||
'can you help me to understand the logic better?'
|
||||
)
|
||||
has_thought2 = False
|
||||
if response2.candidates:
|
||||
for candidate in response2.candidates:
|
||||
for part in candidate.content.parts:
|
||||
if part.thought:
|
||||
has_thought2 = True
|
||||
break
|
||||
assert has_thought2
|
||||
|
||||
|
||||
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,
|
||||
},
|
||||
},
|
||||
)
|
||||
has_thought1 = False
|
||||
for chunk in chat.send_message_stream(
|
||||
'what is the sum of natural numbers from 1 to 100?',
|
||||
):
|
||||
if chunk.candidates:
|
||||
for candidate in chunk.candidates:
|
||||
for part in candidate.content.parts:
|
||||
if part.thought:
|
||||
has_thought1 = True
|
||||
break
|
||||
assert has_thought1
|
||||
|
||||
has_thought2 = False
|
||||
for chunk in chat.send_message_stream(
|
||||
'can you help me to understand the logic better?'
|
||||
):
|
||||
if chunk.candidates:
|
||||
for candidate in chunk.candidates:
|
||||
for part in candidate.content.parts:
|
||||
if part.thought:
|
||||
has_thought2 = True
|
||||
break
|
||||
assert has_thought2
|
||||
|
||||
|
||||
def test_google_cloud_storage_uri(client):
|
||||
chat = client.chats.create(model=MODEL_NAME)
|
||||
with pytest_helper.exception_if_mldev(client, errors.ClientError):
|
||||
chat.send_message(
|
||||
[
|
||||
'what is the image about?',
|
||||
types.Part.from_uri(
|
||||
file_uri='gs://unified-genai-dev/imagen-inputs/google_small.png',
|
||||
mime_type='image/png',
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_uploaded_file_uri(client):
|
||||
chat = client.chats.create(model=MODEL_NAME)
|
||||
with pytest_helper.exception_if_vertex(client, errors.ClientError):
|
||||
chat.send_message(
|
||||
[
|
||||
'what is the image about?',
|
||||
types.Part.from_uri(
|
||||
file_uri='https://generativelanguage.googleapis.com/v1beta/files/az606f58k7zj',
|
||||
mime_type='image/png',
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_config_override(client):
|
||||
chat_config = {'candidate_count': 1}
|
||||
chat = client.chats.create(model=MODEL_NAME, config=chat_config)
|
||||
request_config = {'candidate_count': 2}
|
||||
request_config_response = chat.send_message(
|
||||
'tell me a story in 100 words',
|
||||
config=request_config)
|
||||
default_config_response = chat.send_message(
|
||||
'tell me a story in 100 words')
|
||||
|
||||
assert len(request_config_response.candidates) == 2
|
||||
assert len(default_config_response.candidates) == 1
|
||||
|
||||
|
||||
def test_history(client):
|
||||
history = [
|
||||
types.Content(
|
||||
role='user', parts=[types.Part.from_text(text='define a=5, b=10')]
|
||||
),
|
||||
types.Content(
|
||||
role='model',
|
||||
parts=[types.Part.from_text(text='Hello there! how can I help you?')],
|
||||
),
|
||||
]
|
||||
chat = client.chats.create(model=MODEL_NAME, history=history)
|
||||
chat.send_message('what is a + b?')
|
||||
|
||||
assert len(chat.get_history()) > 2
|
||||
|
||||
|
||||
def test_send_2_messages(client):
|
||||
chat = client.chats.create(model=MODEL_NAME)
|
||||
chat.send_message('write a python function to check if a year is a leap year')
|
||||
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 <Type.OBJECT: "OBJECT"> 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)
|
||||
chat.send_message('Turn this place into a party!')
|
||||
curated_history = chat.get_history()
|
||||
|
||||
assert len(curated_history) == 8
|
||||
assert curated_history[0].role == 'user'
|
||||
assert curated_history[0].parts[0].text == 'Turn this place into a party!'
|
||||
assert curated_history[1].role == 'model'
|
||||
assert len(curated_history[1].parts) == 3
|
||||
for part in curated_history[1].parts:
|
||||
assert part.function_call
|
||||
assert curated_history[2].role == 'user'
|
||||
assert len(curated_history[2].parts) == 3
|
||||
for part in curated_history[2].parts:
|
||||
assert part.function_response
|
||||
assert curated_history[3].role == 'model'
|
||||
assert len(curated_history[3].parts) == 3
|
||||
for part in curated_history[3].parts:
|
||||
assert part.function_call
|
||||
assert curated_history[4].role == 'user'
|
||||
assert len(curated_history[4].parts) == 3
|
||||
for part in curated_history[4].parts:
|
||||
assert part.function_response
|
||||
assert curated_history[5].role == 'model'
|
||||
assert len(curated_history[5].parts) == 3
|
||||
for part in curated_history[5].parts:
|
||||
assert part.function_call
|
||||
assert curated_history[6].role == 'user'
|
||||
assert len(curated_history[6].parts) == 3
|
||||
for part in curated_history[6].parts:
|
||||
assert part.function_response
|
||||
assert curated_history[7].role == 'model'
|
||||
assert len(curated_history[7].parts) == 3
|
||||
for part in curated_history[7].parts:
|
||||
assert part.function_call
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.version_info >= (3, 13),
|
||||
reason=(
|
||||
'object type is dumped as <Type.OBJECT: "OBJECT"> as opposed to'
|
||||
' "OBJECT" in Python 3.13'
|
||||
),
|
||||
)
|
||||
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!')
|
||||
curated_history = chat.get_history()
|
||||
|
||||
assert len(curated_history) == 8
|
||||
assert curated_history[0].role == 'user'
|
||||
assert curated_history[0].parts[0].text == 'Turn this place into a party!'
|
||||
assert curated_history[1].role == 'model'
|
||||
assert len(curated_history[1].parts) == 3
|
||||
for part in curated_history[1].parts:
|
||||
assert part.function_call
|
||||
assert curated_history[2].role == 'user'
|
||||
assert len(curated_history[2].parts) == 3
|
||||
for part in curated_history[2].parts:
|
||||
assert part.function_response
|
||||
assert curated_history[3].role == 'model'
|
||||
assert len(curated_history[3].parts) == 3
|
||||
for part in curated_history[3].parts:
|
||||
assert part.function_call
|
||||
assert curated_history[4].role == 'user'
|
||||
assert len(curated_history[4].parts) == 3
|
||||
for part in curated_history[4].parts:
|
||||
assert part.function_response
|
||||
assert curated_history[5].role == 'model'
|
||||
assert len(curated_history[5].parts) == 3
|
||||
for part in curated_history[5].parts:
|
||||
assert part.function_call
|
||||
assert curated_history[6].role == 'user'
|
||||
assert len(curated_history[6].parts) == 3
|
||||
for part in curated_history[6].parts:
|
||||
assert part.function_response
|
||||
assert curated_history[7].role == 'model'
|
||||
assert len(curated_history[7].parts) == 3
|
||||
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},
|
||||
},
|
||||
)
|
||||
chat.send_message(
|
||||
'Do the square of 3.',
|
||||
)
|
||||
chat_history = chat.get_history()
|
||||
|
||||
assert len(chat_history) == 2
|
||||
assert chat_history[0].role == 'user'
|
||||
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 == 'square_integer'
|
||||
assert chat_history[1].parts[0].function_call.args == {
|
||||
'given_integer': 3,
|
||||
}
|
||||
|
||||
|
||||
@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},
|
||||
},
|
||||
)
|
||||
await chat.send_message(
|
||||
'Do the square of 3.',
|
||||
)
|
||||
chat_history = chat.get_history()
|
||||
|
||||
assert len(chat_history) == 2
|
||||
assert chat_history[0].role == 'user'
|
||||
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 == 'square_integer'
|
||||
assert chat_history[1].parts[0].function_call.args == {
|
||||
'given_integer': 3,
|
||||
}
|
||||
|
||||
|
||||
def test_stream_text(client):
|
||||
chat = client.chats.create(model=MODEL_NAME)
|
||||
chunks = 0
|
||||
for chunk in chat.send_message_stream(
|
||||
'tell me a story in 100 words',
|
||||
):
|
||||
chunks += 1
|
||||
|
||||
assert chunks > 1
|
||||
|
||||
|
||||
def test_stream_part(client):
|
||||
chat = client.chats.create(model=MODEL_NAME)
|
||||
chunks = 0
|
||||
for chunk in chat.send_message_stream(
|
||||
types.Part.from_text(text='tell me a story in 100 words'),
|
||||
):
|
||||
chunks += 1
|
||||
|
||||
assert chunks > 1
|
||||
|
||||
|
||||
def test_stream_parts(client):
|
||||
chat = client.chats.create(model=MODEL_NAME)
|
||||
chunks = 0
|
||||
for chunk in chat.send_message_stream(
|
||||
[
|
||||
types.Part.from_text(text='tell me a story in 100 words'),
|
||||
types.Part.from_text(text='the story is about a car'),
|
||||
],
|
||||
):
|
||||
chunks += 1
|
||||
|
||||
assert chunks > 2
|
||||
|
||||
|
||||
def test_stream_config_override(client):
|
||||
chat_config = {'response_mime_type': 'text/plain'}
|
||||
chat = client.chats.create(model=MODEL_NAME, config=chat_config)
|
||||
request_config = {'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
|
||||
):
|
||||
request_config_text += chunk.text
|
||||
default_config_text = ''
|
||||
for chunk in chat.send_message_stream('tell me a story in 100 words'):
|
||||
default_config_text += chunk.text
|
||||
|
||||
assert json.loads(request_config_text)
|
||||
with pytest.raises(json.JSONDecodeError):
|
||||
json.loads(default_config_text)
|
||||
|
||||
|
||||
def test_stream_function_calling(client):
|
||||
chat = client.chats.create(
|
||||
model='gemini-2.0-flash-exp',
|
||||
config={'tools': [divide_intergers_with_customized_math_rule]},
|
||||
)
|
||||
# Now we support AFC.
|
||||
for chunk in chat.send_message_stream(
|
||||
'what is the result of 100/2?',
|
||||
):
|
||||
pass
|
||||
for chunk in 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,
|
||||
}
|
||||
|
||||
|
||||
def test_stream_send_2_messages(client):
|
||||
chat = client.chats.create(model=MODEL_NAME)
|
||||
for chunk in chat.send_message_stream(
|
||||
'write a python function to check if a year is a leap year'
|
||||
):
|
||||
pass
|
||||
|
||||
for chunk in chat.send_message_stream('write a unit test for the function'):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_text(client):
|
||||
chat = client.aio.chats.create(model=MODEL_NAME)
|
||||
await chat.send_message('tell me a story in 100 words')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_part(client):
|
||||
chat = client.aio.chats.create(model=MODEL_NAME)
|
||||
await chat.send_message(types.Part.from_text(text='tell me a story in 100 words'))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_parts(client):
|
||||
chat = client.aio.chats.create(model=MODEL_NAME)
|
||||
await chat.send_message(
|
||||
[
|
||||
types.Part.from_text(text='tell me a US city'),
|
||||
types.Part.from_text(text='the city is in west coast'),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_config_override(client):
|
||||
chat_config = {'candidate_count': 1}
|
||||
chat = client.aio.chats.create(model=MODEL_NAME, config=chat_config)
|
||||
request_config = {'candidate_count': 2}
|
||||
request_config_response = await chat.send_message(
|
||||
'tell me a story in 100 words',
|
||||
config=request_config)
|
||||
default_config_response = await chat.send_message(
|
||||
'tell me a story in 100 words')
|
||||
|
||||
assert len(request_config_response.candidates) == 2
|
||||
assert len(default_config_response.candidates) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_history(client):
|
||||
history = [
|
||||
types.Content(
|
||||
role='user', parts=[types.Part.from_text(text='define a=5, b=10')]
|
||||
),
|
||||
types.Content(
|
||||
role='model',
|
||||
parts=[types.Part.from_text(text='Hello there! how can I help you?')],
|
||||
),
|
||||
]
|
||||
chat = client.aio.chats.create(model=MODEL_NAME, history=history)
|
||||
await chat.send_message('what is a + b?')
|
||||
|
||||
assert len(chat.get_history()) > 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_stream_text(client):
|
||||
chat = client.aio.chats.create(model=MODEL_NAME)
|
||||
chunks = 0
|
||||
async for chunk in await chat.send_message_stream('tell me a story in 100 words'):
|
||||
chunks += 1
|
||||
|
||||
assert chunks > 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_stream_part(client):
|
||||
chat = client.aio.chats.create(model=MODEL_NAME)
|
||||
chunks = 0
|
||||
async for chunk in await chat.send_message_stream(
|
||||
types.Part.from_text(text='tell me a story in 100 words')
|
||||
):
|
||||
chunks += 1
|
||||
|
||||
assert chunks > 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_stream_parts(client):
|
||||
chat = client.aio.chats.create(model=MODEL_NAME)
|
||||
chunks = 0
|
||||
async for chunk in await chat.send_message_stream(
|
||||
[
|
||||
types.Part.from_text(text='tell me a story in 100 words'),
|
||||
types.Part.from_text(text='the story is about a car'),
|
||||
],
|
||||
):
|
||||
chunks += 1
|
||||
|
||||
assert chunks > 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_stream_config_override(client):
|
||||
chat_config = {'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_text = ''
|
||||
async for chunk in await chat.send_message_stream(
|
||||
'tell me a story in 100 words', config=request_config
|
||||
):
|
||||
request_config_text += chunk.text
|
||||
default_config_text = ''
|
||||
|
||||
async for chunk in await chat.send_message_stream('tell me family friendly story in 100 words'):
|
||||
default_config_text += chunk.text
|
||||
|
||||
assert json.loads(request_config_text)
|
||||
with pytest_helper.exception_if_mldev(client, json.JSONDecodeError):
|
||||
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)
|
||||
async for chunk in await chat.send_message_stream(
|
||||
'write a python function to check if a year is a leap year'
|
||||
):
|
||||
pass
|
||||
async for chunk in await chat.send_message_stream(
|
||||
'write a unit test for the function'
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
def test_mcp_tools(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'}},
|
||||
},
|
||||
)
|
||||
],},
|
||||
)
|
||||
response = chat.send_message('What is the weather in Boston?')
|
||||
response = chat.send_message('What is the weather in San Francisco?')
|
||||
|
||||
|
||||
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'}},
|
||||
},
|
||||
)
|
||||
],
|
||||
},
|
||||
)
|
||||
for chunk in chat.send_message_stream(
|
||||
'What is the weather in Boston?'
|
||||
):
|
||||
pass
|
||||
for chunk in chat.send_message_stream(
|
||||
'What is the weather in San Francisco?'
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
@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'}},
|
||||
},
|
||||
)
|
||||
],},
|
||||
)
|
||||
await chat.send_message('What is the weather in Boston?');
|
||||
await chat.send_message('What is the weather in San Francisco?');
|
||||
|
||||
|
||||
@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'}},
|
||||
},
|
||||
)
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
async for chunk in await chat.send_message_stream(
|
||||
'What is the weather in Boston?'
|
||||
):
|
||||
pass
|
||||
async for chunk in await chat.send_message_stream(
|
||||
'What is the weather in San Francisco?'
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
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': [
|
||||
{
|
||||
'mcp_servers': [
|
||||
{
|
||||
'name': 'weather_server',
|
||||
'streamable_http_transport': {
|
||||
'url': (
|
||||
'https://gemini-api-demos.uc.r.appspot.com/mcp'
|
||||
),
|
||||
'headers': {
|
||||
'AUTHORIZATION': 'Bearer github_pat_XXXX',
|
||||
},
|
||||
'timeout': '10s',
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
)
|
||||
response = chat.send_message('What is the weather in Boston on 02/02/2026?')
|
||||
response = chat.send_message(
|
||||
'What is the weather in San Francisco on 02/02/2026?'
|
||||
)
|
||||
|
||||
|
||||
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': [
|
||||
{
|
||||
'mcp_servers': [
|
||||
{
|
||||
'name': 'weather_server',
|
||||
'streamable_http_transport': {
|
||||
'url': (
|
||||
'https://gemini-api-demos.uc.r.appspot.com/mcp'
|
||||
),
|
||||
'headers': {
|
||||
'AUTHORIZATION': 'Bearer github_pat_XXXX',
|
||||
},
|
||||
'timeout': '10s',
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
)
|
||||
for chunk in chat.send_message_stream(
|
||||
'What is the weather in Boston on 02/02/2026?'
|
||||
):
|
||||
pass
|
||||
for chunk in chat.send_message_stream(
|
||||
'What is the weather in San Francisco on 02/02/2026?'
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
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': [
|
||||
{
|
||||
'mcp_servers': [
|
||||
{
|
||||
'name': 'weather_server',
|
||||
'streamable_http_transport': {
|
||||
'url': (
|
||||
'https://gemini-api-demos.uc.r.appspot.com/mcp'
|
||||
),
|
||||
'headers': {
|
||||
'AUTHORIZATION': 'Bearer github_pat_XXXX',
|
||||
},
|
||||
'timeout': '10s',
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
)
|
||||
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?'
|
||||
)
|
||||
@@ -0,0 +1,90 @@
|
||||
# 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 ... import types
|
||||
from ...chats import _validate_response
|
||||
|
||||
|
||||
def test_validate_response_default_response():
|
||||
response = types.GenerateContentResponse()
|
||||
|
||||
assert not _validate_response(response)
|
||||
|
||||
|
||||
def test_validate_response_empty_content():
|
||||
response = types.GenerateContentResponse(candidates=[])
|
||||
|
||||
assert not _validate_response(response)
|
||||
|
||||
|
||||
def test_validate_response_empty_parts():
|
||||
response = types.GenerateContentResponse(
|
||||
candidates=[types.Candidate(content=types.Content(parts=[]))]
|
||||
)
|
||||
|
||||
assert not _validate_response(response)
|
||||
|
||||
|
||||
def test_validate_response_empty_part():
|
||||
response = types.GenerateContentResponse(
|
||||
candidates=[types.Candidate(content=types.Content(parts=[types.Part()]))]
|
||||
)
|
||||
|
||||
assert not _validate_response(response)
|
||||
|
||||
|
||||
def test_validate_response_part_with_empty_text():
|
||||
response = types.GenerateContentResponse(
|
||||
candidates=[
|
||||
types.Candidate(content=types.Content(parts=[types.Part(text='')]))
|
||||
]
|
||||
)
|
||||
|
||||
assert _validate_response(response)
|
||||
|
||||
|
||||
def test_validate_response_part_with_text():
|
||||
response = types.GenerateContentResponse(
|
||||
candidates=[
|
||||
types.Candidate(
|
||||
content=types.Content(
|
||||
parts=[types.Part(text='response from model')]
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert _validate_response(response)
|
||||
|
||||
|
||||
def test_validate_response_part_with_function_call():
|
||||
response = types.GenerateContentResponse(
|
||||
candidates=[
|
||||
types.Candidate(
|
||||
content=types.Content(
|
||||
parts=[
|
||||
types.Part(
|
||||
function_call=types.FunctionCall(
|
||||
name='foo', args={'bar': 'baz'}
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert _validate_response(response)
|
||||
@@ -0,0 +1,17 @@
|
||||
# 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."""
|
||||
@@ -0,0 +1,427 @@
|
||||
# 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 async stream."""
|
||||
|
||||
import asyncio
|
||||
from typing import List
|
||||
from unittest import mock
|
||||
from unittest.mock import AsyncMock
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import patch
|
||||
import pytest
|
||||
|
||||
|
||||
try:
|
||||
import aiohttp
|
||||
|
||||
AIOHTTP_NOT_INSTALLED = False
|
||||
except ImportError:
|
||||
AIOHTTP_NOT_INSTALLED = True
|
||||
aiohttp = mock.MagicMock()
|
||||
|
||||
import httpx
|
||||
|
||||
from ... import _api_client as api_client
|
||||
from ... import Client
|
||||
from ... import errors
|
||||
from ... import types
|
||||
|
||||
|
||||
EVENT_STREAM_DATA_WITH_ERROR = [
|
||||
b'{"candidates":[{"content":{"parts":[{"text":"test"}],"role":"model"}}]}',
|
||||
b'\n',
|
||||
b'{"error":{"code":500,"message":"Error","status":"INTERNAL"}}',
|
||||
]
|
||||
|
||||
|
||||
class MockHTTPXResponse(httpx.Response):
|
||||
"""Mock httpx.Response class for testing."""
|
||||
|
||||
def __init__(self, lines: List[str]):
|
||||
self.aiter_lines = MagicMock()
|
||||
self.aiter_lines.return_value.__aiter__ = MagicMock(
|
||||
return_value=self._async_line_iterator(lines)
|
||||
)
|
||||
self.aclose = AsyncMock()
|
||||
|
||||
async def _async_line_iterator(self, lines: List[str]):
|
||||
for line in lines:
|
||||
yield line
|
||||
|
||||
|
||||
class MockAIOHTTPResponse(aiohttp.ClientResponse):
|
||||
|
||||
def __init__(self, lines: List[str]):
|
||||
self.content = MagicMock()
|
||||
self.content.readline = AsyncMock()
|
||||
# Simulate reading lines, each ending with newline bytes for readline behavior
|
||||
self._read_data = b"\n".join(line.encode("utf-8") for line in lines) + b"\n"
|
||||
self._read_pos = 0
|
||||
self.content.readline.side_effect = self._async_read_line
|
||||
self.release = MagicMock()
|
||||
|
||||
async def _async_read_line(self) -> bytes:
|
||||
if self._read_pos >= len(self._read_data):
|
||||
return b"" # End of stream
|
||||
|
||||
newline_pos = self._read_data.find(b"\n", self._read_pos)
|
||||
if newline_pos == -1: # Should not happen with the appended '\n'
|
||||
line = self._read_data[self._read_pos :]
|
||||
self._read_pos = len(self._read_data)
|
||||
return line
|
||||
else:
|
||||
line = self._read_data[self._read_pos : newline_pos + 1]
|
||||
self._read_pos = newline_pos + 1
|
||||
return line
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def responses() -> api_client.HttpResponse:
|
||||
return api_client.HttpResponse(headers={})
|
||||
|
||||
|
||||
requires_aiohttp = pytest.mark.skipif(
|
||||
AIOHTTP_NOT_INSTALLED, reason="aiohttp is not installed, skipping test."
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_has_aiohttp():
|
||||
yield
|
||||
api_client.has_aiohttp = False
|
||||
|
||||
|
||||
def test_invalid_response_stream_type(responses: api_client.HttpResponse):
|
||||
"""Tests that an invalid response stream type raises an error."""
|
||||
api_client.has_aiohttp = False
|
||||
with pytest.raises(
|
||||
TypeError,
|
||||
match=(
|
||||
"Expected self.response_stream to be an httpx.Response or"
|
||||
" aiohttp.ClientResponse object"
|
||||
),
|
||||
):
|
||||
|
||||
async def run():
|
||||
async for _ in responses._aiter_response_stream():
|
||||
pass
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_httpx_simple_lines(responses: api_client.HttpResponse):
|
||||
lines = ["hello", "world", "testing"]
|
||||
mock_response = MockHTTPXResponse(lines)
|
||||
responses.response_stream = mock_response
|
||||
|
||||
results = [line async for line in responses._aiter_response_stream()]
|
||||
|
||||
assert results == lines
|
||||
mock_response.aiter_lines.assert_called_once()
|
||||
mock_response.aclose.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_httpx_data_prefix(responses: api_client.HttpResponse):
|
||||
lines = ["data: { 'message': 'hello' }", "data: { 'status': 'ok' }"]
|
||||
mock_response = MockHTTPXResponse(lines)
|
||||
responses.response_stream = mock_response
|
||||
|
||||
results = [line async for line in responses._aiter_response_stream()]
|
||||
|
||||
assert results == ["{ 'message': 'hello' }", "{ 'status': 'ok' }"]
|
||||
mock_response.aiter_lines.assert_called_once()
|
||||
mock_response.aclose.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_httpx_multiple_json_chunk(responses: api_client.HttpResponse):
|
||||
lines = [
|
||||
'{ "id": 1 }',
|
||||
"",
|
||||
'data: { "id": 2 }',
|
||||
'data: { "id": 3 }',
|
||||
]
|
||||
mock_response = MockHTTPXResponse(lines)
|
||||
responses.response_stream = mock_response
|
||||
|
||||
results = [line async for line in responses._aiter_response_stream()]
|
||||
|
||||
assert results == ['{ "id": 1 }', '{ "id": 2 }', '{ "id": 3 }']
|
||||
mock_response.aiter_lines.assert_called_once()
|
||||
mock_response.aclose.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_httpx_incomplete_json_at_end(responses: api_client.HttpResponse):
|
||||
lines = ['{ "partial": "data"'] # Missing closing brace
|
||||
mock_response = MockHTTPXResponse(lines)
|
||||
responses.response_stream = mock_response
|
||||
|
||||
results = [line async for line in responses._aiter_response_stream()]
|
||||
|
||||
# The remaining chunk is yielded
|
||||
assert results == ['{ "partial": "data"']
|
||||
mock_response.aiter_lines.assert_called_once()
|
||||
mock_response.aclose.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_httpx_empty_stream(responses: api_client.HttpResponse):
|
||||
lines: List[str] = []
|
||||
mock_response = MockHTTPXResponse(lines)
|
||||
responses.response_stream = mock_response
|
||||
|
||||
results = [line async for line in responses._aiter_response_stream()]
|
||||
|
||||
assert results == []
|
||||
mock_response.aiter_lines.assert_called_once()
|
||||
mock_response.aclose.assert_called_once()
|
||||
|
||||
|
||||
# Async aiohttp
|
||||
@requires_aiohttp
|
||||
@pytest.mark.asyncio
|
||||
async def test_aiohttp_simple_lines(responses: api_client.HttpResponse):
|
||||
api_client.has_aiohttp = True # Force aiohttp
|
||||
lines = ["hello", "world", "testing"]
|
||||
# Use the mock class that pretends to be aiohttp.ClientResponse
|
||||
mock_response = MockAIOHTTPResponse(lines)
|
||||
responses.response_stream = mock_response
|
||||
|
||||
results = [line async for line in responses._aiter_response_stream()]
|
||||
|
||||
assert results == lines
|
||||
mock_response.content.readline.assert_any_call()
|
||||
mock_response.release.assert_called_once()
|
||||
|
||||
|
||||
@requires_aiohttp
|
||||
@pytest.mark.asyncio
|
||||
async def test_aiohttp_data_prefix(responses: api_client.HttpResponse):
|
||||
api_client.has_aiohttp = True # Force aiohttp
|
||||
lines = ["data: { 'message': 'hello' }", "data: { 'status': 'ok' }"]
|
||||
# Use the mock class that pretends to be aiohttp.ClientResponse
|
||||
mock_response = MockAIOHTTPResponse(lines)
|
||||
responses.response_stream = mock_response
|
||||
|
||||
results = [line async for line in responses._aiter_response_stream()]
|
||||
|
||||
assert results == ["{ 'message': 'hello' }", "{ 'status': 'ok' }"]
|
||||
mock_response.content.readline.assert_any_call()
|
||||
mock_response.release.assert_called_once()
|
||||
|
||||
|
||||
@requires_aiohttp
|
||||
@pytest.mark.asyncio
|
||||
async def test_aiohttp_multiple_json_chunks(responses: api_client.HttpResponse):
|
||||
api_client.has_aiohttp = True # Force aiohttp
|
||||
lines = [
|
||||
'{ "id": 1 }',
|
||||
"", # empty line to check robustness
|
||||
'data: { "id": 2 }',
|
||||
'data: { "id": 3 }',
|
||||
]
|
||||
# Use the mock class that pretends to be aiohttp.ClientResponse
|
||||
mock_response = MockAIOHTTPResponse(lines)
|
||||
responses.response_stream = mock_response
|
||||
|
||||
results = [line async for line in responses._aiter_response_stream()]
|
||||
|
||||
assert results == ['{ "id": 1 }', '{ "id": 2 }', '{ "id": 3 }']
|
||||
mock_response.content.readline.assert_any_call()
|
||||
mock_response.release.assert_called_once()
|
||||
|
||||
|
||||
@requires_aiohttp
|
||||
@pytest.mark.asyncio
|
||||
async def test_aiohttp_incomplete_json_at_end(
|
||||
responses: api_client.HttpResponse,
|
||||
):
|
||||
api_client.has_aiohttp = True # Force aiohttp
|
||||
lines = ['{ "partial": "data"'] # Missing closing brace
|
||||
# Use the mock class that pretends to be aiohttp.ClientResponse
|
||||
mock_response = MockAIOHTTPResponse(lines)
|
||||
responses.response_stream = mock_response
|
||||
|
||||
results = [line async for line in responses._aiter_response_stream()]
|
||||
|
||||
assert results == ['{ "partial": "data"']
|
||||
mock_response.content.readline.assert_any_call()
|
||||
mock_response.release.assert_called_once()
|
||||
|
||||
|
||||
def mock_response(chunks):
|
||||
mock_stream = MagicMock(spec=httpx.SyncByteStream)
|
||||
mock_stream.__iter__.return_value = chunks
|
||||
return httpx.Response(
|
||||
status_code=200,
|
||||
stream=mock_stream,
|
||||
)
|
||||
|
||||
|
||||
@patch('httpx.Client.send')
|
||||
def test_error_event_in_streamed_responses_bad_json(mock_send_method):
|
||||
with_bad_json = [
|
||||
b'{"candidates":[{"content":{"parts":[{"text":"test"}],"role":"model"}}]}',
|
||||
b'\n',
|
||||
b'{"error": bad_json}',
|
||||
]
|
||||
mock_send_method.return_value = mock_response(with_bad_json)
|
||||
|
||||
client = api_client.BaseApiClient(api_key='test_api_key')
|
||||
stream = client.request_streamed('POST', 'models/gemini-2.5-flash', {})
|
||||
|
||||
chunk = next(stream)
|
||||
assert chunk == types.HttpResponse(
|
||||
headers={},
|
||||
body=(
|
||||
'{"candidates": [{"content": {"parts": [{"text": "test"}], "role":'
|
||||
' "model"}}]}'
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(errors.UnknownApiResponseError):
|
||||
next(stream)
|
||||
|
||||
|
||||
@patch('httpx.Client.send')
|
||||
def test_error_event_in_streamed_responses(mock_send_method):
|
||||
mock_send_method.return_value = mock_response(EVENT_STREAM_DATA_WITH_ERROR)
|
||||
|
||||
client = api_client.BaseApiClient(api_key='test_api_key')
|
||||
stream = client.request_streamed('POST', 'models/gemini-2.5-flash', {})
|
||||
|
||||
chunk = next(stream)
|
||||
assert chunk == types.HttpResponse(
|
||||
body=(
|
||||
'{"candidates": [{"content": {"parts": [{"text": "test"}], "role":'
|
||||
' "model"}}]}'
|
||||
),
|
||||
headers={},
|
||||
)
|
||||
|
||||
with pytest.raises(errors.ServerError):
|
||||
next(stream)
|
||||
|
||||
|
||||
@patch('httpx.Client.send')
|
||||
def test_error_event_in_generate_content_stream(mock_send_method):
|
||||
mock_send_method.return_value = mock_response(EVENT_STREAM_DATA_WITH_ERROR)
|
||||
|
||||
client = Client(api_key='test_api_key')
|
||||
generated_response = client.models.generate_content_stream(
|
||||
model='gemini-2.5-flash',
|
||||
contents='Tell me a story in 300 words.',
|
||||
)
|
||||
|
||||
chunk = next(generated_response)
|
||||
assert chunk == types.GenerateContentResponse(
|
||||
candidates=[
|
||||
types.Candidate(
|
||||
content=types.Content(
|
||||
parts=[
|
||||
types.Part(
|
||||
text='test'
|
||||
),
|
||||
],
|
||||
role='model'
|
||||
)
|
||||
),
|
||||
],
|
||||
sdk_http_response=types.HttpResponse(headers={})
|
||||
)
|
||||
|
||||
with pytest.raises(errors.ServerError):
|
||||
next(generated_response)
|
||||
|
||||
|
||||
async def _async_httpx_response(_):
|
||||
mock_stream = MagicMock(spec=httpx.AsyncByteStream)
|
||||
mock_stream.__aiter__.return_value = EVENT_STREAM_DATA_WITH_ERROR
|
||||
mock_stream.aclose = AsyncMock()
|
||||
return httpx.Response(
|
||||
status_code=200,
|
||||
stream=mock_stream,
|
||||
)
|
||||
|
||||
|
||||
@patch('httpx.AsyncBaseTransport')
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_event_in_streamed_responses_async(mock_transport):
|
||||
client = api_client.BaseApiClient(
|
||||
api_key='test_api_key',
|
||||
http_options=types.HttpOptions(
|
||||
async_client_args={'transport': mock_transport}
|
||||
),
|
||||
)
|
||||
mock_transport.handle_async_request = _async_httpx_response
|
||||
mock_transport.aclose = AsyncMock()
|
||||
|
||||
resp = await client.async_request_streamed(
|
||||
'POST', 'models/gemini-2.5-flash', {'key': 'value'}
|
||||
)
|
||||
|
||||
chunk = await anext(resp)
|
||||
assert chunk == types.HttpResponse(
|
||||
headers={},
|
||||
body=(
|
||||
'{"candidates": [{"content": {"parts": [{"text": "test"}], "role":'
|
||||
' "model"}}]}'
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(errors.ServerError):
|
||||
await anext(resp)
|
||||
|
||||
|
||||
@patch('httpx.AsyncBaseTransport')
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_event_in_generate_content_stream_async(mock_transport):
|
||||
client = Client(
|
||||
api_key='test_api_key',
|
||||
http_options=types.HttpOptions(
|
||||
async_client_args={'transport': mock_transport}
|
||||
),
|
||||
)
|
||||
mock_transport.handle_async_request = _async_httpx_response
|
||||
mock_transport.aclose = AsyncMock()
|
||||
|
||||
generated_response = await client.aio.models.generate_content_stream(
|
||||
model='gemini-2.5-flash',
|
||||
contents='Tell me a story in 300 words.',
|
||||
)
|
||||
|
||||
chunk = await anext(generated_response)
|
||||
assert chunk == types.GenerateContentResponse(
|
||||
candidates=[
|
||||
types.Candidate(
|
||||
content=types.Content(
|
||||
parts=[
|
||||
types.Part(
|
||||
text='test'
|
||||
),
|
||||
],
|
||||
role='model'
|
||||
)
|
||||
),
|
||||
],
|
||||
sdk_http_response=types.HttpResponse(headers={})
|
||||
)
|
||||
|
||||
with pytest.raises(errors.ServerError):
|
||||
await anext(generated_response)
|
||||
@@ -0,0 +1,204 @@
|
||||
# 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 closing the clients and context managers."""
|
||||
import asyncio
|
||||
from unittest import mock
|
||||
|
||||
from google.oauth2 import credentials
|
||||
import pytest
|
||||
try:
|
||||
import aiohttp
|
||||
AIOHTTP_NOT_INSTALLED = False
|
||||
except ImportError:
|
||||
AIOHTTP_NOT_INSTALLED = True
|
||||
aiohttp = mock.MagicMock()
|
||||
|
||||
|
||||
from ... import _api_client as api_client
|
||||
from ... import Client
|
||||
|
||||
|
||||
requires_aiohttp = pytest.mark.skipif(
|
||||
AIOHTTP_NOT_INSTALLED, reason='aiohttp is not installed, skipping test.'
|
||||
)
|
||||
|
||||
|
||||
def test_close_httpx_client():
|
||||
"""Tests that the httpx client is closed when the client is closed."""
|
||||
api_client.has_aiohttp = False
|
||||
client = Client(
|
||||
vertexai=True,
|
||||
project='test_project',
|
||||
location='global',
|
||||
http_options=api_client.HttpOptions(client_args={'max_redirects': 10}),
|
||||
)
|
||||
client.close()
|
||||
assert client._api_client._httpx_client.is_closed
|
||||
|
||||
|
||||
def test_httpx_client_context_manager():
|
||||
"""Tests that the httpx client is closed when the client is closed."""
|
||||
api_client.has_aiohttp = False
|
||||
with Client(
|
||||
vertexai=True,
|
||||
project='test_project',
|
||||
location='global',
|
||||
http_options=api_client.HttpOptions(client_args={'max_redirects': 10}),
|
||||
) as client:
|
||||
pass
|
||||
assert not client._api_client._httpx_client.is_closed
|
||||
|
||||
assert client._api_client._httpx_client.is_closed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aclose_httpx_client():
|
||||
"""Tests that the httpx async client is closed when the client is closed."""
|
||||
api_client.has_aiohttp = False
|
||||
async_client = Client(
|
||||
vertexai=True,
|
||||
project='test_project',
|
||||
location='global',
|
||||
).aio
|
||||
await async_client.aclose()
|
||||
assert async_client._api_client._async_httpx_client.is_closed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_httpx_client_context_manager():
|
||||
"""Tests that the httpx async client is closed when the client is closed."""
|
||||
api_client.has_aiohttp = False
|
||||
async with Client(
|
||||
vertexai=True,
|
||||
project='test_project',
|
||||
location='global',
|
||||
).aio as async_client:
|
||||
pass
|
||||
assert not async_client._api_client._async_httpx_client.is_closed
|
||||
|
||||
assert async_client._api_client._async_httpx_client.is_closed
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_request():
|
||||
mock_aiohttp_response = mock.Mock(spec=aiohttp.ClientSession.request)
|
||||
mock_aiohttp_response.return_value = mock_aiohttp_response
|
||||
yield mock_aiohttp_response
|
||||
|
||||
|
||||
def _patch_auth_default():
|
||||
return mock.patch(
|
||||
'google.auth.default',
|
||||
return_value=(credentials.Credentials('magic_token'), 'test_project'),
|
||||
autospec=True,
|
||||
)
|
||||
|
||||
|
||||
async def _aiohttp_async_response(status: int):
|
||||
"""Has to return a coroutine hence async."""
|
||||
response = mock.Mock(spec=aiohttp.ClientResponse)
|
||||
response.status = status
|
||||
response.headers = {'status-code': str(status)}
|
||||
response.json.return_value = {}
|
||||
response.text.return_value = 'test'
|
||||
return response
|
||||
|
||||
|
||||
@requires_aiohttp
|
||||
@mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
|
||||
def test_aclose_aiohttp_session(mock_request):
|
||||
"""Tests that the aiohttp session is closed when the client is closed."""
|
||||
api_client.has_aiohttp = True
|
||||
async def run():
|
||||
mock_request.side_effect = (
|
||||
aiohttp.ClientConnectorError(
|
||||
connection_key=aiohttp.client_reqrep.ConnectionKey(
|
||||
'localhost', 80, False, True, None, None, None
|
||||
),
|
||||
os_error=OSError,
|
||||
),
|
||||
_aiohttp_async_response(200),
|
||||
)
|
||||
with _patch_auth_default():
|
||||
async_client = Client(
|
||||
vertexai=True,
|
||||
project='test_project',
|
||||
location='global',
|
||||
http_options=api_client.HttpOptions(
|
||||
async_client_args={'trust_env': False}
|
||||
),
|
||||
).aio
|
||||
# aiohttp session is created in the first request instead of client
|
||||
# initialization.
|
||||
_ = await async_client._api_client._async_request_once(
|
||||
api_client.HttpRequest(
|
||||
method='GET',
|
||||
url='https://example.com',
|
||||
headers={},
|
||||
data=None,
|
||||
timeout=None,
|
||||
)
|
||||
)
|
||||
assert async_client._api_client._aiohttp_session is not None
|
||||
assert not async_client._api_client._aiohttp_session.closed
|
||||
# Close the client and check that the session is closed.
|
||||
await async_client.aclose()
|
||||
assert async_client._api_client._aiohttp_session.closed
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@requires_aiohttp
|
||||
@mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
|
||||
def test_aiohttp_session_context_manager(mock_request):
|
||||
"""Tests that the aiohttp session is closed when the client is closed."""
|
||||
api_client.has_aiohttp = True
|
||||
async def run():
|
||||
mock_request.side_effect = (
|
||||
aiohttp.ClientConnectorError(
|
||||
connection_key=aiohttp.client_reqrep.ConnectionKey(
|
||||
'localhost', 80, False, True, None, None, None
|
||||
),
|
||||
os_error=OSError,
|
||||
),
|
||||
_aiohttp_async_response(200),
|
||||
)
|
||||
with _patch_auth_default():
|
||||
async with Client(
|
||||
vertexai=True,
|
||||
project='test_project',
|
||||
location='global',
|
||||
http_options=api_client.HttpOptions(
|
||||
async_client_args={'trust_env': False}
|
||||
),
|
||||
).aio as async_client:
|
||||
# aiohttp session is created in the first request instead of client
|
||||
# initialization.
|
||||
_ = await async_client._api_client._async_request_once(
|
||||
api_client.HttpRequest(
|
||||
method='GET',
|
||||
url='https://example.com',
|
||||
headers={},
|
||||
data=None,
|
||||
timeout=None,
|
||||
)
|
||||
)
|
||||
assert async_client._api_client._aiohttp_session is not None
|
||||
assert not async_client._api_client._aiohttp_session.closed
|
||||
|
||||
assert async_client._api_client._aiohttp_session.closed
|
||||
|
||||
asyncio.run(run())
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,221 @@
|
||||
# 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 client behavior when issuing requests."""
|
||||
|
||||
from ... import _api_client as api_client
|
||||
from ... import Client
|
||||
from ... import types
|
||||
|
||||
|
||||
def build_test_client(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'google_api_key')
|
||||
return Client()
|
||||
|
||||
|
||||
def test_join_url_path_base_url_with_trailing_slash_and_path_with_leading_slash():
|
||||
base_url = 'https://fake-url.com/some_path/'
|
||||
path = '/v1beta/models'
|
||||
assert (
|
||||
api_client.join_url_path(base_url, path)
|
||||
== 'https://fake-url.com/some_path/v1beta/models'
|
||||
)
|
||||
|
||||
|
||||
def test_join_url_path_with_base_url_with_trailing_slash_and_path_without_leading_slash():
|
||||
base_url = 'https://fake-url.com/some_path/'
|
||||
path = 'v1beta/models'
|
||||
assert (
|
||||
api_client.join_url_path(base_url, path)
|
||||
== 'https://fake-url.com/some_path/v1beta/models'
|
||||
)
|
||||
|
||||
|
||||
def test_join_url_path_with_base_url_without_trailing_slash_and_path_with_leading_slash():
|
||||
base_url = 'https://fake-url.com/some_path'
|
||||
path = '/v1beta/models'
|
||||
assert (
|
||||
api_client.join_url_path(base_url, path)
|
||||
== 'https://fake-url.com/some_path/v1beta/models'
|
||||
)
|
||||
|
||||
|
||||
def test_join_url_path_with_base_url_without_trailing_slash_and_path_without_leading_slash():
|
||||
base_url = 'https://fake-url.com/some_path'
|
||||
path = 'v1beta/models'
|
||||
assert (
|
||||
api_client.join_url_path(base_url, path)
|
||||
== 'https://fake-url.com/some_path/v1beta/models'
|
||||
)
|
||||
|
||||
|
||||
def test_join_url_path_base_url_without_path_with_trailing_slash():
|
||||
base_url = 'https://fake-url.com/'
|
||||
path = 'v1beta/models'
|
||||
assert (
|
||||
api_client.join_url_path(base_url, path)
|
||||
== 'https://fake-url.com/v1beta/models'
|
||||
)
|
||||
|
||||
|
||||
def test_join_url_path_base_url_without_path_without_trailing_slash():
|
||||
base_url = 'https://fake-url.com'
|
||||
path = 'v1beta/models'
|
||||
assert (
|
||||
api_client.join_url_path(base_url, path)
|
||||
== 'https://fake-url.com/v1beta/models'
|
||||
)
|
||||
|
||||
|
||||
def test_build_request_sets_library_version_headers(monkeypatch):
|
||||
request_client = build_test_client(monkeypatch).models._api_client
|
||||
request = request_client._build_request('GET', 'test/path', {'key': 'value'})
|
||||
assert 'google-genai-sdk/' in request.headers['user-agent']
|
||||
assert 'gl-python/' in request.headers['user-agent']
|
||||
assert 'google-genai-sdk/' in request.headers['x-goog-api-client']
|
||||
assert 'gl-python/' in request.headers['x-goog-api-client']
|
||||
|
||||
|
||||
def test_build_request_appends_to_user_agent_headers(monkeypatch):
|
||||
request_client = build_test_client(monkeypatch).models._api_client
|
||||
request = request_client._build_request(
|
||||
'GET',
|
||||
'test/path',
|
||||
{'key': 'value'},
|
||||
types.HttpOptionsDict(
|
||||
base_url='test/url',
|
||||
api_version='1',
|
||||
headers={'user-agent': 'test-user-agent'},
|
||||
),
|
||||
)
|
||||
assert 'test-user-agent' in request.headers['user-agent']
|
||||
assert 'google-genai-sdk/' in request.headers['user-agent']
|
||||
assert 'gl-python/' in request.headers['user-agent']
|
||||
assert 'google-genai-sdk/' in request.headers['x-goog-api-client']
|
||||
|
||||
|
||||
def test_build_request_appends_to_goog_api_client_headers(monkeypatch):
|
||||
request_client = build_test_client(monkeypatch).models._api_client
|
||||
request = request_client._build_request(
|
||||
'GET',
|
||||
'test/path',
|
||||
{'key': 'value'},
|
||||
types.HttpOptionsDict(
|
||||
base_url='test/url',
|
||||
api_version='1',
|
||||
headers={'x-goog-api-client': 'test-goog-api-client'},
|
||||
),
|
||||
)
|
||||
assert 'google-genai-sdk/' in request.headers['user-agent']
|
||||
assert 'test-goog-api-client' in request.headers['x-goog-api-client']
|
||||
assert 'google-genai-sdk/' in request.headers['x-goog-api-client']
|
||||
assert 'gl-python/' in request.headers['x-goog-api-client']
|
||||
|
||||
|
||||
def test_build_request_keeps_sdk_version_headers(monkeypatch):
|
||||
headers_to_inject = {}
|
||||
api_client.append_library_version_headers(headers_to_inject)
|
||||
assert 'google-genai-sdk/' in headers_to_inject['user-agent']
|
||||
request_client = build_test_client(monkeypatch).models._api_client
|
||||
request = request_client._build_request(
|
||||
'GET',
|
||||
'test/path',
|
||||
{'key': 'value'},
|
||||
types.HttpOptionsDict(
|
||||
base_url='test/url',
|
||||
api_version='1',
|
||||
headers=headers_to_inject,
|
||||
),
|
||||
)
|
||||
assert 'google-genai-sdk/' in request.headers['user-agent']
|
||||
assert 'gl-python/' in request.headers['x-goog-api-client']
|
||||
assert 'google-genai-sdk/' in request.headers['x-goog-api-client']
|
||||
assert 'gl-python/' in request.headers['x-goog-api-client']
|
||||
|
||||
|
||||
def test_build_request_with_resource_scope(monkeypatch):
|
||||
monkeypatch.delenv('GOOGLE_API_KEY', raising=False)
|
||||
monkeypatch.delenv('GEMINI_API_KEY', raising=False)
|
||||
monkeypatch.delenv('GOOGLE_CLOUD_PROJECT', raising=False)
|
||||
monkeypatch.delenv('GOOGLE_CLOUD_LOCATION', raising=False)
|
||||
|
||||
client = Client(
|
||||
vertexai=True,
|
||||
http_options=types.HttpOptionsDict(
|
||||
base_url='https://custom-base-url.com',
|
||||
base_url_resource_scope=types.ResourceScope.COLLECTION,
|
||||
),
|
||||
)
|
||||
|
||||
request = client.models._api_client._build_request(
|
||||
'post',
|
||||
'publishers/google/models/gemini-3-pro-preview',
|
||||
{'key': 'value'},
|
||||
)
|
||||
assert request.url == 'https://custom-base-url.com/publishers/google/models/gemini-3-pro-preview'
|
||||
|
||||
|
||||
def test_build_request_with_resource_scope_with_project_and_location(
|
||||
monkeypatch,
|
||||
):
|
||||
monkeypatch.delenv('GOOGLE_API_KEY', raising=False)
|
||||
monkeypatch.delenv('GEMINI_API_KEY', raising=False)
|
||||
monkeypatch.delenv('GOOGLE_CLOUD_PROJECT', raising=False)
|
||||
monkeypatch.delenv('GOOGLE_CLOUD_LOCATION', raising=False)
|
||||
|
||||
client = Client(
|
||||
vertexai=True,
|
||||
project='test-project',
|
||||
location='test-location',
|
||||
http_options=types.HttpOptionsDict(
|
||||
base_url='https://custom-base-url.com',
|
||||
base_url_resource_scope=types.ResourceScope.COLLECTION,
|
||||
),
|
||||
)
|
||||
|
||||
request = client.models._api_client._build_request(
|
||||
'post',
|
||||
'publishers/google/models/gemini-3-pro-preview',
|
||||
{'key': 'value'},
|
||||
)
|
||||
assert request.url == 'https://custom-base-url.com/publishers/google/models/gemini-3-pro-preview'
|
||||
|
||||
|
||||
|
||||
def build_test_client_no_env_vars(monkeypatch):
|
||||
monkeypatch.delenv('GOOGLE_API_KEY', raising=False)
|
||||
monkeypatch.delenv('GEMINI_API_KEY', raising=False)
|
||||
monkeypatch.delenv('GOOGLE_CLOUD_PROJECT', raising=False)
|
||||
monkeypatch.delenv('GOOGLE_CLOUD_LOCATION', raising=False)
|
||||
return Client(
|
||||
vertexai=True,
|
||||
http_options=types.HttpOptionsDict(
|
||||
base_url='https://custom-base-url.com',
|
||||
headers={'Authorization': 'Bearer fake_access_token'},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_build_request_with_custom_base_url_no_env_vars(monkeypatch):
|
||||
request_client = (
|
||||
build_test_client_no_env_vars(monkeypatch).models._api_client
|
||||
)
|
||||
request = request_client._build_request(
|
||||
'GET',
|
||||
'test/path',
|
||||
{'key': 'value'},
|
||||
)
|
||||
assert request.url == 'https://custom-base-url.com'
|
||||
@@ -0,0 +1,104 @@
|
||||
# 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 custom clients."""
|
||||
import asyncio
|
||||
from unittest import mock
|
||||
|
||||
from google.oauth2 import credentials
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from ... import _api_client as api_client
|
||||
from ... import Client
|
||||
|
||||
|
||||
try:
|
||||
import aiohttp
|
||||
|
||||
AIOHTTP_NOT_INSTALLED = False
|
||||
except ImportError:
|
||||
AIOHTTP_NOT_INSTALLED = True
|
||||
aiohttp = mock.MagicMock()
|
||||
|
||||
requires_aiohttp = pytest.mark.skipif(
|
||||
AIOHTTP_NOT_INSTALLED, reason='aiohttp is not installed, skipping test.'
|
||||
)
|
||||
|
||||
|
||||
# Httpx
|
||||
def test_constructor_with_httpx_clients():
|
||||
mldev_http_options = {
|
||||
'httpx_client': httpx.Client(trust_env=False),
|
||||
'httpx_async_client': httpx.AsyncClient(trust_env=False),
|
||||
}
|
||||
vertexai_http_options = {
|
||||
'httpx_client': httpx.Client(trust_env=False),
|
||||
'httpx_async_client': httpx.AsyncClient(trust_env=False),
|
||||
}
|
||||
|
||||
# Even if aiohttp is installed, expect it to be disabled when httpx clients
|
||||
# are provided.
|
||||
api_client.has_aiohttp = True
|
||||
|
||||
mldev_client = Client(
|
||||
api_key='google_api_key', http_options=mldev_http_options
|
||||
)
|
||||
assert not mldev_client.models._api_client._httpx_client.trust_env
|
||||
assert not mldev_client.models._api_client._async_httpx_client.trust_env
|
||||
# Expect aiohttp to be disabled when httpx clients are provided, regardless of
|
||||
# whether aiohttp is installed.
|
||||
assert not mldev_client.models._api_client._use_aiohttp()
|
||||
|
||||
vertexai_client = Client(
|
||||
vertexai=True,
|
||||
project='fake_project_id',
|
||||
location='fake-location',
|
||||
http_options=vertexai_http_options,
|
||||
)
|
||||
assert not vertexai_client.models._api_client._httpx_client.trust_env
|
||||
assert not vertexai_client.models._api_client._async_httpx_client.trust_env
|
||||
# Expect aiohttp to be disabled when httpx clients are provided, regardless of
|
||||
# whether aiohttp is installed.
|
||||
assert not mldev_client.models._api_client._use_aiohttp()
|
||||
|
||||
|
||||
# Aiohttp
|
||||
@requires_aiohttp
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skipif(
|
||||
AIOHTTP_NOT_INSTALLED, reason='aiohttp is not installed, skipping test.'
|
||||
)
|
||||
async def test_constructor_with_aiohttp_clients():
|
||||
api_client.has_aiohttp = True
|
||||
mldev_http_options = {
|
||||
'aiohttp_client': aiohttp.ClientSession(trust_env=False),
|
||||
}
|
||||
vertexai_http_options = {
|
||||
'aiohttp_client': aiohttp.ClientSession(trust_env=False),
|
||||
}
|
||||
mldev_client = Client(
|
||||
api_key='google_api_key', http_options=mldev_http_options
|
||||
)
|
||||
assert not mldev_client.models._api_client._aiohttp_session.trust_env
|
||||
|
||||
vertexai_client = Client(
|
||||
vertexai=True,
|
||||
project='fake_project_id',
|
||||
location='fake-location',
|
||||
http_options=vertexai_http_options,
|
||||
)
|
||||
assert not vertexai_client.models._api_client._aiohttp_session.trust_env
|
||||
|
||||
@@ -0,0 +1,182 @@
|
||||
# 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 client behavior when issuing requests."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import _api_client
|
||||
from ... import types
|
||||
|
||||
|
||||
def test_patch_http_options_with_copies_all_fields():
|
||||
patch_options = types.HttpOptions(
|
||||
base_url='https://fake-url.com/',
|
||||
api_version='v1',
|
||||
headers={'X-Custom-Header': 'custom_value'},
|
||||
timeout=10000,
|
||||
client_args={'http2': True},
|
||||
async_client_args={'http1': True},
|
||||
extra_body={'key': 'value'},
|
||||
retry_options=types.HttpRetryOptions(attempts=10),
|
||||
base_url_resource_scope=types.ResourceScope.COLLECTION,
|
||||
)
|
||||
options = types.HttpOptions()
|
||||
patched = _api_client.patch_http_options(options, patch_options)
|
||||
http_options_keys = types.HttpOptions.model_fields.keys()
|
||||
|
||||
for key in http_options_keys:
|
||||
assert hasattr(patched, key)
|
||||
if key not in [
|
||||
'httpx_client',
|
||||
'httpx_async_client',
|
||||
'aiohttp_client',
|
||||
]:
|
||||
assert getattr(patched, key) is not None
|
||||
assert patched.base_url == 'https://fake-url.com/'
|
||||
assert patched.api_version == 'v1'
|
||||
assert patched.headers['X-Custom-Header'] == 'custom_value'
|
||||
assert patched.timeout == 10000
|
||||
assert patched.retry_options.attempts == 10
|
||||
assert patched.client_args['http2']
|
||||
assert patched.async_client_args['http1']
|
||||
|
||||
|
||||
def test_patch_http_options_merges_headers():
|
||||
original_options = types.HttpOptions(
|
||||
headers={
|
||||
'X-Custom-Header': 'different_value',
|
||||
'X-different-header': 'different_value',
|
||||
}
|
||||
)
|
||||
patch_options = types.HttpOptions(
|
||||
base_url='https://fake-url.com/',
|
||||
api_version='v1',
|
||||
headers={'X-Custom-Header': 'custom_value'},
|
||||
timeout=10000,
|
||||
)
|
||||
patched = _api_client.patch_http_options(original_options, patch_options)
|
||||
# If the header is present in both the original and patch options, the patch
|
||||
# options value should be used
|
||||
assert patched.headers['X-Custom-Header'] == 'custom_value'
|
||||
assert patched.headers['X-different-header'] == 'different_value'
|
||||
assert patched.base_url == 'https://fake-url.com/'
|
||||
assert patched.api_version == 'v1'
|
||||
assert patched.timeout == 10000
|
||||
|
||||
|
||||
def test_patch_http_options_appends_version_headers():
|
||||
original_options = types.HttpOptions(
|
||||
headers={
|
||||
'X-Custom-Header': 'different_value',
|
||||
'X-different-header': 'different_value',
|
||||
}
|
||||
)
|
||||
patch_options = types.HttpOptions(
|
||||
base_url='https://fake-url.com/',
|
||||
api_version='v1',
|
||||
headers={'X-Custom-Header': 'custom_value'},
|
||||
timeout=10000,
|
||||
)
|
||||
patched = _api_client.patch_http_options(original_options, patch_options)
|
||||
assert 'user-agent' in patched.headers
|
||||
assert 'x-goog-api-client' in patched.headers
|
||||
|
||||
|
||||
def test_setting_timeout_populates_server_timeout_header():
|
||||
api_client = _api_client.BaseApiClient(
|
||||
vertexai=False,
|
||||
api_key='test_api_key',
|
||||
http_options=types.HttpOptions(timeout=10000),
|
||||
)
|
||||
request = api_client._build_request(
|
||||
http_method='POST',
|
||||
path='sample/path',
|
||||
request_dict={},
|
||||
)
|
||||
assert 'X-Server-Timeout' in request.headers
|
||||
assert request.headers['X-Server-Timeout'] == '10'
|
||||
|
||||
|
||||
def test_timeout_rounded_to_nearest_second():
|
||||
api_client = _api_client.BaseApiClient(
|
||||
vertexai=False,
|
||||
api_key='test_api_key',
|
||||
)
|
||||
http_options = types.HttpOptions(timeout=7300)
|
||||
request = api_client._build_request(
|
||||
http_method='POST',
|
||||
path='sample/path',
|
||||
request_dict={},
|
||||
http_options=http_options,
|
||||
)
|
||||
assert request.headers['X-Server-Timeout'] == '8'
|
||||
|
||||
|
||||
def test_server_timeout_not_overwritten():
|
||||
api_client = _api_client.BaseApiClient(
|
||||
vertexai=False,
|
||||
api_key='test_api_key',
|
||||
)
|
||||
http_options = types.HttpOptions(
|
||||
headers={'X-Server-Timeout': '3'},
|
||||
timeout=11000)
|
||||
request = api_client._build_request(
|
||||
http_method='POST',
|
||||
path='sample/path',
|
||||
request_dict={},
|
||||
http_options=http_options,
|
||||
)
|
||||
assert request.headers['X-Server-Timeout'] == '3'
|
||||
|
||||
|
||||
def test_server_timeout_not_set_by_default():
|
||||
api_client = _api_client.BaseApiClient(
|
||||
vertexai=False,
|
||||
api_key='test_api_key',
|
||||
)
|
||||
request = api_client._build_request(
|
||||
http_method='POST',
|
||||
path='sample/path',
|
||||
request_dict={},
|
||||
)
|
||||
assert not 'X-Server-Timeout' in request.headers
|
||||
|
||||
|
||||
def test_resource_scope_without_base_url_raises_error():
|
||||
with pytest.raises(ValueError):
|
||||
_api_client.BaseApiClient(
|
||||
vertexai=True,
|
||||
http_options=types.HttpOptions(
|
||||
base_url_resource_scope=types.ResourceScope.COLLECTION,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_base_url_resource_scope_not_set_by_default():
|
||||
api_client = _api_client.BaseApiClient(
|
||||
vertexai=True,
|
||||
http_options=types.HttpOptions(
|
||||
base_url='https://fake-url.com/',
|
||||
),
|
||||
)
|
||||
|
||||
assert api_client._http_options.base_url_resource_scope is None
|
||||
|
||||
|
||||
def test_retry_options_not_set_by_default():
|
||||
options = types.HttpOptions()
|
||||
assert options.retry_options is None
|
||||
@@ -0,0 +1,168 @@
|
||||
# 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 replay client correctly compares requests ignoring key casing."""
|
||||
|
||||
from ... import _replay_api_client
|
||||
from ... import types
|
||||
|
||||
|
||||
def test_equal_objects_with_same_casing_returns_true():
|
||||
obj1 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generationConfig': {'topK': 2.0, 'maxOutputTokens': 3},
|
||||
}]
|
||||
obj2 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generationConfig': {'topK': 2.0, 'maxOutputTokens': 3},
|
||||
}]
|
||||
|
||||
assert _replay_api_client._equals_ignore_key_case(obj1, obj2)
|
||||
|
||||
|
||||
def test_equal_objects_with_different_casing_returns_true():
|
||||
obj1 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generationConfig': {'topK': 2.0, 'maxOutputTokens': 3},
|
||||
}]
|
||||
obj2 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generation_config': {'topK': 2.0, 'max_output_tokens': 3},
|
||||
}]
|
||||
|
||||
assert _replay_api_client._equals_ignore_key_case(obj1, obj2)
|
||||
|
||||
|
||||
def test_equal_objects_with_different_nested_casing_returns_true():
|
||||
obj1 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generationConfig': {'topK': 2.0, 'maxOutputTokens': 3},
|
||||
}]
|
||||
obj2 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generationConfig': {'top_k': 2.0, 'maxOutputTokens': 3},
|
||||
}]
|
||||
|
||||
assert _replay_api_client._equals_ignore_key_case(obj1, obj2)
|
||||
|
||||
def test_equal_int_float_values_returns_true():
|
||||
obj1 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generationConfig': {'topK': 2.0, 'maxOutputTokens': 3},
|
||||
}]
|
||||
obj2 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generation_config': {'topK': 2, 'max_output_tokens': 3},
|
||||
}]
|
||||
|
||||
assert _replay_api_client._equals_ignore_key_case(obj1, obj2)
|
||||
|
||||
|
||||
def test_equal_enum_string_values_returns_true():
|
||||
obj1 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generationConfig': {
|
||||
'topK': 2.0,
|
||||
'maxOutputTokens': 3,
|
||||
'type': types.Type.STRING,
|
||||
},
|
||||
}]
|
||||
obj2 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generation_config': {
|
||||
'topK': 2,
|
||||
'max_output_tokens': 3,
|
||||
'type': 'STRING',
|
||||
},
|
||||
}]
|
||||
|
||||
assert _replay_api_client._equals_ignore_key_case(obj1, obj2)
|
||||
|
||||
|
||||
def test_equal_enum_values_returns_true():
|
||||
obj1 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generationConfig': {
|
||||
'topK': 2.0,
|
||||
'maxOutputTokens': 3,
|
||||
'type': types.Type.STRING,
|
||||
},
|
||||
}]
|
||||
obj2 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generation_config': {
|
||||
'topK': 2,
|
||||
'max_output_tokens': 3,
|
||||
'type': types.Type.STRING,
|
||||
},
|
||||
}]
|
||||
|
||||
assert _replay_api_client._equals_ignore_key_case(obj1, obj2)
|
||||
|
||||
|
||||
def test_different_number_value_returns_false():
|
||||
obj1 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generationConfig': {'topK': 3, 'maxOutputTokens': 3},
|
||||
}]
|
||||
obj2 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generation_config': {'topK': 2.0, 'max_output_tokens': 3},
|
||||
}]
|
||||
|
||||
assert not _replay_api_client._equals_ignore_key_case(
|
||||
obj1, obj2
|
||||
)
|
||||
|
||||
|
||||
def test_different_string_value_returns_false():
|
||||
obj1 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generationConfig': {'topK': 3, 'maxOutputTokens': 3},
|
||||
}]
|
||||
obj2 = [{
|
||||
'contents': [
|
||||
{'parts': [{'text': 'What is your name?'}], 'role': 'model'}
|
||||
],
|
||||
'generation_config': {'topK': 3, 'max_output_tokens': 3},
|
||||
}]
|
||||
|
||||
assert not _replay_api_client._equals_ignore_key_case(
|
||||
obj1, obj2
|
||||
)
|
||||
|
||||
|
||||
def test_different_enum_values_returns_false():
|
||||
obj1 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generationConfig': {
|
||||
'topK': 2.0,
|
||||
'maxOutputTokens': 3,
|
||||
'type': types.Type.STRING,
|
||||
},
|
||||
}]
|
||||
obj2 = [{
|
||||
'contents': [{'parts': [{'text': 'What is your name?'}], 'role': 'user'}],
|
||||
'generation_config': {
|
||||
'topK': 2,
|
||||
'max_output_tokens': 3,
|
||||
'type': types.Type.OBJECT,
|
||||
},
|
||||
}]
|
||||
|
||||
assert not _replay_api_client._equals_ignore_key_case(
|
||||
obj1, obj2
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,197 @@
|
||||
# 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 error handling in the client api client layer for uploads."""
|
||||
|
||||
import io
|
||||
import json
|
||||
from unittest import mock
|
||||
import pytest
|
||||
|
||||
import httpx
|
||||
|
||||
from ... import _api_client as api_client
|
||||
from ... import errors
|
||||
from ... import types
|
||||
|
||||
try:
|
||||
import aiohttp
|
||||
AIOHTTP_NOT_INSTALLED = False
|
||||
except ImportError:
|
||||
AIOHTTP_NOT_INSTALLED = True
|
||||
|
||||
def _httpx_response(code: int, headers: dict = None, content: bytes = b''):
|
||||
headers = headers or {}
|
||||
headers['status-code'] = str(code)
|
||||
return httpx.Response(
|
||||
status_code=code,
|
||||
headers=headers,
|
||||
content=content,
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
def client():
|
||||
return api_client.BaseApiClient(vertexai=False, api_key='test_api_key')
|
||||
|
||||
|
||||
def test_upload_url_rewrite(client: api_client.BaseApiClient):
|
||||
mock_httpx_client = mock.MagicMock(spec=httpx.Client)
|
||||
mock_httpx_client.request.side_effect = [
|
||||
_httpx_response(
|
||||
200, headers={"X-Goog-Upload-Status": "final"}
|
||||
), # Upload request succeeding
|
||||
]
|
||||
client._httpx_client = mock_httpx_client
|
||||
|
||||
http_options = types.HttpOptions(base_url="https://my-proxy.company.com")
|
||||
|
||||
with io.BytesIO(b"test") as f:
|
||||
client._upload_fd(
|
||||
f,
|
||||
"https://generativelanguage.googleapis.com/upload/v1beta/files?uploadType=resumable",
|
||||
4,
|
||||
http_options=http_options,
|
||||
)
|
||||
|
||||
assert mock_httpx_client.request.call_count == 1
|
||||
call_args = mock_httpx_client.request.call_args[1]
|
||||
assert (
|
||||
call_args["url"]
|
||||
== "https://my-proxy.company.com/upload/v1beta/files?uploadType=resumable"
|
||||
)
|
||||
|
||||
|
||||
def test_upload_fd_error(client: api_client.BaseApiClient):
|
||||
error_content = json.dumps({
|
||||
"error": {
|
||||
"code": 400,
|
||||
"message": "Unsupported MIME type: bad/mime_type",
|
||||
"status": "INVALID_ARGUMENT"
|
||||
}
|
||||
}).encode('utf-8')
|
||||
|
||||
mock_httpx_client = mock.MagicMock(spec=httpx.Client)
|
||||
mock_httpx_client.request.side_effect = [
|
||||
_httpx_response(200, headers={"X-Goog-Upload-URL": "http://fake/upload"}), # Initial request to get upload URL
|
||||
_httpx_response(400, content=error_content, headers={"X-Goog-Upload-Status": "final"}), # Upload request failing
|
||||
]
|
||||
client._httpx_client = mock_httpx_client
|
||||
|
||||
with pytest.raises(errors.APIError, match='Unsupported MIME type: bad/mime_type'), io.BytesIO(b"test") as f:
|
||||
client._upload_fd(
|
||||
f, # Pass file object directly
|
||||
"http://fake/upload",
|
||||
4,
|
||||
http_options=types.HttpOptions(),
|
||||
)
|
||||
|
||||
assert mock_httpx_client.request.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_upload_url_rewrite_httpx(client: api_client.BaseApiClient):
|
||||
mock_async_httpx_client = mock.MagicMock(spec=httpx.AsyncClient)
|
||||
mock_async_httpx_client.request = mock.AsyncMock(
|
||||
side_effect=[
|
||||
_httpx_response(
|
||||
200, headers={"X-Goog-Upload-Status": "final"}
|
||||
), # Upload request
|
||||
]
|
||||
)
|
||||
client._async_httpx_client = mock_async_httpx_client
|
||||
|
||||
http_options = types.HttpOptions(base_url="https://my-proxy.company.com")
|
||||
|
||||
with mock.patch.object(
|
||||
client, "_use_aiohttp", return_value=False
|
||||
), io.BytesIO(b"test") as f:
|
||||
await client._async_upload_fd(
|
||||
f,
|
||||
"https://generativelanguage.googleapis.com/upload/v1beta/files?uploadType=resumable",
|
||||
4,
|
||||
http_options=http_options,
|
||||
)
|
||||
|
||||
assert mock_async_httpx_client.request.call_count == 1
|
||||
call_args = mock_async_httpx_client.request.call_args[1]
|
||||
assert (
|
||||
call_args["url"]
|
||||
== "https://my-proxy.company.com/upload/v1beta/files?uploadType=resumable"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_upload_fd_error_httpx(client: api_client.BaseApiClient):
|
||||
error_content = json.dumps({
|
||||
"error": {
|
||||
"code": 400,
|
||||
"message": "Unsupported MIME type: bad/mime_type",
|
||||
"status": "INVALID_ARGUMENT"
|
||||
}
|
||||
}).encode('utf-8')
|
||||
|
||||
# Mock for httpx.AsyncClient
|
||||
mock_async_httpx_client = mock.MagicMock(spec=httpx.AsyncClient)
|
||||
mock_async_httpx_client.request = mock.AsyncMock(side_effect=[
|
||||
_httpx_response(200, headers={"X-Goog-Upload-URL": "http://fake/upload"}), # Initial request
|
||||
_httpx_response(400, content=error_content, headers={"X-Goog-Upload-Status": "final"}), # Upload request
|
||||
])
|
||||
client._async_httpx_client = mock_async_httpx_client
|
||||
|
||||
# Patch the _use_aiohttp method to control which client is used
|
||||
with mock.patch.object(client, '_use_aiohttp', return_value=False), \
|
||||
pytest.raises(errors.APIError, match='Unsupported MIME type: bad/mime_type'), \
|
||||
io.BytesIO(b"test") as f:
|
||||
await client._async_upload_fd(
|
||||
f, # Pass file object directly
|
||||
"http://fake/upload",
|
||||
4,
|
||||
http_options=types.HttpOptions(),
|
||||
|
||||
)
|
||||
assert mock_async_httpx_client.request.call_count == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skipif(
|
||||
AIOHTTP_NOT_INSTALLED, reason="aiohttp is not installed, skipping test."
|
||||
)
|
||||
async def test_async_upload_fd_error_aiohttp(client: api_client.BaseApiClient):
|
||||
error_content = json.dumps({
|
||||
"error": {
|
||||
"code": 400,
|
||||
"message": "Unsupported MIME type: bad/mime_type",
|
||||
"status": "INVALID_ARGUMENT"
|
||||
}
|
||||
}).encode('utf-8')
|
||||
|
||||
# Mock for aiohttp.ClientSession
|
||||
mock_aiohttp_session = mock.MagicMock()
|
||||
mock_aiohttp_session.request = mock.AsyncMock(side_effect=[
|
||||
_httpx_response(200, headers={"X-Goog-Upload-URL": "http://fake/upload"}), # Initial request
|
||||
_httpx_response(400, content=error_content, headers={"X-Goog-Upload-Status": "final"}), # Upload request
|
||||
])
|
||||
|
||||
with mock.patch.object(client, '_use_aiohttp', return_value=True), \
|
||||
mock.patch.object(client, '_get_aiohttp_session', return_value=mock_aiohttp_session), \
|
||||
pytest.raises(errors.APIError, match='Unsupported MIME type: bad/mime_type'), \
|
||||
io.BytesIO(b"test") as f:
|
||||
await client._async_upload_fd(
|
||||
f, # Pass file object directly
|
||||
"http://fake/upload",
|
||||
4,
|
||||
http_options=types.HttpOptions(),
|
||||
)
|
||||
assert mock_aiohttp_session.request.call_count == 2
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
# 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 _common module."""
|
||||
@@ -0,0 +1,954 @@
|
||||
# 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 tools in the _common module."""
|
||||
|
||||
from enum import Enum
|
||||
import inspect
|
||||
import logging
|
||||
import textwrap
|
||||
import typing
|
||||
from typing import List, Optional
|
||||
import warnings
|
||||
|
||||
import pydantic
|
||||
import pytest
|
||||
|
||||
from ... import _common
|
||||
from ... import types
|
||||
from ... import errors
|
||||
|
||||
|
||||
def test_warn_once():
|
||||
@_common.experimental_warning('Warning!')
|
||||
def func():
|
||||
pass
|
||||
|
||||
with warnings.catch_warnings(record=True) as w:
|
||||
func()
|
||||
func()
|
||||
|
||||
assert len(w) == 1
|
||||
assert w[0].category == errors.ExperimentalWarning
|
||||
|
||||
def test_warn_at_call_line():
|
||||
@_common.experimental_warning('Warning!')
|
||||
def func():
|
||||
pass
|
||||
|
||||
with warnings.catch_warnings(record=True) as captured_warnings:
|
||||
call_line = inspect.currentframe().f_lineno + 1
|
||||
func()
|
||||
|
||||
assert captured_warnings[0].lineno == call_line
|
||||
|
||||
|
||||
def test_is_struct_type():
|
||||
assert _common._is_struct_type(list[dict[str, typing.Any]])
|
||||
assert _common._is_struct_type(typing.List[typing.Dict[str, typing.Any]])
|
||||
assert not _common._is_struct_type(list[dict[str, int]])
|
||||
assert not _common._is_struct_type(list[dict[int, typing.Any]])
|
||||
assert not _common._is_struct_type(list[str])
|
||||
assert not _common._is_struct_type(dict[str, typing.Any])
|
||||
assert not _common._is_struct_type(typing.List[typing.Dict[str, int]])
|
||||
assert not _common._is_struct_type(typing.List[typing.Dict[int, typing.Any]])
|
||||
assert not _common._is_struct_type(typing.List[str])
|
||||
assert not _common._is_struct_type(typing.Dict[str, typing.Any])
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"test_id, initial_target, update_dict, expected_target",
|
||||
[
|
||||
(
|
||||
"simple_update",
|
||||
{"a": 1, "b": 2},
|
||||
{"b": 3, "c": 4},
|
||||
{"a": 1, "b": 3, "c": 4},
|
||||
),
|
||||
(
|
||||
"nested_update",
|
||||
{"a": 1, "b": {"x": 10, "y": 20}},
|
||||
{"b": {"y": 30, "z": 40}, "c": 3},
|
||||
{"a": 1, "b": {"x": 10, "y": 30, "z": 40}, "c": 3},
|
||||
),
|
||||
(
|
||||
"add_new_nested_dict",
|
||||
{"a": 1},
|
||||
{"b": {"x": 10, "y": 20}},
|
||||
{"a": 1, "b": {"x": 10, "y": 20}},
|
||||
),
|
||||
(
|
||||
"empty_target",
|
||||
{},
|
||||
{"a": 1, "b": {"x": 10}},
|
||||
{"a": 1, "b": {"x": 10}},
|
||||
),
|
||||
(
|
||||
"empty_update",
|
||||
{"a": 1, "b": {"x": 10}},
|
||||
{},
|
||||
{"a": 1, "b": {"x": 10}},
|
||||
),
|
||||
(
|
||||
"overwrite_non_dict_with_dict",
|
||||
{"a": 1, "b": 2},
|
||||
{"b": {"x": 10}},
|
||||
{"a": 1, "b": {"x": 10}},
|
||||
),
|
||||
(
|
||||
"overwrite_dict_with_non_dict",
|
||||
{"a": 1, "b": {"x": 10}},
|
||||
{"b": 2},
|
||||
{"a": 1, "b": 2},
|
||||
),
|
||||
(
|
||||
"deeper_nesting",
|
||||
{"a": {"b": {"c": 1, "d": 2}, "e": 3}},
|
||||
{"a": {"b": {"d": 4, "f": 5}, "g": 6}, "h": 7},
|
||||
{"a": {"b": {"c": 1, "d": 4, "f": 5}, "e": 3, "g": 6}, "h": 7},
|
||||
),
|
||||
(
|
||||
"different_value_types",
|
||||
{"key1": "string_val", "key2": {"nested_int": 100}},
|
||||
{"key1": 123, "key2": {"nested_list": [1, 2, 3]}, "key3": True},
|
||||
{
|
||||
"key1": 123,
|
||||
"key2": {"nested_int": 100, "nested_list": [1, 2, 3]},
|
||||
"key3": True,
|
||||
},
|
||||
),
|
||||
(
|
||||
"update_with_empty_nested_dict", # Existing nested dict in target should not be cleared
|
||||
{"a": {"b": 1}},
|
||||
{"a": {}},
|
||||
{"a": {"b": 1}},
|
||||
),
|
||||
(
|
||||
"target_with_empty_nested_dict",
|
||||
{"a": {}},
|
||||
{"a": {"b": 1}},
|
||||
{"a": {"b": 1}},
|
||||
),
|
||||
(
|
||||
"key_case_alignment_check",
|
||||
{"first_name": "John", "contact_info": {"email_address": "john@example.com"}},
|
||||
{"firstName": "Jane", "contact_info": {"email_address": "jane@example.com", "phone_number": "123"}},
|
||||
{"first_name": "Jane", "contact_info": {"email_address": "jane@example.com", "phone_number": "123"}},
|
||||
)
|
||||
],
|
||||
)
|
||||
def test_recursive_dict_update(
|
||||
test_id: str, initial_target: dict, update_dict: dict, expected_target: dict
|
||||
):
|
||||
_common.recursive_dict_update(initial_target, update_dict)
|
||||
assert initial_target == expected_target
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"test_id, initial_target, update_dict, expected_target, expect_warning, expected_log_message_part",
|
||||
[
|
||||
(
|
||||
"type_match_int",
|
||||
{"a": 1},
|
||||
{"a": 2},
|
||||
{"a": 2},
|
||||
False,
|
||||
"",
|
||||
),
|
||||
(
|
||||
"type_match_dict",
|
||||
{"a": {"b": 1}},
|
||||
{"a": {"b": 2}},
|
||||
{"a": {"b": 2}},
|
||||
False,
|
||||
"",
|
||||
),
|
||||
(
|
||||
"type_mismatch_int_to_str",
|
||||
{"a": 1},
|
||||
{"a": "hello"},
|
||||
{"a": "hello"},
|
||||
True,
|
||||
"Type mismatch for key 'a'. Existing type: <class 'int'>, new type: <class 'str'>. Overwriting.",
|
||||
),
|
||||
(
|
||||
"type_mismatch_dict_to_int",
|
||||
{"a": {"b": 1}},
|
||||
{"a": 100},
|
||||
{"a": 100},
|
||||
True,
|
||||
"Type mismatch for key 'a'. Existing type: <class 'dict'>, new type: <class 'int'>. Overwriting.",
|
||||
),
|
||||
(
|
||||
"type_mismatch_int_to_dict",
|
||||
{"a": 100},
|
||||
{"a": {"b": 1}},
|
||||
{"a": {"b": 1}},
|
||||
True,
|
||||
"Type mismatch for key 'a'. Existing type: <class 'int'>, new type: <class 'dict'>. Overwriting.",
|
||||
),
|
||||
("add_new_key", {"a": 1}, {"b": "new"}, {"a": 1, "b": "new"}, False, ""),
|
||||
],
|
||||
)
|
||||
def test_recursive_dict_update_type_warnings(test_id, initial_target, update_dict, expected_target, expect_warning, expected_log_message_part, caplog):
|
||||
_common.recursive_dict_update(initial_target, update_dict)
|
||||
assert initial_target == expected_target
|
||||
if expect_warning:
|
||||
assert len(caplog.records) == 1
|
||||
assert caplog.records[0].levelname == "WARNING"
|
||||
assert expected_log_message_part in caplog.records[0].message
|
||||
else:
|
||||
for record in caplog.records:
|
||||
if record.levelname == "WARNING" and expected_log_message_part in record.message:
|
||||
pytest.fail(f"Unexpected warning logged for {test_id}: {record.message}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"test_id, target_dict, update_dict, expected_aligned_dict",
|
||||
[
|
||||
(
|
||||
"simple_snake_to_camel",
|
||||
{"first_name": "John", "last_name": "Doe"},
|
||||
{"firstName": "Jane", "lastName": "Doe"},
|
||||
{"first_name": "Jane", "last_name": "Doe"},
|
||||
),
|
||||
(
|
||||
"simple_camel_to_snake",
|
||||
{"firstName": "John", "lastName": "Doe"},
|
||||
{"first_name": "Jane", "last_name": "Doe"},
|
||||
{"firstName": "Jane", "lastName": "Doe"},
|
||||
),
|
||||
(
|
||||
"nested_dict_alignment",
|
||||
{"user_info": {"contact_details": {"email_address": ""}}},
|
||||
{"userInfo": {"contactDetails": {"emailAddress": "test@example.com"}}},
|
||||
{"user_info": {"contact_details": {"email_address": "test@example.com"}}},
|
||||
),
|
||||
(
|
||||
"list_of_dicts_alignment",
|
||||
{"users_list": [{"user_id": 0, "user_name": ""}]},
|
||||
{"usersList": [{"userId": 1, "userName": "Alice"}]},
|
||||
{"users_list": [{"userId": 1, "userName": "Alice"}]},
|
||||
),
|
||||
(
|
||||
"list_of_dicts_alignment_mixed_case_in_update",
|
||||
{"users_list": [{"user_id": 0, "user_name": ""}]},
|
||||
{"usersList": [{"user_id": 1, "UserName": "Alice"}]},
|
||||
{"users_list": [{"user_id": 1, "UserName": "Alice"}]},
|
||||
),
|
||||
(
|
||||
"list_of_dicts_different_lengths_update_longer",
|
||||
{"items_data": [{"item_id": 0}]},
|
||||
{"itemsData": [{"itemId": 1}, {"item_id": 2, "itemName": "Extra"}]},
|
||||
{"items_data": [{"itemId": 1}, {"item_id": 2, "itemName": "Extra"}]},
|
||||
),
|
||||
(
|
||||
"list_of_dicts_different_lengths_target_longer",
|
||||
{"items_data": [{"item_id": 0, "item_name": ""}, {"item_id": 1}]},
|
||||
{"itemsData": [{"itemId": 10}]},
|
||||
{"items_data": [{"itemId": 10}]},
|
||||
),
|
||||
(
|
||||
"no_matching_keys_preserves_update_case",
|
||||
{"key_one": 1},
|
||||
{"KEY_TWO": 2, "keyThree": 3},
|
||||
{"KEY_TWO": 2, "keyThree": 3},
|
||||
),
|
||||
(
|
||||
"mixed_match_and_no_match",
|
||||
{"first_name": "John", "age_years": 30},
|
||||
{"firstName": "Jane", "AGE_YEARS": 28, "occupation_title": "Engineer"},
|
||||
{"first_name": "Jane", "age_years": 28, "occupation_title": "Engineer"},
|
||||
),
|
||||
(
|
||||
"empty_target_dict",
|
||||
{},
|
||||
{"new_key": "new_value", "anotherKey": "anotherValue"},
|
||||
{"new_key": "new_value", "anotherKey": "anotherValue"},
|
||||
),
|
||||
(
|
||||
"empty_update_dict",
|
||||
{"existing_key": "value"},
|
||||
{},
|
||||
{},
|
||||
),
|
||||
(
|
||||
"target_has_non_dict_value_for_nested_key",
|
||||
{"config_settings": 123},
|
||||
{"configSettings": {"themeName": "dark"}},
|
||||
{"config_settings": {"themeName": "dark"}}, # Overwrites as per recursive_dict_update logic
|
||||
),
|
||||
(
|
||||
"update_has_non_dict_value_for_nested_key",
|
||||
{"config_settings": {"theme_name": "light"}},
|
||||
{"configSettings": "dark_theme_string"},
|
||||
{"config_settings": "dark_theme_string"}, # Overwrites
|
||||
),
|
||||
(
|
||||
"deeply_nested_with_lists",
|
||||
{"level_one": {"list_items": [{"item_name": "", "item_value": 0}]}},
|
||||
{"levelOne": {"listItems": [{"itemName": "Test", "itemValue": 100}, {"itemName": "Test2", "itemValue": 200}]}},
|
||||
{"level_one": {"list_items": [{"itemName": "Test", "itemValue": 100}, {"itemName": "Test2", "itemValue": 200}]}},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_align_key_case(
|
||||
test_id: str, target_dict: dict, update_dict: dict, expected_aligned_dict: dict
|
||||
):
|
||||
aligned_dict = _common.align_key_case(target_dict, update_dict)
|
||||
assert aligned_dict == expected_aligned_dict, f"Test failed for: {test_id}"
|
||||
|
||||
|
||||
|
||||
class SimpleModel(_common.BaseModel):
|
||||
name: str
|
||||
value: int
|
||||
is_active: bool = True
|
||||
none_field: Optional[str] = None
|
||||
|
||||
|
||||
class Chain(_common.BaseModel):
|
||||
id: int
|
||||
child: Optional["Chain"] = None
|
||||
|
||||
|
||||
Chain.model_rebuild()
|
||||
|
||||
|
||||
class Tree(_common.BaseModel):
|
||||
id: int
|
||||
children: List["Tree"] = pydantic.Field(default_factory=list)
|
||||
|
||||
|
||||
Tree.model_rebuild()
|
||||
|
||||
|
||||
class ReprFalseModel(_common.BaseModel):
|
||||
visible: str
|
||||
hidden: str = pydantic.Field("secret", repr=False)
|
||||
|
||||
|
||||
class NonPydantic:
|
||||
|
||||
def __repr__(self):
|
||||
return "NonPydantic(\n attr='value'\n)"
|
||||
|
||||
|
||||
class MyEnum(Enum):
|
||||
ONE = 1
|
||||
TWO = 2
|
||||
|
||||
|
||||
class EmptyModel(_common.BaseModel):
|
||||
pass
|
||||
|
||||
|
||||
def test_repr_simple_model_defaults_and_no_none():
|
||||
obj = SimpleModel(name="Test Name", value=123)
|
||||
expected = textwrap.dedent("""
|
||||
SimpleModel(
|
||||
is_active=True,
|
||||
name='Test Name',
|
||||
value=123
|
||||
)
|
||||
""").strip()
|
||||
assert repr(obj) == expected
|
||||
|
||||
|
||||
def test_repr_empty_model():
|
||||
obj = EmptyModel()
|
||||
expected = "EmptyModel()"
|
||||
assert repr(obj) == expected
|
||||
|
||||
|
||||
def test_repr_nested_model():
|
||||
obj = Chain(id=1, child=Chain(id=2))
|
||||
expected = textwrap.dedent("""
|
||||
Chain(
|
||||
child=Chain(
|
||||
id=2
|
||||
),
|
||||
id=1
|
||||
)
|
||||
""").strip()
|
||||
assert repr(obj) == expected
|
||||
|
||||
|
||||
def test_repr_circular_model():
|
||||
obj1 = Chain(id=1)
|
||||
obj2 = Chain(id=2)
|
||||
obj1.child = obj2
|
||||
obj2.child = obj1 # Circular reference
|
||||
expected = textwrap.dedent("""
|
||||
Chain(
|
||||
child=Chain(
|
||||
child=<... Circular reference ...>,
|
||||
id=2
|
||||
),
|
||||
id=1
|
||||
)
|
||||
""").strip()
|
||||
|
||||
assert repr(obj1) == expected
|
||||
|
||||
|
||||
def test_repr_circular_list():
|
||||
my_list = [1, 2]
|
||||
my_list.append(my_list)
|
||||
expected = textwrap.dedent("""
|
||||
[
|
||||
1,
|
||||
2,
|
||||
<... Circular reference ...>,
|
||||
]
|
||||
""").strip()
|
||||
assert _common._pretty_repr(my_list) == expected
|
||||
|
||||
|
||||
def test_repr_circular_dict():
|
||||
my_dict = {"a": 1}
|
||||
my_dict["self"] = my_dict
|
||||
expected = textwrap.dedent("""
|
||||
{
|
||||
'a': 1,
|
||||
'self': <... Circular reference ...>
|
||||
}
|
||||
""").strip()
|
||||
assert _common._pretty_repr(my_dict) == expected
|
||||
|
||||
|
||||
def test_repr_max_items():
|
||||
lst = list(range(10))
|
||||
dct = {i: i for i in range(10)}
|
||||
st = set(range(10))
|
||||
tpl = tuple(range(10))
|
||||
|
||||
assert (
|
||||
"<... 5 more items ...>" in
|
||||
_common._pretty_repr(lst, max_items=5)
|
||||
)
|
||||
assert (
|
||||
"<dict len=10>" in _common._pretty_repr(dct, max_items=5))
|
||||
assert (
|
||||
"<... 5 more items ...>" in _common._pretty_repr(st, max_items=5)
|
||||
)
|
||||
assert (
|
||||
"<... 5 more items ...>" in _common._pretty_repr(tpl, max_items=5)
|
||||
)
|
||||
|
||||
|
||||
def test_repr_max_len_bytes():
|
||||
b_data = b"a" * 100
|
||||
assert len(_common._pretty_repr(b_data, max_len=90)) == 90 + 3
|
||||
assert repr(b_data) == _common._pretty_repr(b_data, max_len=200)
|
||||
|
||||
|
||||
def test_repr_max_depth_dict():
|
||||
nested = {'a': {'a': {'a': {'a': 'a', 'b': 'b'}}}}
|
||||
assert "{<... 2 items at Max depth ...>}" in _common._pretty_repr(nested, depth=3)
|
||||
|
||||
|
||||
def test_repr_max_depth_list():
|
||||
nested = [[[["d", "e", "e", "p"]]]]
|
||||
assert "[<... 4 items at Max depth ...>]" in _common._pretty_repr(nested, depth=3)
|
||||
|
||||
|
||||
def test_repr_collections():
|
||||
obj = {
|
||||
"set": {3, 1, 2},
|
||||
"tuple": (4, 5, 6),
|
||||
"dict": {"b": 2, "a": 1},
|
||||
"list": [7, 8, 9],
|
||||
}
|
||||
expected = textwrap.dedent("""
|
||||
{
|
||||
'dict': {
|
||||
'a': 1,
|
||||
'b': 2
|
||||
},
|
||||
'list': [
|
||||
7,
|
||||
8,
|
||||
9,
|
||||
],
|
||||
'set': {
|
||||
1,
|
||||
2,
|
||||
3,
|
||||
},
|
||||
'tuple': (
|
||||
4,
|
||||
5,
|
||||
6,
|
||||
)
|
||||
}
|
||||
""").strip()
|
||||
assert _common._pretty_repr(obj) == expected
|
||||
|
||||
|
||||
def test_tuple_collections():
|
||||
obj = {
|
||||
"tuple0": (),
|
||||
"tuple1": (1,),
|
||||
"tuple2": (1, 2),
|
||||
}
|
||||
expected = textwrap.dedent("""
|
||||
{
|
||||
'tuple0': (),
|
||||
'tuple1': (
|
||||
1,
|
||||
),
|
||||
'tuple2': (
|
||||
1,
|
||||
2,
|
||||
)
|
||||
}
|
||||
""").strip()
|
||||
assert _common._pretty_repr(obj) == expected
|
||||
|
||||
|
||||
def test_repr_empty_collections():
|
||||
assert _common._pretty_repr([]) == "[]"
|
||||
assert _common._pretty_repr({}) == "{}"
|
||||
assert (
|
||||
_common._pretty_repr(set()) == "set()"
|
||||
)
|
||||
assert _common._pretty_repr(tuple()) == "()"
|
||||
assert (
|
||||
_common._pretty_repr({"empty_set": set()}) ==
|
||||
textwrap.dedent("""
|
||||
{
|
||||
'empty_set': set()
|
||||
}
|
||||
""").strip()
|
||||
)
|
||||
|
||||
|
||||
def test_repr_strings():
|
||||
s1 = "line one"
|
||||
exp1 = "'line one'"
|
||||
assert _common._pretty_repr(s1) == exp1
|
||||
|
||||
s2 = 'line one\nline two with """ inside'
|
||||
exp2 = '"""line one\nline two with \\"\\"\\" inside"""'
|
||||
assert _common._pretty_repr(s2) == exp2
|
||||
|
||||
s3 = 'A string with """ inside'
|
||||
exp3 = '\'A string with """ inside\''
|
||||
assert _common._pretty_repr(s3) == exp3
|
||||
|
||||
|
||||
def test_repr_repr_false():
|
||||
obj = ReprFalseModel(visible="show", hidden="hide")
|
||||
result = repr(obj)
|
||||
assert "visible='show'" in result
|
||||
assert "hidden" not in result
|
||||
expected = textwrap.dedent("""
|
||||
ReprFalseModel(
|
||||
visible='show'
|
||||
)
|
||||
""").strip()
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_repr_none_fields():
|
||||
obj = SimpleModel(name="Only Name", value=0, none_field=None)
|
||||
result = repr(obj)
|
||||
assert "none_field" not in result
|
||||
expected = textwrap.dedent("""
|
||||
SimpleModel(
|
||||
is_active=True,
|
||||
name='Only Name',
|
||||
value=0
|
||||
)
|
||||
""").strip()
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_repr_other_types():
|
||||
np = NonPydantic()
|
||||
en = MyEnum.TWO
|
||||
obj = {"np": np, "en": en}
|
||||
expected = textwrap.dedent("""
|
||||
{
|
||||
'en': <MyEnum.TWO: 2>,
|
||||
'np': NonPydantic(
|
||||
attr='value'
|
||||
)
|
||||
}
|
||||
""").strip()
|
||||
assert _common._pretty_repr(obj) == expected
|
||||
|
||||
|
||||
def test_repr_indent_delta():
|
||||
obj = SimpleModel(name="Indent Test", value=1)
|
||||
expected = textwrap.dedent("""
|
||||
SimpleModel(
|
||||
is_active=True,
|
||||
name='Indent Test',
|
||||
value=1
|
||||
)
|
||||
""").strip()
|
||||
assert _common._pretty_repr(obj, indent_delta=4) == expected
|
||||
|
||||
|
||||
def test_repr_complex_object():
|
||||
obj = types.GenerateContentResponse(
|
||||
automatic_function_calling_history=[],
|
||||
candidates=[
|
||||
types.Candidate(
|
||||
content=types.Content(
|
||||
parts=[
|
||||
types.Part(
|
||||
text="""There isn't a single "best" LLM, as the ideal choice highly depends on your specific needs, use case, budget, and priorities. The field is evolving incredibly fast, with new models and improvements being released constantly.
|
||||
|
||||
However, we can talk about the **leading contenders** and what they are generally known for:..."""
|
||||
)
|
||||
],
|
||||
role="model"
|
||||
),
|
||||
finish_reason=types.FinishReason.STOP,
|
||||
index=0
|
||||
)
|
||||
],
|
||||
model_version='models/gemini-2.5-flash-preview-05-20',
|
||||
usage_metadata=types.GenerateContentResponseUsageMetadata(
|
||||
candidates_token_count=1086,
|
||||
prompt_token_count=7,
|
||||
prompt_tokens_details=[
|
||||
types.ModalityTokenCount(
|
||||
modality=types.MediaModality.TEXT,
|
||||
token_count=7
|
||||
)
|
||||
],
|
||||
thoughts_token_count=860,
|
||||
total_token_count=1953
|
||||
)
|
||||
)
|
||||
|
||||
expected = textwrap.dedent("""
|
||||
GenerateContentResponse(
|
||||
automatic_function_calling_history=[],
|
||||
candidates=[
|
||||
Candidate(
|
||||
content=Content(
|
||||
parts=[
|
||||
Part(
|
||||
text=\"\"\"There isn't a single "best" LLM, as the ideal choice highly depends on your specific needs, use case, budget, and priorities. The field is evolving incredibly fast, with new models and improvements being released constantly.
|
||||
|
||||
However, we can talk about the **leading contenders** and what they are generally known for:...\"\"\"
|
||||
),
|
||||
],
|
||||
role='model'
|
||||
),
|
||||
finish_reason=<FinishReason.STOP: 'STOP'>,
|
||||
index=0
|
||||
),
|
||||
],
|
||||
model_version='models/gemini-2.5-flash-preview-05-20',
|
||||
usage_metadata=GenerateContentResponseUsageMetadata(
|
||||
candidates_token_count=1086,
|
||||
prompt_token_count=7,
|
||||
prompt_tokens_details=[
|
||||
ModalityTokenCount(
|
||||
modality=<MediaModality.TEXT: 'TEXT'>,
|
||||
token_count=7
|
||||
),
|
||||
],
|
||||
thoughts_token_count=860,
|
||||
total_token_count=1953
|
||||
)
|
||||
)
|
||||
""").strip()
|
||||
assert repr(obj) == expected
|
||||
|
||||
|
||||
def test_move_value_by_path():
|
||||
"""Test move_value_by_path function with array wildcard notation."""
|
||||
data = {
|
||||
"requests": [
|
||||
{
|
||||
"request": {
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"text": "1"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"outputDimensionality": 64
|
||||
},
|
||||
{
|
||||
"request": {
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"text": "2"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"outputDimensionality": 64
|
||||
},
|
||||
{
|
||||
"request": {
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"text": "3"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"outputDimensionality": 64
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
paths = {'requests[].*': 'requests[].request.*'}
|
||||
_common.move_value_by_path(data, paths)
|
||||
|
||||
expected = {
|
||||
"requests": [
|
||||
{
|
||||
"request": {
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"text": "1"
|
||||
}
|
||||
]
|
||||
},
|
||||
"outputDimensionality": 64
|
||||
}
|
||||
},
|
||||
{
|
||||
"request": {
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"text": "2"
|
||||
}
|
||||
]
|
||||
},
|
||||
"outputDimensionality": 64
|
||||
}
|
||||
},
|
||||
{
|
||||
"request": {
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"text": "3"
|
||||
}
|
||||
]
|
||||
},
|
||||
"outputDimensionality": 64
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
assert data == expected
|
||||
|
||||
|
||||
def test_check_field_type_mismatches_no_warning_for_correct_types(caplog):
|
||||
"""Test that no warning is logged when types match."""
|
||||
|
||||
class ModelA(_common.BaseModel):
|
||||
value: int
|
||||
|
||||
class TestModel(_common.BaseModel):
|
||||
model_a: ModelA
|
||||
|
||||
# Should not warn - dict will be converted to ModelA by Pydantic
|
||||
data = {"model_a": {"value": 123}}
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
|
||||
result = TestModel.model_validate(data)
|
||||
|
||||
assert result.model_a.value == 123
|
||||
assert len(caplog.records) == 0
|
||||
|
||||
|
||||
def test_check_field_type_mismatches_warns_on_pydantic_type_mismatch(caplog):
|
||||
"""Test that warning is logged when Pydantic model types mismatch."""
|
||||
|
||||
class ModelA(_common.BaseModel):
|
||||
value: int
|
||||
|
||||
class ModelB(_common.BaseModel):
|
||||
value: str
|
||||
|
||||
class TestModel(_common.BaseModel):
|
||||
model_field: ModelA
|
||||
|
||||
# Create an instance of ModelB (wrong type)
|
||||
model_b_instance = ModelB(value="test")
|
||||
|
||||
# Pass the wrong Pydantic model instance
|
||||
data = {"model_field": model_b_instance}
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
|
||||
TestModel._check_field_type_mismatches(data)
|
||||
|
||||
assert len(caplog.records) == 1
|
||||
assert "Type mismatch in TestModel.model_field" in caplog.records[0].message
|
||||
assert "expected ModelA, got ModelB" in caplog.records[0].message
|
||||
|
||||
|
||||
def test_check_field_type_mismatches_no_warning_for_none_values(caplog):
|
||||
"""Test that no warning is logged for None values."""
|
||||
|
||||
class ModelA(_common.BaseModel):
|
||||
value: int
|
||||
|
||||
class TestModel(_common.BaseModel):
|
||||
model_field: Optional[ModelA] = None
|
||||
|
||||
data = {"model_field": None}
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
|
||||
result = TestModel.model_validate(data)
|
||||
|
||||
assert result.model_field is None
|
||||
assert len(caplog.records) == 0
|
||||
|
||||
|
||||
def test_check_field_type_mismatches_no_warning_for_missing_fields(caplog):
|
||||
"""Test that no warning is logged for missing fields."""
|
||||
|
||||
class ModelA(_common.BaseModel):
|
||||
value: int
|
||||
|
||||
class TestModel(_common.BaseModel):
|
||||
model_field: Optional[ModelA] = None
|
||||
|
||||
data = {}
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
|
||||
result = TestModel.model_validate(data)
|
||||
|
||||
assert result.model_field is None
|
||||
assert len(caplog.records) == 0
|
||||
|
||||
|
||||
def test_check_field_type_mismatches_no_warning_for_primitive_types(caplog):
|
||||
"""Test that no warning is logged for primitive type mismatches."""
|
||||
|
||||
class TestModel(_common.BaseModel):
|
||||
int_field: int
|
||||
str_field: str
|
||||
|
||||
# Even though we're passing wrong primitive types, we should not warn
|
||||
# (Pydantic will handle validation)
|
||||
data = {"int_field": "123", "str_field": "test"}
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
|
||||
# This will succeed because Pydantic can coerce "123" to int
|
||||
result = TestModel.model_validate(data)
|
||||
|
||||
assert result.int_field == 123
|
||||
assert result.str_field == "test"
|
||||
assert len(caplog.records) == 0
|
||||
|
||||
|
||||
def test_check_field_type_mismatches_handles_optional_unwrapping(caplog):
|
||||
"""Test that Optional types are properly unwrapped before checking."""
|
||||
|
||||
class ModelA(_common.BaseModel):
|
||||
value: int
|
||||
|
||||
class ModelB(_common.BaseModel):
|
||||
value: str
|
||||
|
||||
class TestModel(_common.BaseModel):
|
||||
model_field: Optional[ModelA] = None
|
||||
|
||||
# Pass wrong Pydantic model type
|
||||
model_b_instance = ModelB(value="test")
|
||||
data = {"model_field": model_b_instance}
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
|
||||
TestModel._check_field_type_mismatches(data)
|
||||
|
||||
assert len(caplog.records) == 1
|
||||
assert "expected ModelA, got ModelB" in caplog.records[0].message
|
||||
|
||||
|
||||
def test_check_field_type_mismatches_no_warning_for_correct_pydantic_instance(caplog):
|
||||
"""Test that no warning is logged when correct Pydantic instance is provided."""
|
||||
|
||||
class ModelA(_common.BaseModel):
|
||||
value: int
|
||||
|
||||
class TestModel(_common.BaseModel):
|
||||
model_field: ModelA
|
||||
|
||||
# Pass correct Pydantic model instance
|
||||
model_a_instance = ModelA(value=42)
|
||||
data = {"model_field": model_a_instance}
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
|
||||
result = TestModel.model_validate(data)
|
||||
|
||||
assert result.model_field.value == 42
|
||||
assert len(caplog.records) == 0
|
||||
|
||||
|
||||
def test_check_field_type_mismatches_with_multiple_fields(caplog):
|
||||
"""Test checking multiple fields with mixed scenarios."""
|
||||
|
||||
class ModelA(_common.BaseModel):
|
||||
value: int
|
||||
|
||||
class ModelB(_common.BaseModel):
|
||||
value: str
|
||||
|
||||
class TestModel(_common.BaseModel):
|
||||
field_a: ModelA
|
||||
field_b: Optional[ModelA] = None
|
||||
field_c: str
|
||||
|
||||
model_b_instance = ModelB(value="wrong")
|
||||
data = {
|
||||
"field_a": model_b_instance, # Wrong type - should warn
|
||||
"field_b": None, # None - should not warn
|
||||
"field_c": "test", # Primitive - should not warn
|
||||
}
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
|
||||
TestModel._check_field_type_mismatches(data)
|
||||
|
||||
# Should only warn about field_a
|
||||
assert len(caplog.records) == 1
|
||||
assert "field_a" in caplog.records[0].message
|
||||
assert "expected ModelA, got ModelB" in caplog.records[0].message
|
||||
|
||||
|
||||
def test_check_field_type_mismatches_generic_type_no_error(caplog):
|
||||
"""Test that validation doesn't crash on generic types like list[str]."""
|
||||
class TestModel(_common.BaseModel):
|
||||
tags: list[str]
|
||||
|
||||
data = {"tags": ["a", "b"]}
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="google_genai._common"):
|
||||
TestModel.model_validate(data)
|
||||
|
||||
assert len(caplog.records) == 0
|
||||
@@ -0,0 +1,96 @@
|
||||
import unittest
|
||||
|
||||
import pydantic
|
||||
|
||||
from ... import _common
|
||||
|
||||
|
||||
class TestIsDuckTypeOf(unittest.TestCase):
|
||||
|
||||
class FakePydanticModel(pydantic.BaseModel):
|
||||
field1: str
|
||||
field2: int
|
||||
field3: str
|
||||
field4: str
|
||||
field5: str
|
||||
|
||||
class FakePydanticModelWithLessFields(pydantic.BaseModel):
|
||||
field1: str
|
||||
field2: int
|
||||
field3: str
|
||||
field4: str
|
||||
|
||||
def test_is_duck_type_of_true_for_pydantic_object(self):
|
||||
obj = self.FakePydanticModel(
|
||||
field1="a",
|
||||
field2=1,
|
||||
field3="b",
|
||||
field4="c",
|
||||
field5="d",
|
||||
)
|
||||
self.assertTrue(_common.is_duck_type_of(obj, self.FakePydanticModel))
|
||||
|
||||
def test_is_duck_type_of_true_for_duck_typed_object(self):
|
||||
class DuckTypedObject:
|
||||
|
||||
def __init__(self):
|
||||
self.field1 = "a"
|
||||
self.field2 = 1
|
||||
self.field3 = "b"
|
||||
self.field4 = "c"
|
||||
self.field5 = "d"
|
||||
|
||||
obj = DuckTypedObject()
|
||||
self.assertTrue(_common.is_duck_type_of(obj, self.FakePydanticModel))
|
||||
|
||||
def test_is_duck_type_of_false_for_different_many_fields(self):
|
||||
class DifferentFieldsObject:
|
||||
|
||||
def __init__(self):
|
||||
self.fielda = "a"
|
||||
|
||||
obj = DifferentFieldsObject()
|
||||
self.assertFalse(
|
||||
_common.is_duck_type_of(obj, self.FakePydanticModel)
|
||||
)
|
||||
|
||||
def test_is_duck_type_of_false_for_missing_fields(self):
|
||||
|
||||
obj = self.FakePydanticModelWithLessFields(
|
||||
field1="a", field2=1, field3="b", field4="c"
|
||||
)
|
||||
self.assertFalse(_common.is_duck_type_of(obj, self.FakePydanticModel))
|
||||
|
||||
def test_is_duck_type_of_false_for_dict(self):
|
||||
obj = {"field1": "a", "field2": 1}
|
||||
self.assertFalse(
|
||||
_common.is_duck_type_of(obj, self.FakePydanticModel)
|
||||
)
|
||||
|
||||
def test_is_duck_type_of_false_for_non_pydantic_class(self):
|
||||
class NonPydanticModel:
|
||||
pass
|
||||
|
||||
class SomeObject:
|
||||
pass
|
||||
|
||||
obj = SomeObject()
|
||||
self.assertFalse(_common.is_duck_type_of(obj, NonPydanticModel))
|
||||
|
||||
def test_is_duck_type_of_true_with_extra_fields(self):
|
||||
class ExtraFieldsObject:
|
||||
|
||||
def __init__(self):
|
||||
self.field1 = "a"
|
||||
self.field2 = 1
|
||||
self.field3 = "extra"
|
||||
self.field4 = "extra"
|
||||
self.field5 = "extra"
|
||||
self.field6 = "extra"
|
||||
|
||||
obj = ExtraFieldsObject()
|
||||
self.assertTrue(_common.is_duck_type_of(obj, self.FakePydanticModel))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
162
venv/lib/python3.12/site-packages/google/genai/tests/conftest.py
Normal file
162
venv/lib/python3.12/site-packages/google/genai/tests/conftest.py
Normal file
@@ -0,0 +1,162 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""Conftest for google.genai tests."""
|
||||
|
||||
import datetime
|
||||
import os
|
||||
from unittest import mock
|
||||
import uuid
|
||||
import pytest
|
||||
|
||||
from .. import _common
|
||||
from .. import _replay_api_client
|
||||
from .. import client as google_genai_client_module
|
||||
|
||||
|
||||
def pytest_addoption(parser):
|
||||
parser.addoption(
|
||||
'--mode',
|
||||
action='store',
|
||||
default='auto',
|
||||
help="""Replay mode.
|
||||
One of:
|
||||
* auto: Replay if replay files exist, otherwise record.
|
||||
* record: Always call the API and record.
|
||||
* replay: Always replay, fail if replay files do not exist.
|
||||
* api: Always call the API and do not record.
|
||||
* tap: Always replay, fail if replay files do not exist. Also sets default values for the API key and replay directory.
|
||||
""",
|
||||
)
|
||||
parser.addoption(
|
||||
'--private',
|
||||
action='store_true',
|
||||
default=False,
|
||||
help='Run private tests.',
|
||||
)
|
||||
|
||||
|
||||
# Overridden via parameterized test.
|
||||
@pytest.fixture
|
||||
def use_vertex():
|
||||
return False
|
||||
|
||||
|
||||
# Overridden at the module level for each test file.
|
||||
@pytest.fixture
|
||||
def replays_prefix():
|
||||
return 'test'
|
||||
|
||||
|
||||
def _get_replay_id(use_vertex: bool, replays_prefix: str) -> str:
|
||||
test_name_ending = os.environ.get('PYTEST_CURRENT_TEST').split('::')[-1]
|
||||
test_name = (
|
||||
test_name_ending.split(' ')[0].split('[')[0]
|
||||
+ '.'
|
||||
+ ('vertex' if use_vertex else 'mldev')
|
||||
)
|
||||
return '/'.join([replays_prefix, test_name])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(use_vertex, replays_prefix, http_options, request):
|
||||
mode = request.config.getoption('--mode')
|
||||
if mode not in ['auto', 'record', 'replay', 'api', 'tap']:
|
||||
raise ValueError('Invalid mode: ' + mode)
|
||||
test_function_name = request.function.__name__
|
||||
test_filename = os.path.splitext(os.path.basename(request.path))[0]
|
||||
if test_function_name.startswith(test_filename):
|
||||
raise ValueError(f"""
|
||||
{test_function_name}:
|
||||
Do not include the test filename in the test function name.
|
||||
keep the test function name short.""")
|
||||
for word in ['mldev', 'ml_dev', 'vertex']:
|
||||
if word in test_function_name:
|
||||
raise ValueError(f"""
|
||||
{test_function_name}:
|
||||
You should never ever include api types in the file name (mldev or vertex).
|
||||
All tests should run against all apis.
|
||||
Assert an exception if the test is not supported in an API.""")
|
||||
replay_id = _get_replay_id(use_vertex, replays_prefix)
|
||||
|
||||
if mode == 'tap':
|
||||
mode = 'replay'
|
||||
|
||||
# Set various environment variables to ensure that the test runs.
|
||||
os.environ['GOOGLE_API_KEY'] = 'dummy-api-key'
|
||||
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
'credentials.json',
|
||||
)
|
||||
os.environ['GOOGLE_CLOUD_PROJECT'] = 'project-id'
|
||||
os.environ['GOOGLE_CLOUD_LOCATION'] = 'location'
|
||||
|
||||
# Set the replay directory to the root directory of the replays.
|
||||
# This is needed to ensure that the replay files are found.
|
||||
replays_root_directory = os.path.abspath(
|
||||
os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
'../../../../../google/cloud/aiplatform/sdk/genai/replays',
|
||||
)
|
||||
)
|
||||
os.environ['GOOGLE_GENAI_REPLAYS_DIRECTORY'] = replays_root_directory
|
||||
# Get private arg.
|
||||
private = request.config.getoption('--private')
|
||||
replay_client = _replay_api_client.ReplayApiClient(
|
||||
mode=mode,
|
||||
replay_id=replay_id,
|
||||
vertexai=use_vertex,
|
||||
http_options=http_options,
|
||||
private=private,
|
||||
)
|
||||
|
||||
with mock.patch.object(
|
||||
google_genai_client_module.Client, '_get_api_client'
|
||||
) as patch_method:
|
||||
patch_method.return_value = replay_client
|
||||
google_genai_client = google_genai_client_module.Client(vertexai=use_vertex)
|
||||
|
||||
# Yield the client so that cleanup can be completed at the end of the test.
|
||||
yield google_genai_client
|
||||
|
||||
# Save the replay after the test if we were in recording mode.
|
||||
google_genai_client._api_client.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_timestamped_unique_name():
|
||||
with mock.patch.object(
|
||||
_common,
|
||||
'timestamped_unique_name',
|
||||
return_value='20240101000000_bd656',
|
||||
) as unique_name_mock:
|
||||
yield unique_name_mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def image_jpeg():
|
||||
import PIL.Image
|
||||
image_path = os.path.join(os.path.dirname(__file__), 'data', 'google.jpg')
|
||||
with PIL.Image.open(image_path) as img:
|
||||
yield img
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def image_png():
|
||||
import PIL.Image
|
||||
image_path = os.path.join(os.path.dirname(__file__), 'data', 'google.png')
|
||||
with PIL.Image.open(image_path) as img:
|
||||
yield img
|
||||
@@ -0,0 +1,17 @@
|
||||
# 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 documents module."""
|
||||
@@ -0,0 +1,51 @@
|
||||
# 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 documents.delete()."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
# All tests will be run for both Vertex and MLDev.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_delete",
|
||||
parameters=types._DeleteDocumentParameters(
|
||||
name=(
|
||||
"fileSearchStores/fr3l0ri2so25-a3r1ump9x821/documents/asurveyofmodernistpoetrytxt-01zyb2rig5gu"
|
||||
),
|
||||
config=types.DeleteDocumentConfig(force=True),
|
||||
),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method="file_search_stores.documents.delete",
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_delete(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
await client.aio.file_search_stores.documents.delete(
|
||||
name="fileSearchStores/fr3l0ri2so25-a3r1ump9x821/documents/asurveyofmodernistpoetrytxt-uvmqjtmkm1h2",
|
||||
config=types.DeleteDocumentConfig(force=True),
|
||||
)
|
||||
@@ -0,0 +1,85 @@
|
||||
# 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 file_search_stores.documents.get()."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
# A Document name known to exist in the test environment.
|
||||
# Replace with a real one from your test setup.
|
||||
_EXISTING_DOCUMENT_NAME = "fileSearchStores/fr3l0ri2so25-a3r1ump9x821/documents/asurveyofmodernistpoetrytxt-uvmqjtmkm1h2"
|
||||
_NON_EXISTENT_DOCUMENT_NAME = (
|
||||
"fileSearchStores/fr3l0ri2so25-a3r1ump9x821/documents/non-existent-document"
|
||||
)
|
||||
_INVALID_DOCUMENT_NAME = "fileSearchStores/fr3l0ri2so25-a3r1ump9x821/documents/_invalid_document_name"
|
||||
_NOT_A_DOCUMENT_NAME = "genai-test-document"
|
||||
|
||||
# All tests will be run for both Vertex and MLDev.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_get_success",
|
||||
parameters=types._GetDocumentParameters(name=_EXISTING_DOCUMENT_NAME),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_get_not_found",
|
||||
parameters=types._GetDocumentParameters(
|
||||
name=_NON_EXISTENT_DOCUMENT_NAME
|
||||
),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
exception_if_mldev="Documents does not exist",
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_get_invalid_name",
|
||||
parameters=types._GetDocumentParameters(name=_INVALID_DOCUMENT_NAME),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
# Validation should catch this before the API call
|
||||
exception_if_mldev="INVALID_ARGUMENT",
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_get_not_a_document_name",
|
||||
parameters=types._GetDocumentParameters(name=_NOT_A_DOCUMENT_NAME),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
exception_if_mldev="Not Found",
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method="file_search_stores.documents.get",
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_get(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
# This test relies on the mocking/replays set up by pytest_helper
|
||||
# For a success case:
|
||||
try:
|
||||
document = await client.aio.file_search_stores.documents.get(
|
||||
name=_EXISTING_DOCUMENT_NAME
|
||||
)
|
||||
assert document is not None
|
||||
assert document.name == _EXISTING_DOCUMENT_NAME
|
||||
except Exception as e:
|
||||
# Depending on the mock setup, errors might be raised directly
|
||||
print(f"Async get failed: {e}")
|
||||
raise
|
||||
@@ -0,0 +1,72 @@
|
||||
# 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 file_search_stores.documents.list()."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
_FILE_SEARCH_STORE_NAME = "fileSearchStores/gzn7kdl2wpxl-4z2yqvuxbcxw"
|
||||
|
||||
# All tests will be run for both Vertex and MLDev.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_list_default",
|
||||
parameters=types._ListDocumentsParameters(
|
||||
parent=_FILE_SEARCH_STORE_NAME
|
||||
),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_list_with_page_size",
|
||||
parameters=types._ListDocumentsParameters(
|
||||
parent=_FILE_SEARCH_STORE_NAME,
|
||||
config=types.ListDocumentsConfig(
|
||||
page_size=2,
|
||||
),
|
||||
),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method="file_search_stores.documents.list",
|
||||
test_table=test_table,
|
||||
http_options={
|
||||
"base_url": "https://autopush-generativelanguage.sandbox.googleapis.com/",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_pager(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
# Iterate through the first page.
|
||||
for document in client.file_search_stores.documents.list(
|
||||
parent=_FILE_SEARCH_STORE_NAME, config={"page_size": 2}
|
||||
):
|
||||
assert isinstance(document, types.Document)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pager(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
# Iterate through the first page.
|
||||
async for document in await client.aio.file_search_stores.documents.list(
|
||||
parent=_FILE_SEARCH_STORE_NAME, config={"page_size": 2}
|
||||
):
|
||||
assert isinstance(document, types.Document)
|
||||
@@ -0,0 +1 @@
|
||||
"""Unit Tests for the error modules."""
|
||||
@@ -0,0 +1,454 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""Unit Tests for the APIError class.
|
||||
|
||||
End to end tests should be in models/test_generate_content.py.
|
||||
"""
|
||||
|
||||
import pickle
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from ... import errors
|
||||
|
||||
|
||||
def test_constructor_code_none_error_in_json_code_in_error():
|
||||
|
||||
actual_error = errors.APIError(
|
||||
None,
|
||||
{
|
||||
'error': {
|
||||
'code': 400,
|
||||
'message': 'error message',
|
||||
'status': 'INVALID_ARGUMENT',
|
||||
}
|
||||
},
|
||||
httpx.Response(status_code=400),
|
||||
)
|
||||
|
||||
assert actual_error.code == 400
|
||||
assert actual_error.message == 'error message'
|
||||
assert actual_error.status == 'INVALID_ARGUMENT'
|
||||
assert actual_error.details == {
|
||||
'error': {
|
||||
'code': 400,
|
||||
'message': 'error message',
|
||||
'status': 'INVALID_ARGUMENT',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_constructor_code_none_error_in_json_code_outside_error():
|
||||
|
||||
actual_error = errors.APIError(
|
||||
None,
|
||||
{
|
||||
'code': 400,
|
||||
'error': {
|
||||
'code': 500,
|
||||
'message': 'error message',
|
||||
'status': 'INVALID_ARGUMENT',
|
||||
},
|
||||
},
|
||||
httpx.Response(status_code=400),
|
||||
)
|
||||
|
||||
assert actual_error.code == 400
|
||||
assert actual_error.message == 'error message'
|
||||
assert actual_error.status == 'INVALID_ARGUMENT'
|
||||
assert actual_error.details == {
|
||||
'code': 400,
|
||||
'error': {
|
||||
'code': 500,
|
||||
'message': 'error message',
|
||||
'status': 'INVALID_ARGUMENT',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_constructor_code_not_present():
|
||||
|
||||
actual_error = errors.APIError(
|
||||
None,
|
||||
{
|
||||
'error': {
|
||||
'message': 'error message',
|
||||
'status': 'INVALID_ARGUMENT',
|
||||
}
|
||||
},
|
||||
httpx.Response(status_code=400),
|
||||
)
|
||||
|
||||
assert actual_error.code is None
|
||||
assert actual_error.message == 'error message'
|
||||
assert actual_error.status == 'INVALID_ARGUMENT'
|
||||
assert actual_error.details == {
|
||||
'error': {
|
||||
'message': 'error message',
|
||||
'status': 'INVALID_ARGUMENT',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_constructor_code_exist_error_in_json():
|
||||
|
||||
actual_error = errors.APIError(
|
||||
400,
|
||||
{
|
||||
'error': {
|
||||
'code': 400,
|
||||
'message': 'error message',
|
||||
'status': 'INVALID_ARGUMENT',
|
||||
}
|
||||
},
|
||||
httpx.Response(status_code=400),
|
||||
)
|
||||
|
||||
assert actual_error.code == 400
|
||||
assert actual_error.message == 'error message'
|
||||
assert actual_error.status == 'INVALID_ARGUMENT'
|
||||
assert actual_error.details == {
|
||||
'error': {
|
||||
'code': 400,
|
||||
'message': 'error message',
|
||||
'status': 'INVALID_ARGUMENT',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_constructor_error_not_in_json():
|
||||
|
||||
actual_error = errors.APIError(
|
||||
400,
|
||||
{
|
||||
'message': 'error message',
|
||||
'status': 'INVALID_ARGUMENT',
|
||||
'code': 400,
|
||||
},
|
||||
httpx.Response(status_code=400),
|
||||
)
|
||||
|
||||
assert actual_error.code == 400
|
||||
assert actual_error.message == 'error message'
|
||||
assert actual_error.status == 'INVALID_ARGUMENT'
|
||||
assert actual_error.details == {
|
||||
'message': 'error message',
|
||||
'status': 'INVALID_ARGUMENT',
|
||||
'code': 400,
|
||||
}
|
||||
|
||||
|
||||
def test_constructor_error_in_json_status_outside_error():
|
||||
|
||||
actual_error = errors.APIError(
|
||||
400,
|
||||
{
|
||||
'status': 'OUTER_INVALID_ARGUMENT_STATUS',
|
||||
'error': {
|
||||
'code': 400,
|
||||
'message': 'error message',
|
||||
'status': 'INNER_INVALID_ARGUMENT_STATUS',
|
||||
},
|
||||
},
|
||||
httpx.Response(status_code=400),
|
||||
)
|
||||
|
||||
assert actual_error.code == 400
|
||||
assert actual_error.message == 'error message'
|
||||
assert actual_error.status == 'OUTER_INVALID_ARGUMENT_STATUS'
|
||||
assert actual_error.details == {
|
||||
'status': 'OUTER_INVALID_ARGUMENT_STATUS',
|
||||
'error': {
|
||||
'code': 400,
|
||||
'message': 'error message',
|
||||
'status': 'INNER_INVALID_ARGUMENT_STATUS',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_constructor_status_not_present():
|
||||
|
||||
actual_error = errors.APIError(
|
||||
400,
|
||||
{
|
||||
'error': {
|
||||
'code': 400,
|
||||
'message': 'error message',
|
||||
}
|
||||
},
|
||||
httpx.Response(status_code=400),
|
||||
)
|
||||
|
||||
assert actual_error.code == 400
|
||||
assert actual_error.message == 'error message'
|
||||
assert actual_error.status == None
|
||||
assert actual_error.details == {
|
||||
'error': {
|
||||
'code': 400,
|
||||
'message': 'error message',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_constructor_error_in_json_message_outside_error():
|
||||
|
||||
actual_error = errors.APIError(
|
||||
400,
|
||||
{
|
||||
'message': 'OUTER_ERROR_MESSAGE',
|
||||
'error': {
|
||||
'code': 400,
|
||||
'message': 'INNER_ERROR_MESSAGE',
|
||||
'status': 'INVALID_ARGUMENT',
|
||||
},
|
||||
},
|
||||
httpx.Response(status_code=400),
|
||||
)
|
||||
|
||||
assert actual_error.code == 400
|
||||
assert actual_error.message == 'OUTER_ERROR_MESSAGE'
|
||||
assert actual_error.status == 'INVALID_ARGUMENT'
|
||||
assert actual_error.details == {
|
||||
'message': 'OUTER_ERROR_MESSAGE',
|
||||
'error': {
|
||||
'code': 400,
|
||||
'message': 'INNER_ERROR_MESSAGE',
|
||||
'status': 'INVALID_ARGUMENT',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_constructor_message_not_present():
|
||||
|
||||
actual_error = errors.APIError(
|
||||
400,
|
||||
{
|
||||
'error': {
|
||||
'code': 400,
|
||||
'status': 'INVALID_ARGUMENT',
|
||||
}
|
||||
},
|
||||
httpx.Response(status_code=400),
|
||||
)
|
||||
|
||||
assert actual_error.code == 400
|
||||
assert actual_error.message is None
|
||||
assert actual_error.status == 'INVALID_ARGUMENT'
|
||||
assert actual_error.details == {
|
||||
'error': {
|
||||
'code': 400,
|
||||
'status': 'INVALID_ARGUMENT',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_constructor_with_websocket_connection_closed_error():
|
||||
actual_error = errors.APIError(
|
||||
1007,
|
||||
'At most one response modality can be specified in the setup request.'
|
||||
' To enable simultaneous transcription and audio output,',
|
||||
None,
|
||||
)
|
||||
assert actual_error.code == 1007
|
||||
assert (
|
||||
actual_error.details
|
||||
== 'At most one response modality can be specified in the setup request.'
|
||||
' To enable simultaneous transcription and audio output,'
|
||||
)
|
||||
assert actual_error.status == None
|
||||
assert actual_error.message == None
|
||||
|
||||
|
||||
def test_raise_for_websocket_connection_closed_error():
|
||||
try:
|
||||
errors.APIError.raise_error(
|
||||
1007,
|
||||
'At most one response modality can be specified in the setup request.'
|
||||
' To enable simultaneous transcription and audio output,',
|
||||
None,
|
||||
)
|
||||
except errors.APIError as actual_error:
|
||||
assert actual_error.code == 1007
|
||||
assert (
|
||||
actual_error.details
|
||||
== 'At most one response modality can be specified in the setup'
|
||||
' request.'
|
||||
' To enable simultaneous transcription and audio output,'
|
||||
)
|
||||
assert actual_error.status == None
|
||||
assert actual_error.message == None
|
||||
|
||||
|
||||
def test_raise_for_response_code_exist_json_decoder_error():
|
||||
class FakeResponse(httpx.Response):
|
||||
|
||||
def read(self) -> bytes:
|
||||
self._content = b'{"data": {"key1": "value1", "key2"}'
|
||||
return self._content
|
||||
|
||||
try:
|
||||
errors.APIError.raise_for_response(
|
||||
FakeResponse(
|
||||
status_code=503,
|
||||
extensions={'reason_phrase': b'Service Unavailable'},
|
||||
)
|
||||
)
|
||||
except errors.ServerError as actual_error:
|
||||
assert actual_error.code == 503
|
||||
assert actual_error.message == '{"data": {"key1": "value1", "key2"}'
|
||||
assert actual_error.status == 'Service Unavailable'
|
||||
assert actual_error.details == {
|
||||
'message': '{"data": {"key1": "value1", "key2"}',
|
||||
'status': 'Service Unavailable',
|
||||
}
|
||||
|
||||
|
||||
def test_raise_for_response_client_error():
|
||||
class FakeResponse(httpx.Response):
|
||||
|
||||
def read(self) -> bytes:
|
||||
self._content = (
|
||||
b'{"error": {"code": 400, "message": "error message", "status":'
|
||||
b' "INVALID_ARGUMENT"}}'
|
||||
)
|
||||
return self._content
|
||||
|
||||
try:
|
||||
errors.APIError.raise_for_response(FakeResponse(status_code=400))
|
||||
except errors.ClientError as actual_error:
|
||||
assert actual_error.code == 400
|
||||
assert actual_error.message == 'error message'
|
||||
assert actual_error.status == 'INVALID_ARGUMENT'
|
||||
assert actual_error.details == {
|
||||
'error': {
|
||||
'code': 400,
|
||||
'message': 'error message',
|
||||
'status': 'INVALID_ARGUMENT',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_raise_for_response_server_error():
|
||||
class FakeResponse(httpx.Response):
|
||||
|
||||
def read(self) -> bytes:
|
||||
self._content = (
|
||||
b'{"error": {"code": 500, "message": "error message", "status":'
|
||||
b' "SERVER_INTERNAL ERROR"}}'
|
||||
)
|
||||
return self._content
|
||||
|
||||
try:
|
||||
errors.APIError.raise_for_response(FakeResponse(status_code=500))
|
||||
except errors.ServerError as actual_error:
|
||||
assert actual_error.code == 500
|
||||
assert actual_error.message == 'error message'
|
||||
assert actual_error.status == 'SERVER_INTERNAL ERROR'
|
||||
assert actual_error.details == {
|
||||
'error': {
|
||||
'code': 500,
|
||||
'message': 'error message',
|
||||
'status': 'SERVER_INTERNAL ERROR',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_api_error_is_picklable():
|
||||
pickled_error = pickle.loads(pickle.dumps(errors.APIError(1, {})))
|
||||
assert isinstance(pickled_error, errors.APIError)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raise_for_async_response_client_error():
|
||||
class FakeResponse(httpx.Response):
|
||||
|
||||
async def aread(self) -> bytes:
|
||||
self._content = (
|
||||
b'{"error": {"code": 400, "message": "error message", "status":'
|
||||
b' "INVALID_ARGUMENT"}}'
|
||||
)
|
||||
return self._content
|
||||
|
||||
try:
|
||||
await errors.APIError.raise_for_async_response(
|
||||
FakeResponse(status_code=400)
|
||||
)
|
||||
except errors.ClientError as actual_error:
|
||||
assert actual_error.code == 400
|
||||
assert actual_error.message == 'error message'
|
||||
assert actual_error.status == 'INVALID_ARGUMENT'
|
||||
assert actual_error.details == {
|
||||
'error': {
|
||||
'code': 400,
|
||||
'message': 'error message',
|
||||
'status': 'INVALID_ARGUMENT',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raise_for_async_response_server_error():
|
||||
class FakeResponse(httpx.Response):
|
||||
|
||||
async def aread(self) -> bytes:
|
||||
self._content = (
|
||||
b'{"error": {"code": 500, "message": "error message", "status":'
|
||||
b' "SERVER_INTERNAL ERROR"}}'
|
||||
)
|
||||
return self._content
|
||||
|
||||
try:
|
||||
await errors.APIError.raise_for_async_response(
|
||||
FakeResponse(status_code=500)
|
||||
)
|
||||
except errors.ServerError as actual_error:
|
||||
assert actual_error.code == 500
|
||||
assert actual_error.message == 'error message'
|
||||
assert actual_error.status == 'SERVER_INTERNAL ERROR'
|
||||
assert actual_error.details == {
|
||||
'error': {
|
||||
'code': 500,
|
||||
'message': 'error message',
|
||||
'status': 'SERVER_INTERNAL ERROR',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raise_for_async_response_code_exist_json_decoder_error():
|
||||
class FakeResponse(httpx.Response):
|
||||
|
||||
async def aread(self) -> bytes:
|
||||
self._content = b'{"data": {"key1": "value1", "key2"}'
|
||||
return self._content
|
||||
|
||||
try:
|
||||
await errors.APIError.raise_for_async_response(
|
||||
FakeResponse(
|
||||
status_code=503,
|
||||
extensions={'reason_phrase': b'Service Unavailable'},
|
||||
)
|
||||
)
|
||||
except errors.ServerError as actual_error:
|
||||
assert actual_error.code == 503
|
||||
assert actual_error.message == '{"data": {"key1": "value1", "key2"}'
|
||||
assert actual_error.status == 'Service Unavailable'
|
||||
assert actual_error.details == {
|
||||
'message': '{"data": {"key1": "value1", "key2"}',
|
||||
'status': 'Service Unavailable',
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
# 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 file search stores module."""
|
||||
@@ -0,0 +1,66 @@
|
||||
# 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 file_search_stores.create()."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
# All tests will be run for both Vertex and MLDev.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_display_name",
|
||||
parameters=types._CreateFileSearchStoreParameters(
|
||||
config=types.CreateFileSearchStoreConfig(
|
||||
display_name="My File Search Store"
|
||||
)
|
||||
),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_basic",
|
||||
parameters=types._CreateFileSearchStoreParameters(),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.usefixtures("mock_timestamped_unique_name"),
|
||||
pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method="file_search_stores.create",
|
||||
test_table=test_table,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_display_name(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file_search_store = await client.aio.file_search_stores.create(
|
||||
config=types.CreateFileSearchStoreConfig(
|
||||
display_name="My File Search Store"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_basic(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file_search_store = await client.aio.file_search_stores.create()
|
||||
@@ -0,0 +1,64 @@
|
||||
# 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 file_search_stores.delete()."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
# All tests will be run for both Vertex and MLDev.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_delete_with_name",
|
||||
parameters=types._DeleteFileSearchStoreParameters(
|
||||
name="fileSearchStores/acxjj7m366ln-aw6xyp94icll",
|
||||
),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_delete_with_name_and_force",
|
||||
parameters=types._DeleteFileSearchStoreParameters(
|
||||
name="fileSearchStores/7igesc9r2zw9-0mpxpsqubv7s",
|
||||
config=types.DeleteFileSearchStoreConfig(force=True),
|
||||
),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method="file_search_stores.delete",
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_delete(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
await client.aio.file_search_stores.delete(
|
||||
name="fileSearchStores/my-file-search-store-l65kcyel9lkz"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_force_delete(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
await client.aio.file_search_stores.delete(
|
||||
name="fileSearchStores/my-file-search-store-vjtrjw6re8oz",
|
||||
config=types.DeleteFileSearchStoreConfig(force=True),
|
||||
)
|
||||
@@ -0,0 +1,94 @@
|
||||
# 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 file_search_stores.get()."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
# A FileSearchStore name known to exist in the test environment.
|
||||
# Replace with a real one from your test setup.
|
||||
_EXISTING_FILE_SEARCH_STORE_NAME = "fileSearchStores/5en07ei3kojo-yo8sjqgvx2xf"
|
||||
_NON_EXISTENT_FILE_SEARCH_STORE_NAME = (
|
||||
"fileSearchStores/non-existent-file-search-store"
|
||||
)
|
||||
_INVALID_FILE_SEARCH_STORE_NAME = (
|
||||
"fileSearchStores/_invalid_file_search_store_name"
|
||||
)
|
||||
_NOT_A_FILE_SEARCH_STORE_NAME = "genai-test-file-search-store"
|
||||
|
||||
# All tests will be run for both Vertex and MLDev.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_get_success",
|
||||
parameters=types._GetFileSearchStoreParameters(
|
||||
name=_EXISTING_FILE_SEARCH_STORE_NAME
|
||||
),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_get_not_found",
|
||||
parameters=types._GetFileSearchStoreParameters(
|
||||
name=_NON_EXISTENT_FILE_SEARCH_STORE_NAME
|
||||
),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
# Expect an exception to be raised by the mock
|
||||
exception_if_mldev="PERMISSION_DENIED",
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_get_invalid_name",
|
||||
parameters=types._GetFileSearchStoreParameters(
|
||||
name=_INVALID_FILE_SEARCH_STORE_NAME
|
||||
),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
# Validation should catch this before the API call
|
||||
exception_if_mldev="INVALID_ARGUMENT",
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_get_not_a_file_search_store_name",
|
||||
parameters=types._GetFileSearchStoreParameters(
|
||||
name=_NOT_A_FILE_SEARCH_STORE_NAME
|
||||
),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
exception_if_mldev="Not Found",
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method="file_search_stores.get",
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_get(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
# This test relies on the mocking/replays set up by pytest_helper
|
||||
# For a success case:
|
||||
try:
|
||||
file_search_store = await client.aio.file_search_stores.get(
|
||||
name=_EXISTING_FILE_SEARCH_STORE_NAME
|
||||
)
|
||||
assert file_search_store is not None
|
||||
assert file_search_store.name == _EXISTING_FILE_SEARCH_STORE_NAME
|
||||
except Exception as e:
|
||||
# Depending on the mock setup, errors might be raised directly
|
||||
print(f"Async get failed: {e}")
|
||||
raise
|
||||
@@ -0,0 +1,112 @@
|
||||
# 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 file_search_stores.upload_to_file_search_store()."""
|
||||
|
||||
import io
|
||||
import pathlib
|
||||
import pytest
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
_FILE_SEARCH_STORE_NAME = 'fileSearchStores/my-store-37cbhu1nw16r'
|
||||
_FILE_NAME = 'files/mk4h34zkv33d'
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = []
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='filesearchstores.import_file',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
def test_import(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
operation = client.file_search_stores.import_file(
|
||||
file_search_store_name=_FILE_SEARCH_STORE_NAME, file_name=_FILE_NAME
|
||||
)
|
||||
|
||||
|
||||
def test_chunking(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
operation = client.file_search_stores.import_file(
|
||||
file_search_store_name=_FILE_SEARCH_STORE_NAME,
|
||||
file_name=_FILE_NAME,
|
||||
config=types.ImportFileConfig(
|
||||
chunking_config=types.ChunkingConfig(
|
||||
white_space_config=types.WhiteSpaceConfig(
|
||||
max_tokens_per_chunk=200,
|
||||
max_overlap_tokens=20,
|
||||
)
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_metadata(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
operation = client.file_search_stores.import_file(
|
||||
file_search_store_name=_FILE_SEARCH_STORE_NAME,
|
||||
file_name=_FILE_NAME,
|
||||
config=types.ImportFileConfig(
|
||||
custom_metadata=[
|
||||
types.CustomMetadata(key='year', numeric_value=2024),
|
||||
types.CustomMetadata(key='tag', string_value='story'),
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_import(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
operation = await client.aio.file_search_stores.import_file(
|
||||
file_search_store_name=_FILE_SEARCH_STORE_NAME,
|
||||
file_name=_FILE_NAME,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_chunking(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
operation = await client.aio.file_search_stores.import_file(
|
||||
file_search_store_name=_FILE_SEARCH_STORE_NAME,
|
||||
file_name=_FILE_NAME,
|
||||
config=types.ImportFileConfig(
|
||||
chunking_config=types.ChunkingConfig(
|
||||
white_space_config=types.WhiteSpaceConfig(
|
||||
max_tokens_per_chunk=200,
|
||||
max_overlap_tokens=20,
|
||||
)
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_metadata(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
operation = await client.aio.file_search_stores.import_file(
|
||||
file_search_store_name=_FILE_SEARCH_STORE_NAME,
|
||||
file_name=_FILE_NAME,
|
||||
config=types.ImportFileConfig(
|
||||
custom_metadata=[
|
||||
types.CustomMetadata(key='year', numeric_value=2024),
|
||||
types.CustomMetadata(key='tag', string_value='story'),
|
||||
],
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,57 @@
|
||||
# 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 file_search_stores.list()."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
# All tests will be run for both Vertex and MLDev.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_list_default",
|
||||
parameters=types._ListFileSearchStoresParameters(),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name="test_list_with_page_size",
|
||||
parameters=types._ListFileSearchStoresParameters(
|
||||
config=types.ListFileSearchStoresConfig(
|
||||
page_size=5,
|
||||
),
|
||||
),
|
||||
exception_if_vertex="only supported in the Gemini Developer client",
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method="file_search_stores.list",
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pager(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
# Iterate through the first page.
|
||||
async for file_search_store in await client.aio.file_search_stores.list(
|
||||
config={"page_size": 2}
|
||||
):
|
||||
assert isinstance(file_search_store, types.FileSearchStore)
|
||||
break # Only check one item
|
||||
@@ -0,0 +1,141 @@
|
||||
# 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 file_search_stores.upload_to_file_search_store()."""
|
||||
|
||||
import io
|
||||
import pathlib
|
||||
import pytest
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
_FILE_SEARCH_STORE_NAME = 'fileSearchStores/my-store-37cbhu1nw16r'
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = []
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='filesearchstores.upload_to_file_search_store',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
def test_file_path(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
operation = client.file_search_stores.upload_to_file_search_store(
|
||||
file_search_store_name=_FILE_SEARCH_STORE_NAME,
|
||||
file='tests/data/story.pdf',
|
||||
)
|
||||
|
||||
|
||||
def test_bytesio(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
with open('tests/data/story.pdf', 'rb') as f:
|
||||
with io.BytesIO(f.read()) as buffer:
|
||||
operation = client.file_search_stores.upload_to_file_search_store(
|
||||
file_search_store_name=_FILE_SEARCH_STORE_NAME,
|
||||
file=buffer,
|
||||
config=types.UploadToFileSearchStoreConfig(
|
||||
mime_type='application/pdf',
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_chunking(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
operation = client.file_search_stores.upload_to_file_search_store(
|
||||
file_search_store_name=_FILE_SEARCH_STORE_NAME,
|
||||
file='tests/data/story.pdf',
|
||||
config=types.UploadToFileSearchStoreConfig(
|
||||
chunking_config=types.ChunkingConfig(
|
||||
white_space_config=types.WhiteSpaceConfig(
|
||||
max_tokens_per_chunk=200,
|
||||
max_overlap_tokens=20,
|
||||
)
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_metadata(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
operation = client.file_search_stores.upload_to_file_search_store(
|
||||
file_search_store_name=_FILE_SEARCH_STORE_NAME,
|
||||
file='tests/data/story.pdf',
|
||||
config=types.UploadToFileSearchStoreConfig(
|
||||
custom_metadata=[
|
||||
types.CustomMetadata(key='year', numeric_value=2024),
|
||||
types.CustomMetadata(key='tag', string_value='story'),
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_file_path(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
operation = await client.aio.file_search_stores.upload_to_file_search_store(
|
||||
file_search_store_name=_FILE_SEARCH_STORE_NAME,
|
||||
file='tests/data/story.pdf',
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_bytesio(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
with open('tests/data/story.pdf', 'rb') as f:
|
||||
with io.BytesIO(f.read()) as buffer:
|
||||
operation = (
|
||||
await client.aio.file_search_stores.upload_to_file_search_store(
|
||||
file_search_store_name=_FILE_SEARCH_STORE_NAME,
|
||||
file=buffer,
|
||||
config=types.UploadToFileSearchStoreConfig(
|
||||
mime_type='application/pdf',
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_chunking(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
operation = await client.aio.file_search_stores.upload_to_file_search_store(
|
||||
file_search_store_name=_FILE_SEARCH_STORE_NAME,
|
||||
file='tests/data/story.pdf',
|
||||
config=types.UploadToFileSearchStoreConfig(
|
||||
chunking_config=types.ChunkingConfig(
|
||||
white_space_config=types.WhiteSpaceConfig(
|
||||
max_tokens_per_chunk=200,
|
||||
max_overlap_tokens=20,
|
||||
)
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_metadata(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
operation = await client.aio.file_search_stores.upload_to_file_search_store(
|
||||
file_search_store_name=_FILE_SEARCH_STORE_NAME,
|
||||
file='tests/data/story.pdf',
|
||||
config=types.UploadToFileSearchStoreConfig(
|
||||
custom_metadata=[
|
||||
types.CustomMetadata(key='year', numeric_value=2024),
|
||||
types.CustomMetadata(key='tag', string_value='story'),
|
||||
],
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,17 @@
|
||||
# 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 files module."""
|
||||
@@ -0,0 +1,46 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""Test files delete method."""
|
||||
|
||||
import pytest
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_delete',
|
||||
parameters=types._DeleteFileParameters(name='files/1g583ke2xdsn'),
|
||||
exception_if_vertex='only supported in the Gemini Developer client',
|
||||
skip_in_api_mode=(
|
||||
'The files have a TTL, they cannot be reliably retrieved for a long'
|
||||
' time.'
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='files.delete',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = await client.aio.files.get(name='files/n1gls7dyh90q')
|
||||
@@ -0,0 +1,85 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""Test files upload method."""
|
||||
|
||||
|
||||
import pathlib
|
||||
import pytest
|
||||
from ... import _transformers as t
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = []
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='t.t_file_name',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
pytest_plugins = ('pytest_asyncio',)
|
||||
|
||||
|
||||
def test_name_transform_name(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
for file in client.files.list():
|
||||
if file.download_uri is not None:
|
||||
break
|
||||
else:
|
||||
raise ValueError('No files found with a `download_uri`.')
|
||||
|
||||
file_id = file.name.split('/')[-1]
|
||||
video = types.Video(uri=file.download_uri)
|
||||
generated_video = types.GeneratedVideo(video=video)
|
||||
for f in [
|
||||
file,
|
||||
file_id,
|
||||
file.name,
|
||||
file.uri,
|
||||
file.download_uri,
|
||||
video,
|
||||
generated_video,
|
||||
]:
|
||||
name = t.t_file_name(f)
|
||||
assert name == file_id
|
||||
|
||||
|
||||
def test_basic_download(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
for file in client.files.list():
|
||||
if file.download_uri is not None:
|
||||
break
|
||||
else:
|
||||
raise ValueError('No files found with a `download_uri`.')
|
||||
|
||||
content = client.files.download(file=file)
|
||||
assert content[4:8] == b'ftyp'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_basic_download_async(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
async for file in await client.aio.files.list():
|
||||
if file.download_uri is not None:
|
||||
break
|
||||
else:
|
||||
raise ValueError('No files found with a `download_uri`.')
|
||||
|
||||
content = await client.aio.files.download(file=file)
|
||||
assert content[4:8] == b'ftyp'
|
||||
@@ -0,0 +1,46 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""Test files get method."""
|
||||
|
||||
import pytest
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_get',
|
||||
parameters=types._GetFileParameters(name='files/vjvu9fwk2qj8'),
|
||||
exception_if_vertex='only supported in the Gemini Developer client',
|
||||
skip_in_api_mode=(
|
||||
'The files have a TTL, they cannot be reliably retrieved for a long'
|
||||
' time.'
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='files.get',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = await client.aio.files.get(name='files/vjvu9fwk2qj8')
|
||||
@@ -0,0 +1,72 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""Test files list method."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_not_empty',
|
||||
exception_if_vertex='only supported in the Gemini Developer client',
|
||||
parameters=types._ListFilesParameters(
|
||||
config=types.ListFilesConfig(
|
||||
page_size=2,
|
||||
),
|
||||
),
|
||||
),
|
||||
]
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='files.list',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
def test_pager(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
files = client.files.list(config={'page_size': 2})
|
||||
assert 'content-type' in files.sdk_http_response.headers
|
||||
assert files.name == 'files'
|
||||
assert files.page_size == 2
|
||||
assert len(files) <= 2
|
||||
|
||||
# Iterate through all the pages. Then next_page() should raise an exception.
|
||||
for _ in files:
|
||||
pass
|
||||
with pytest.raises(IndexError, match='No more pages to fetch.'):
|
||||
files.next_page()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pager(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
files = await client.aio.files.list(config={'page_size': 2})
|
||||
|
||||
assert 'Content-Type' in files.sdk_http_response.headers
|
||||
assert files.name == 'files'
|
||||
assert files.page_size == 2
|
||||
assert len(files) <= 2
|
||||
|
||||
# Iterate through all the pages. Then next_page() should raise an exception.
|
||||
async for _ in files:
|
||||
pass
|
||||
with pytest.raises(IndexError, match='No more pages to fetch.'):
|
||||
await files.next_page()
|
||||
@@ -0,0 +1,272 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""Test files register method."""
|
||||
|
||||
import json
|
||||
from unittest import mock
|
||||
|
||||
from google.auth import credentials
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from ... import _api_client
|
||||
from ... import Client
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
|
||||
class FakeCredentials(credentials.Credentials):
|
||||
|
||||
def __init__(self, token="fake_token", expired=False, quota_project_id=None):
|
||||
super().__init__()
|
||||
self.token = token
|
||||
self._expired = expired
|
||||
self._quota_project_id = quota_project_id
|
||||
self.refresh_count = 0
|
||||
|
||||
@property
|
||||
def expired(self):
|
||||
return self._expired
|
||||
|
||||
@property
|
||||
def quota_project_id(self):
|
||||
return self._quota_project_id
|
||||
|
||||
def refresh(self, request):
|
||||
self.refresh_count += 1
|
||||
self.token = "refreshed_token"
|
||||
self._expired = False
|
||||
|
||||
|
||||
@mock.patch.object(_api_client.BaseApiClient, "_request_once", autospec=True)
|
||||
def test_simple_token(mock_request):
|
||||
client = Client(api_key="dummy_key")
|
||||
captured_request = None
|
||||
|
||||
def side_effect(self, http_request, stream=False):
|
||||
nonlocal captured_request
|
||||
captured_request = http_request
|
||||
return _api_client.HttpResponse(
|
||||
headers={},
|
||||
response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
|
||||
)
|
||||
|
||||
mock_request.side_effect = side_effect
|
||||
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
response = client.files.register_files(
|
||||
auth=FakeCredentials(token="test_token"),
|
||||
uris=["gs://test-bucket/test-file-1.txt"],
|
||||
)
|
||||
|
||||
assert len(response.files) == 1
|
||||
assert response.files[0].uri == "files/abc"
|
||||
assert captured_request.headers["authorization"] == "Bearer test_token"
|
||||
|
||||
|
||||
@mock.patch.object(_api_client.BaseApiClient, "_request_once", autospec=True)
|
||||
def test_token_refresh(mock_request):
|
||||
client = Client(api_key="dummy_key")
|
||||
captured_request = None
|
||||
|
||||
def side_effect(self, http_request, stream=False):
|
||||
nonlocal captured_request
|
||||
captured_request = http_request
|
||||
return _api_client.HttpResponse(
|
||||
headers={},
|
||||
response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
|
||||
)
|
||||
|
||||
mock_request.side_effect = side_effect
|
||||
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
creds = FakeCredentials(expired=True)
|
||||
response = client.files.register_files(
|
||||
auth=creds,
|
||||
uris=["gs://test-bucket/test-file-1.txt"],
|
||||
)
|
||||
assert creds.refresh_count == 1
|
||||
assert len(response.files) == 1
|
||||
assert response.files[0].uri == "files/abc"
|
||||
assert captured_request.headers["authorization"] == "Bearer refreshed_token"
|
||||
|
||||
|
||||
@mock.patch.object(_api_client.BaseApiClient, "_request_once", autospec=True)
|
||||
def test_quota_project(mock_request):
|
||||
client = Client(api_key="dummy_key")
|
||||
captured_request = None
|
||||
|
||||
def side_effect(self, http_request, stream=False):
|
||||
nonlocal captured_request
|
||||
captured_request = http_request
|
||||
return _api_client.HttpResponse(
|
||||
headers={},
|
||||
response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
|
||||
)
|
||||
|
||||
mock_request.side_effect = side_effect
|
||||
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
creds = FakeCredentials(quota_project_id="test_project")
|
||||
response = client.files.register_files(
|
||||
auth=creds,
|
||||
uris=["gs://test-bucket/test-file-1.txt"],
|
||||
)
|
||||
assert len(response.files) == 1
|
||||
assert response.files[0].uri == "files/abc"
|
||||
assert captured_request.headers["x-goog-user-project"] == "test_project"
|
||||
|
||||
|
||||
@mock.patch.object(_api_client.BaseApiClient, "_request_once", autospec=True)
|
||||
def test_multiple_uris(mock_request):
|
||||
client = Client(api_key="dummy_key")
|
||||
|
||||
def side_effect(self, http_request, stream=False):
|
||||
return _api_client.HttpResponse(
|
||||
headers={},
|
||||
response_stream=[
|
||||
json.dumps({"files": [{"uri": "files/abc"}, {"uri": "files/def"}]})
|
||||
],
|
||||
)
|
||||
|
||||
mock_request.side_effect = side_effect
|
||||
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
response = client.files.register_files(
|
||||
auth=FakeCredentials(),
|
||||
uris=[
|
||||
"gs://test-bucket/test-file-1.txt",
|
||||
"gs://test-bucket/test-file-2.txt",
|
||||
],
|
||||
)
|
||||
assert len(response.files) == 2
|
||||
assert response.files[0].uri == "files/abc"
|
||||
assert response.files[1].uri == "files/def"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@mock.patch.object(
|
||||
_api_client.BaseApiClient, "_async_request_once", autospec=True
|
||||
)
|
||||
async def test_async_single(mock_request):
|
||||
client = Client(api_key="dummy_key")
|
||||
|
||||
async def side_effect(self, http_request, stream=False):
|
||||
return _api_client.HttpResponse(
|
||||
headers={},
|
||||
response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
|
||||
)
|
||||
|
||||
mock_request.side_effect = side_effect
|
||||
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
response = await client.aio.files.register_files(
|
||||
auth=FakeCredentials(),
|
||||
uris=["gs://test-bucket/test-file-1.txt"],
|
||||
)
|
||||
|
||||
assert len(response.files) == 1
|
||||
assert response.files[0].uri == "files/abc"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@mock.patch.object(
|
||||
_api_client.BaseApiClient, "_async_request_once", autospec=True
|
||||
)
|
||||
async def test_async_token_refresh(mock_request):
|
||||
client = Client(api_key="dummy_key")
|
||||
captured_request = None
|
||||
|
||||
async def side_effect(self, http_request, stream=False):
|
||||
nonlocal captured_request
|
||||
captured_request = http_request
|
||||
return _api_client.HttpResponse(
|
||||
headers={},
|
||||
response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
|
||||
)
|
||||
|
||||
mock_request.side_effect = side_effect
|
||||
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
creds = FakeCredentials(expired=True)
|
||||
response = await client.aio.files.register_files(
|
||||
auth=creds,
|
||||
uris=["gs://test-bucket/test-file-1.txt"],
|
||||
)
|
||||
assert creds.refresh_count == 1
|
||||
assert len(response.files) == 1
|
||||
assert response.files[0].uri == "files/abc"
|
||||
assert captured_request.headers["authorization"] == "Bearer refreshed_token"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@mock.patch.object(
|
||||
_api_client.BaseApiClient, "_async_request_once", autospec=True
|
||||
)
|
||||
async def test_async_quota_project(mock_request):
|
||||
client = Client(api_key="dummy_key")
|
||||
captured_request = None
|
||||
|
||||
async def side_effect(self, http_request, stream=False):
|
||||
nonlocal captured_request
|
||||
captured_request = http_request
|
||||
return _api_client.HttpResponse(
|
||||
headers={},
|
||||
response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
|
||||
)
|
||||
|
||||
mock_request.side_effect = side_effect
|
||||
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
creds = FakeCredentials(quota_project_id="test_project")
|
||||
response = await client.aio.files.register_files(
|
||||
auth=creds,
|
||||
uris=["gs://test-bucket/test-file-1.txt"],
|
||||
)
|
||||
assert len(response.files) == 1
|
||||
assert response.files[0].uri == "files/abc"
|
||||
assert captured_request.headers["x-goog-user-project"] == "test_project"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@mock.patch.object(
|
||||
_api_client.BaseApiClient, "_async_request_once", autospec=True
|
||||
)
|
||||
async def test_async_multiple_uris(mock_request):
|
||||
client = Client(api_key="dummy_key")
|
||||
|
||||
async def side_effect(self, http_request, stream=False):
|
||||
return _api_client.HttpResponse(
|
||||
headers={},
|
||||
response_stream=[
|
||||
json.dumps({"files": [{"uri": "files/abc"}, {"uri": "files/def"}]})
|
||||
],
|
||||
)
|
||||
|
||||
mock_request.side_effect = side_effect
|
||||
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
response = await client.aio.files.register_files(
|
||||
auth=FakeCredentials(),
|
||||
uris=[
|
||||
"gs://test-bucket/test-file-1.txt",
|
||||
"gs://test-bucket/test-file-2.txt",
|
||||
],
|
||||
)
|
||||
assert len(response.files) == 2
|
||||
assert response.files[0].uri == "files/abc"
|
||||
assert response.files[1].uri == "files/def"
|
||||
@@ -0,0 +1,70 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""Test files get method."""
|
||||
|
||||
import pytest
|
||||
from ... import types
|
||||
from ... import Client
|
||||
from ... import _api_client
|
||||
from .. import pytest_helper
|
||||
import google.auth
|
||||
|
||||
|
||||
# $ gcloud config set project vertex-sdk-dev
|
||||
# $ gcloud auth application-default login --no-launch-browser --scopes="https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/devstorage.read_only"
|
||||
def get_headers():
|
||||
try:
|
||||
credentials, _ = google.auth.default()
|
||||
token = _api_client.get_token_from_credentials(None, credentials)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {token}",}
|
||||
if credentials.quota_project_id:
|
||||
headers["x-goog-user-project"] = credentials.quota_project_id
|
||||
except google.auth.exceptions.DefaultCredentialsError:
|
||||
# So this can run in replay mode without credentials.
|
||||
headers = {}
|
||||
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_register',
|
||||
parameters=types._InternalRegisterFilesParameters(uris=['gs://unified-genai-dev/image.jpg']),
|
||||
exception_if_vertex='only supported in the Gemini Developer client',
|
||||
skip_in_api_mode=(
|
||||
'The files have a TTL, they cannot be reliably retrieved for a long'
|
||||
' time.'
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='files._register_files',
|
||||
test_table=test_table,
|
||||
http_options={
|
||||
'headers': get_headers(),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
files = await client.aio.files._register_files(uris=['gs://unified-genai-dev/image.jpg'])
|
||||
assert files.files
|
||||
assert files.files[0].mime_type == 'image/jpeg'
|
||||
@@ -0,0 +1,255 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""Test files upload method."""
|
||||
|
||||
|
||||
import io
|
||||
import pathlib
|
||||
import pytest
|
||||
from ... import types
|
||||
from ... import errors
|
||||
from .. import pytest_helper
|
||||
|
||||
# Upload method is not pydantic.
|
||||
test_table: list[pytest_helper.TestTableItem] = []
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='files.upload',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
def test_image_png_upload(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = client.files.upload(file='tests/data/google.png')
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
def test_image_png_upload_with_path(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
p = pathlib.Path('tests/data/google.png')
|
||||
file = client.files.upload(
|
||||
file=p,
|
||||
config=types.UploadFileConfig(display_name='test_image_png_path'),
|
||||
)
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
def test_image_png_upload_with_bytesio(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
with open('tests/data/google.png', 'rb') as f:
|
||||
with io.BytesIO(f.read()) as buffer:
|
||||
file = client.files.upload(
|
||||
file=buffer,
|
||||
config=types.UploadFileConfig(mime_type='image/png'),
|
||||
)
|
||||
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
def test_image_png_upload_with_fd(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
with open('tests/data/google.png', 'rb') as f:
|
||||
file = client.files.upload(
|
||||
file=f,
|
||||
config=types.UploadFileConfig(mime_type='image/png'),
|
||||
)
|
||||
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
def test_image_png_upload_with_config(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = client.files.upload(
|
||||
file='tests/data/google.png',
|
||||
config=types.UploadFileConfig(display_name='test_image_png'),
|
||||
)
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
def test_image_png_upload_with_config_dict(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = client.files.upload(
|
||||
file='tests/data/google.png', config={'display_name': 'test_image_png'}
|
||||
)
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
def test_image_jpg_upload(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = client.files.upload(file='tests/data/google.jpg')
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
def test_image_jpg_upload_with_config(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = client.files.upload(
|
||||
file='tests/data/google.jpg',
|
||||
config=types.UploadFileConfig(display_name='test_image_jpg'),
|
||||
)
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
def test_image_jpg_upload_with_config_dict(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = client.files.upload(
|
||||
file='tests/data/google.jpg', config={'display_name': 'test_image_jpg'}
|
||||
)
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
def test_application_pdf_file_upload(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = client.files.upload(file='tests/data/story.pdf')
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
def test_application_pdf_upload_with_config(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = client.files.upload(
|
||||
file='tests/data/story.pdf',
|
||||
config=types.UploadFileConfig(display_name='test_application_pdf'),
|
||||
)
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
def test_application_pdf_upload_with_config_dict(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = client.files.upload(
|
||||
file='tests/data/story.pdf',
|
||||
config={'display_name': 'test_application_pdf'},
|
||||
)
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
def test_video_mp4_file_upload(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = client.files.upload(file='tests/data/animal.mp4')
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
def test_video_mp4_upload_with_config(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = client.files.upload(
|
||||
file='tests/data/animal.mp4',
|
||||
config=types.UploadFileConfig(display_name='test_video_mp4'),
|
||||
)
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
def test_video_mp4_upload_with_config_dict(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = client.files.upload(
|
||||
file='tests/data/animal.mp4', config={'display_name': 'test_video_mp4'}
|
||||
)
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
def test_audio_m4a_file_upload(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = client.files.upload(
|
||||
file='tests/data/pixel.m4a',
|
||||
config=types.UploadFileConfig(mime_type='audio/mp4'),
|
||||
)
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
def test_audio_m4a_upload_with_config(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = client.files.upload(
|
||||
file='tests/data/pixel.m4a',
|
||||
config=types.UploadFileConfig(
|
||||
display_name='test_audio_m4a', mime_type='audio/mp4'
|
||||
),
|
||||
)
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
def test_audio_m4a_upload_with_config_dict(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = client.files.upload(
|
||||
file='tests/data/pixel.m4a',
|
||||
config={'display_name': 'test_audio_m4a', 'mime_type': 'audio/mp4'},
|
||||
)
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
def test_bad_mime_type(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
with pytest.raises(errors.APIError, match="Unsupported MIME"):
|
||||
file = client.files.upload(
|
||||
file=io.BytesIO(b'test'),
|
||||
config={'mime_type': 'bad/mime_type'},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_upload_async(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = await client.aio.files.upload(file='tests/data/google.png')
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_upload_with_config_async(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = await client.aio.files.upload(
|
||||
file='tests/data/google.png',
|
||||
config=types.UploadFileConfig(display_name='test_image'),
|
||||
)
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_upload_with_config_dict_async(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = await client.aio.files.upload(
|
||||
file='tests/data/google.png',
|
||||
config={
|
||||
'display_name': 'test_image',
|
||||
'http_options': {'timeout': '8000'},
|
||||
},
|
||||
)
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_upload_with_bytesio_async(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
with open('tests/data/google.png', 'rb') as f:
|
||||
buffer = io.BytesIO(f.read())
|
||||
file = await client.aio.files.upload(
|
||||
file=buffer,
|
||||
config=types.UploadFileConfig(
|
||||
mime_type='image/png'),
|
||||
)
|
||||
assert file.name.startswith('files/')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_path_upload_async(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
try:
|
||||
await client.aio.files.upload(file='unknown_path')
|
||||
except FileNotFoundError as e:
|
||||
assert 'is not a valid file path' in str(e)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bad_mime_type_async(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
with pytest.raises(errors.APIError, match="Unsupported MIME"):
|
||||
file = await client.aio.files.upload(
|
||||
file=io.BytesIO(b'test'),
|
||||
config={'mime_type': 'bad/mime_type'},
|
||||
)
|
||||
@@ -0,0 +1,28 @@
|
||||
# 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 os
|
||||
import pytest
|
||||
|
||||
IS_NOT_GITHUB_ACTIONS = os.getenv('GITHUB_ACTIONS') != 'true'
|
||||
|
||||
|
||||
@pytest.mark.skipif(IS_NOT_GITHUB_ACTIONS,
|
||||
reason='This test is only run on GitHub Actions.')
|
||||
def test_library_can_be_imported_without_optional_dependencies():
|
||||
"""Tests that the library can be imported without optional dependencies.
|
||||
"""
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
@@ -0,0 +1,476 @@
|
||||
# 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 Interactions API."""
|
||||
|
||||
from ... import Client
|
||||
from ... import _base_url
|
||||
from unittest import mock
|
||||
import pytest
|
||||
from httpx import Request, Response
|
||||
from ..._api_client import AsyncHttpxClient, BaseApiClient
|
||||
from httpx import Client as HTTPClient
|
||||
import os
|
||||
|
||||
ENV_VARS = [
|
||||
"GOOGLE_CLOUD_PROJECT",
|
||||
"GEMINI_API_KEY",
|
||||
"GOOGLE_API_KEY",
|
||||
"GOOGLE_CLOUD_LOCATION",
|
||||
]
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_env_vars(monkeypatch):
|
||||
for var in ENV_VARS:
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
|
||||
def test_interactions_gemini_url(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
|
||||
client = Client()
|
||||
|
||||
|
||||
with mock.patch.object(HTTPClient, "send") as mock_send:
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''))
|
||||
client.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url).endswith('/v1beta/interactions')
|
||||
assert request.headers['x-goog-api-key'] == 'test-api-key'
|
||||
|
||||
|
||||
def test_interactions_gemini_no_vertex_auth(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
|
||||
client = Client()
|
||||
|
||||
with (
|
||||
mock.patch.object(BaseApiClient, "_access_token") as mock_access_token,
|
||||
mock.patch.object(HTTPClient, "send") as mock_send,
|
||||
):
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''))
|
||||
client.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
)
|
||||
mock_access_token.assert_not_called()
|
||||
|
||||
def test_interactions_gemini_retry(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
|
||||
client = Client()
|
||||
client._api_client.max_retries = 2
|
||||
|
||||
with mock.patch.object(HTTPClient, "send") as mock_send:
|
||||
mock_send.side_effect = [
|
||||
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1"}),
|
||||
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1"}),
|
||||
Response(200, request=Request('POST', '')),
|
||||
]
|
||||
client.interactions.create(model='gemini-1.5-flash', input='Hello')
|
||||
assert mock_send.call_count == 3
|
||||
|
||||
def test_interactions_gemini_extra_headers(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
|
||||
client = Client()
|
||||
|
||||
with mock.patch.object(HTTPClient, "send") as mock_send:
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''))
|
||||
client.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
extra_headers={'X-Custom-Header': 'TestValue'}
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert request.headers['x-custom-header'] == 'TestValue'
|
||||
assert request.headers['x-goog-api-key'] == 'test-api-key'
|
||||
|
||||
|
||||
def test_interactions_vertex_auth_header():
|
||||
from ..._api_client import BaseApiClient
|
||||
from ..._interactions._base_client import SyncAPIClient
|
||||
from httpx import Client as HTTPClient
|
||||
|
||||
creds = mock.Mock()
|
||||
creds.quota_project_id = "test-quota-project"
|
||||
client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
|
||||
|
||||
with (
|
||||
mock.patch.object(
|
||||
BaseApiClient, "_access_token", return_value='fake-vertex-token'
|
||||
) as mock_access_token,
|
||||
mock.patch.object(
|
||||
HTTPClient, "send",
|
||||
return_value=mock.Mock(),
|
||||
) as mock_send,
|
||||
):
|
||||
|
||||
response = client.interactions.create(
|
||||
model='gemini-2.5-flash',
|
||||
input='What is the largest planet in our solar system?',
|
||||
)
|
||||
|
||||
mock_send.assert_called_once()
|
||||
mock_access_token.assert_called_once()
|
||||
args, kwargs = mock_send.call_args
|
||||
headers = args[0].headers
|
||||
assert any(
|
||||
key == "authorization" and value == 'Bearer fake-vertex-token'
|
||||
for key, value in headers.items())
|
||||
assert any(
|
||||
key == "x-goog-user-project" and value == 'test-quota-project'
|
||||
for key, value in headers.items())
|
||||
|
||||
def test_interactions_vertex_key_no_auth_header():
|
||||
from ..._api_client import BaseApiClient
|
||||
from httpx import Client as HTTPClient
|
||||
|
||||
creds = mock.Mock()
|
||||
client = Client(vertexai=True, api_key='test-api-key')
|
||||
|
||||
with (
|
||||
mock.patch.object(
|
||||
BaseApiClient, "_access_token", return_value='fake-vertex-token'
|
||||
) as mock_access_token,
|
||||
mock.patch.object(
|
||||
HTTPClient, "send",
|
||||
return_value=mock.Mock(),
|
||||
) as mock_send,
|
||||
):
|
||||
|
||||
response = client.interactions.create(
|
||||
model='gemini-2.5-flash',
|
||||
input='What is the largest planet in our solar system?',
|
||||
)
|
||||
|
||||
mock_send.assert_called_once()
|
||||
mock_access_token.assert_not_called()
|
||||
args, kwargs = mock_send.call_args
|
||||
headers = args[0].headers
|
||||
assert any(
|
||||
key == "x-goog-api-key" and value == 'test-api-key'
|
||||
for key, value in headers.items())
|
||||
|
||||
def test_interactions_vertex_url():
|
||||
creds = mock.Mock()
|
||||
creds.quota_project_id = "test-quota-project"
|
||||
client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
|
||||
|
||||
with mock.patch("httpx.Client.send") as mock_send:
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''))
|
||||
client.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url) == 'https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/interactions'
|
||||
|
||||
def test_interactions_vertex_auth_refresh_on_retry():
|
||||
from ..._api_client import BaseApiClient
|
||||
from httpx import Client as HTTPClient
|
||||
|
||||
creds = mock.Mock()
|
||||
creds.quota_project_id = "test-quota-project"
|
||||
client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
|
||||
client._api_client.max_retries = 2
|
||||
|
||||
token_values = ['token1', 'token2', 'token3']
|
||||
token_iter = iter(token_values)
|
||||
|
||||
def get_token():
|
||||
return next(token_iter)
|
||||
|
||||
with (
|
||||
mock.patch.object(BaseApiClient, "_access_token", side_effect=get_token) as mock_access_token,
|
||||
mock.patch.object(HTTPClient, "send") as mock_send,
|
||||
):
|
||||
mock_send.side_effect = [
|
||||
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1"}),
|
||||
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1"}),
|
||||
Response(200, request=Request('POST', '')),
|
||||
]
|
||||
|
||||
client.interactions.create(model='gemini-1.5-flash', input='Hello')
|
||||
|
||||
assert mock_access_token.call_count == 3
|
||||
assert mock_send.call_count == 3
|
||||
# Check headers of each call
|
||||
for i in range(3):
|
||||
headers = mock_send.call_args_list[i][0][0].headers
|
||||
assert headers['authorization'] == f'Bearer {token_values[i]}'
|
||||
|
||||
|
||||
def test_interactions_vertex_extra_headers_override():
|
||||
from ..._api_client import BaseApiClient
|
||||
from httpx import Client as HTTPClient
|
||||
|
||||
creds = mock.Mock()
|
||||
creds.quota_project_id = "test-quota-project"
|
||||
client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
|
||||
|
||||
with (
|
||||
mock.patch.object(BaseApiClient, "_access_token", return_value='default-token') as mock_access_token,
|
||||
mock.patch.object(HTTPClient, "send") as mock_send,
|
||||
):
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''))
|
||||
|
||||
# Override Authorization
|
||||
client.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
extra_headers={'Authorization': 'Bearer manual-token'}
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
headers = mock_send.call_args[0][0].headers
|
||||
assert headers['authorization'] == 'Bearer manual-token'
|
||||
mock_access_token.assert_not_called() # Should not fetch default token
|
||||
|
||||
mock_send.reset_mock()
|
||||
mock_access_token.reset_mock()
|
||||
|
||||
# Provide API Key
|
||||
client.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
extra_headers={'x-goog-api-key': 'manual-key'}
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
headers = mock_send.call_args[0][0].headers
|
||||
assert headers['x-goog-api-key'] == 'manual-key'
|
||||
assert 'authorization' not in headers
|
||||
mock_access_token.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_gemini_url(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
|
||||
client = Client()
|
||||
|
||||
with mock.patch.object(AsyncHttpxClient, "send") as mock_send:
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''))
|
||||
await client.aio.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url).endswith('/v1beta/interactions')
|
||||
assert request.headers['x-goog-api-key'] == 'test-api-key'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_gemini_no_vertex_auth(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
|
||||
client = Client()
|
||||
|
||||
with (
|
||||
mock.patch.object(BaseApiClient, "_async_access_token") as mock_access_token,
|
||||
mock.patch.object(AsyncHttpxClient, "send") as mock_send,
|
||||
):
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''))
|
||||
await client.aio.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
)
|
||||
mock_access_token.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_gemini_retry(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
|
||||
client = Client()
|
||||
client.aio._api_client.max_retries = 2
|
||||
|
||||
with mock.patch.object(AsyncHttpxClient, "send") as mock_send:
|
||||
mock_send.side_effect = [
|
||||
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1"}),
|
||||
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1"}),
|
||||
Response(200, request=Request('POST', '')),
|
||||
]
|
||||
await client.aio.interactions.create(model='gemini-1.5-flash', input='Hello')
|
||||
assert mock_send.call_count == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_gemini_extra_headers(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
|
||||
client = Client()
|
||||
|
||||
with mock.patch.object(AsyncHttpxClient, "send") as mock_send:
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''))
|
||||
await client.aio.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
extra_headers={'X-Custom-Header': 'TestValue'}
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert request.headers['x-custom-header'] == 'TestValue'
|
||||
assert request.headers['x-goog-api-key'] == 'test-api-key'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_vertex_auth_header():
|
||||
from ..._api_client import BaseApiClient
|
||||
from ..._interactions._base_client import SyncAPIClient
|
||||
from ..._api_client import AsyncHttpxClient
|
||||
|
||||
creds = mock.Mock()
|
||||
creds.quota_project_id = "test-quota-project"
|
||||
client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
|
||||
|
||||
with (
|
||||
mock.patch.object(
|
||||
BaseApiClient, "_async_access_token", return_value='fake-vertex-token'
|
||||
) as mock_access_token,
|
||||
mock.patch.object(
|
||||
AsyncHttpxClient, "send",
|
||||
return_value=mock.Mock(),
|
||||
) as mock_send,
|
||||
):
|
||||
|
||||
response = await client.aio.interactions.create(
|
||||
model='gemini-2.5-flash',
|
||||
input='What is the largest planet in our solar system?',
|
||||
)
|
||||
|
||||
mock_send.assert_called_once()
|
||||
mock_access_token.assert_called_once()
|
||||
args, kwargs = mock_send.call_args
|
||||
headers = args[0].headers
|
||||
assert any(
|
||||
key == "authorization" and value == 'Bearer fake-vertex-token'
|
||||
for key, value in headers.items())
|
||||
assert any(
|
||||
key == "x-goog-user-project" and value == 'test-quota-project'
|
||||
for key, value in headers.items())
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_vertex_key_no_auth_header():
|
||||
from ..._api_client import BaseApiClient
|
||||
client = Client(vertexai=True, api_key='test-api-key')
|
||||
|
||||
with (
|
||||
mock.patch.object(
|
||||
BaseApiClient, "_async_access_token", return_value='fake-vertex-token'
|
||||
) as mock_access_token,
|
||||
mock.patch.object(
|
||||
AsyncHttpxClient, "send",
|
||||
return_value=mock.Mock(),
|
||||
) as mock_send,
|
||||
):
|
||||
|
||||
response = await client.aio.interactions.create(
|
||||
model='gemini-2.5-flash',
|
||||
input='What is the largest planet in our solar system?',
|
||||
)
|
||||
|
||||
mock_send.assert_called_once()
|
||||
mock_access_token.assert_not_called()
|
||||
args, kwargs = mock_send.call_args
|
||||
headers = args[0].headers
|
||||
assert any(
|
||||
key == "x-goog-api-key" and value == 'test-api-key'
|
||||
for key, value in headers.items())
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_vertex_url():
|
||||
from ..._api_client import AsyncHttpxClient
|
||||
creds = mock.Mock()
|
||||
creds.quota_project_id = "test-quota-project"
|
||||
client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
|
||||
|
||||
with mock.patch.object(AsyncHttpxClient, "send") as mock_send:
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''))
|
||||
await client.aio.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url) == 'https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/interactions'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_vertex_auth_refresh_on_retry():
|
||||
from ..._api_client import BaseApiClient
|
||||
from ..._api_client import AsyncHttpxClient
|
||||
|
||||
creds = mock.Mock()
|
||||
creds.quota_project_id = "test-quota-project"
|
||||
client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
|
||||
client.aio._api_client.max_retries = 2
|
||||
|
||||
token_values = ['token1', 'token2', 'token3']
|
||||
token_iter = iter(token_values)
|
||||
|
||||
async def get_token():
|
||||
return next(token_iter)
|
||||
|
||||
with (
|
||||
mock.patch.object(BaseApiClient, "_async_access_token", side_effect=get_token) as mock_access_token,
|
||||
mock.patch.object(AsyncHttpxClient, "send") as mock_send,
|
||||
):
|
||||
mock_send.side_effect = [
|
||||
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1"}),
|
||||
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1"}),
|
||||
Response(200, request=Request('POST', '')),
|
||||
]
|
||||
|
||||
await client.aio.interactions.create(model='gemini-1.5-flash', input='Hello')
|
||||
|
||||
assert mock_access_token.call_count == 3
|
||||
assert mock_send.call_count == 3
|
||||
for i in range(3):
|
||||
headers = mock_send.call_args_list[i][0][0].headers
|
||||
assert headers['authorization'] == f'Bearer {token_values[i]}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_vertex_extra_headers_override():
|
||||
from ..._api_client import BaseApiClient
|
||||
from ..._api_client import AsyncHttpxClient
|
||||
|
||||
creds = mock.Mock()
|
||||
creds.quota_project_id = "test-quota-project"
|
||||
client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
|
||||
|
||||
with (
|
||||
mock.patch.object(BaseApiClient, "_async_access_token", return_value='default-token') as mock_access_token,
|
||||
mock.patch.object(AsyncHttpxClient, "send") as mock_send,
|
||||
):
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''))
|
||||
|
||||
# Override Authorization
|
||||
await client.aio.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
extra_headers={'Authorization': 'Bearer manual-token'}
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
headers = mock_send.call_args[0][0].headers
|
||||
assert headers['authorization'] == 'Bearer manual-token'
|
||||
mock_access_token.assert_not_called()
|
||||
|
||||
mock_send.reset_mock()
|
||||
mock_access_token.reset_mock()
|
||||
|
||||
# Provide API Key
|
||||
await client.aio.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
extra_headers={'x-goog-api-key': 'manual-key'}
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
headers = mock_send.call_args[0][0].headers
|
||||
assert headers['x-goog-api-key'] == 'manual-key'
|
||||
assert 'authorization' not in headers
|
||||
mock_access_token.assert_not_called()
|
||||
@@ -0,0 +1,84 @@
|
||||
# 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 client as client_lib
|
||||
|
||||
pytest_plugins = ("pytest_asyncio",)
|
||||
|
||||
|
||||
def test_client_future_warning():
|
||||
with mock.patch.object(
|
||||
client_lib, "_interactions_experimental_warned", new=False
|
||||
):
|
||||
client = client_lib.Client(
|
||||
api_key="placeholder",
|
||||
http_options={
|
||||
"api_version": "v1alpha",
|
||||
}
|
||||
)
|
||||
with pytest.warns(
|
||||
UserWarning, match="Interactions.*experimental"
|
||||
):
|
||||
_ = client.interactions
|
||||
|
||||
|
||||
def test_client_timeout():
|
||||
with mock.patch.object(
|
||||
client_lib, "GeminiNextGenAPIClient", spec_set=True
|
||||
) as mock_nextgen_client:
|
||||
|
||||
client = client_lib.Client(
|
||||
api_key="placeholder",
|
||||
http_options={"api_version": "v1alpha", "timeout": 5000},
|
||||
)
|
||||
|
||||
_ = client.interactions
|
||||
|
||||
mock_nextgen_client.assert_called_once_with(
|
||||
base_url=mock.ANY,
|
||||
api_key="placeholder",
|
||||
api_version="v1alpha",
|
||||
default_headers=mock.ANY,
|
||||
http_client=mock.ANY,
|
||||
timeout=5.0,
|
||||
max_retries=mock.ANY,
|
||||
client_adapter=mock.ANY,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_client_timeout():
|
||||
with mock.patch.object(
|
||||
client_lib, "AsyncGeminiNextGenAPIClient", spec_set=True
|
||||
) as mock_nextgen_client:
|
||||
|
||||
client = client_lib.Client(
|
||||
api_key="placeholder",
|
||||
http_options={"api_version": "v1alpha", "timeout": 5000},
|
||||
)
|
||||
|
||||
_ = client.aio.interactions
|
||||
|
||||
mock_nextgen_client.assert_called_once_with(
|
||||
base_url=mock.ANY,
|
||||
api_key="placeholder",
|
||||
api_version="v1alpha",
|
||||
default_headers=mock.ANY,
|
||||
http_client=mock.ANY,
|
||||
timeout=5.0,
|
||||
max_retries=mock.ANY,
|
||||
client_adapter=mock.ANY,
|
||||
)
|
||||
@@ -0,0 +1,105 @@
|
||||
|
||||
# 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 Interactions API URL paths."""
|
||||
|
||||
from unittest import mock
|
||||
import pytest
|
||||
from httpx import Request, Response
|
||||
from ..._api_client import AsyncHttpxClient
|
||||
from httpx import Client as HTTPClient
|
||||
from .. import pytest_helper
|
||||
import google.auth
|
||||
|
||||
@mock.patch.object(google.auth, "default", autospec=True)
|
||||
def test_interactions_paths(mock_auth_default, client):
|
||||
interaction_id = "test-interaction-id"
|
||||
|
||||
mock_creds = mock.Mock()
|
||||
mock_creds.token = "test-token"
|
||||
mock_creds.expired = False
|
||||
mock_creds.quota_project_id = "test-quota-project"
|
||||
mock_auth_default.return_value = (mock_creds, "test-project")
|
||||
|
||||
if client._api_client.vertexai:
|
||||
expected_base_url = f'https://{client._api_client.location}-aiplatform.googleapis.com/v1beta1/projects/{client._api_client.project}/locations/{client._api_client.location}'
|
||||
else:
|
||||
expected_base_url = "https://generativelanguage.googleapis.com/v1beta"
|
||||
|
||||
with mock.patch.object(HTTPClient, "send") as mock_send:
|
||||
mock_send.return_value = Response(200, request=Request('GET', ''))
|
||||
client.interactions.get(id=interaction_id)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}'
|
||||
|
||||
mock_send.reset_mock()
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''))
|
||||
client.interactions.cancel(id=interaction_id)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}/cancel'
|
||||
|
||||
mock_send.reset_mock()
|
||||
mock_send.return_value = Response(200, request=Request('DELETE', ''))
|
||||
client.interactions.delete(id=interaction_id)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@mock.patch.object(google.auth, "default", autospec=True)
|
||||
async def test_async_interactions_paths(mock_auth_default, client):
|
||||
interaction_id = "test-interaction-id"
|
||||
|
||||
mock_creds = mock.Mock()
|
||||
mock_creds.token = "test-token"
|
||||
mock_creds.expired = False
|
||||
mock_creds.quota_project_id = "test-quota-project"
|
||||
mock_auth_default.return_value = (mock_creds, "test-project")
|
||||
|
||||
if client._api_client.vertexai:
|
||||
expected_base_url = f'https://{client._api_client.location}-aiplatform.googleapis.com/v1beta1/projects/{client._api_client.project}/locations/{client._api_client.location}'
|
||||
else:
|
||||
expected_base_url = "https://generativelanguage.googleapis.com/v1beta"
|
||||
|
||||
with mock.patch.object(AsyncHttpxClient, "send") as mock_send:
|
||||
mock_send.return_value = Response(200, request=Request('GET', ''))
|
||||
await client.aio.interactions.get(id=interaction_id)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}'
|
||||
|
||||
mock_send.reset_mock()
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''))
|
||||
await client.aio.interactions.cancel(id=interaction_id)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}/cancel'
|
||||
|
||||
mock_send.reset_mock()
|
||||
mock_send.return_value = Response(200, request=Request('DELETE', ''))
|
||||
await client.aio.interactions.delete(id=interaction_id)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}'
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_table=[],
|
||||
)
|
||||
@@ -0,0 +1,16 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,362 @@
|
||||
# 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 live_music.py."""
|
||||
import contextlib
|
||||
import json
|
||||
from typing import AsyncIterator
|
||||
from unittest import mock
|
||||
from unittest.mock import AsyncMock
|
||||
from unittest.mock import Mock
|
||||
from unittest.mock import patch
|
||||
import warnings
|
||||
|
||||
from google.oauth2.credentials import Credentials
|
||||
import pytest
|
||||
from websockets import client
|
||||
|
||||
from ... import _api_client as api_client
|
||||
from ... import _common
|
||||
from ... import Client
|
||||
from ... import client as gl_client
|
||||
from ... import live
|
||||
from ... import live_music
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
try:
|
||||
import aiohttp
|
||||
AIOHTTP_NOT_INSTALLED = False
|
||||
except ImportError:
|
||||
AIOHTTP_NOT_INSTALLED = True
|
||||
aiohttp = mock.MagicMock()
|
||||
|
||||
|
||||
requires_aiohttp = pytest.mark.skipif(
|
||||
AIOHTTP_NOT_INSTALLED, reason="aiohttp is not installed, skipping test."
|
||||
)
|
||||
|
||||
|
||||
def mock_api_client(vertexai=False, credentials=None):
|
||||
api_client = mock.MagicMock(spec=gl_client.BaseApiClient)
|
||||
if not vertexai:
|
||||
api_client.api_key = 'TEST_API_KEY'
|
||||
api_client.location = None
|
||||
api_client.project = None
|
||||
else:
|
||||
api_client.api_key = None
|
||||
api_client.location = 'us-central1'
|
||||
api_client.project = 'test_project'
|
||||
|
||||
api_client._host = lambda: 'test_host'
|
||||
api_client._credentials = credentials
|
||||
api_client._http_options = types.HttpOptions.model_validate(
|
||||
{'headers': {}}
|
||||
) # Ensure headers exist
|
||||
api_client.vertexai = vertexai
|
||||
api_client._api_client = api_client
|
||||
return api_client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_websocket():
|
||||
websocket = AsyncMock(spec=client.ClientConnection)
|
||||
websocket.send = AsyncMock()
|
||||
websocket.recv = AsyncMock(
|
||||
return_value=b"""{
|
||||
"serverContent": {
|
||||
"audioChunks": [
|
||||
{
|
||||
"data": "Z2VsYmFuYW5h",
|
||||
"mimeType": "audio/l16;rate=48000;channels=2",
|
||||
"sourceMetadata": {
|
||||
"clientContent": {
|
||||
"weightedPrompts": [
|
||||
{
|
||||
"text": "Jazz",
|
||||
"weight": 1
|
||||
}
|
||||
]
|
||||
},
|
||||
"musicGenerationConfig": {
|
||||
"seed": -957124937,
|
||||
"bpm": 140,
|
||||
"scale": "A_FLAT_MAJOR_F_MINOR"
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}"""
|
||||
) # Default response
|
||||
websocket.close = AsyncMock()
|
||||
return websocket
|
||||
|
||||
|
||||
async def get_connect_message(api_client, model):
|
||||
mock_ws = AsyncMock()
|
||||
mock_ws.send = AsyncMock()
|
||||
mock_ws.recv = AsyncMock(return_value=b'some response')
|
||||
|
||||
mock_google_auth_default = Mock(return_value=(None, None))
|
||||
mock_creds = Mock(token='test_token')
|
||||
mock_google_auth_default.return_value = (mock_creds, None)
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def mock_connect(uri, additional_headers=None):
|
||||
yield mock_ws
|
||||
|
||||
@patch('google.auth.default', new=mock_google_auth_default)
|
||||
@patch.object(live_music, 'connect', new=mock_connect)
|
||||
async def _test_connect():
|
||||
live_module = live.AsyncLive(api_client)
|
||||
async with live_module.music.connect(
|
||||
model=model,
|
||||
):
|
||||
pass
|
||||
|
||||
mock_ws.send.assert_called_once()
|
||||
return json.loads(mock_ws.send.call_args[0][0])
|
||||
|
||||
return await _test_connect()
|
||||
|
||||
|
||||
def test_mldev_from_env(monkeypatch):
|
||||
api_key = 'google_api_key'
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', api_key)
|
||||
|
||||
client = Client()
|
||||
|
||||
assert not client.aio.live.music._api_client.vertexai
|
||||
assert client.aio.live.music._api_client.api_key == api_key
|
||||
assert isinstance(client.aio.live._api_client, api_client.BaseApiClient)
|
||||
|
||||
|
||||
@requires_aiohttp
|
||||
def test_vertex_from_env(monkeypatch):
|
||||
project_id = 'fake_project_id'
|
||||
location = 'fake-location'
|
||||
monkeypatch.setenv('GOOGLE_GENAI_USE_VERTEXAI', 'true')
|
||||
monkeypatch.setenv('GOOGLE_CLOUD_PROJECT', project_id)
|
||||
monkeypatch.setenv('GOOGLE_CLOUD_LOCATION', location)
|
||||
|
||||
client = Client()
|
||||
|
||||
assert client.aio.live.music._api_client.vertexai
|
||||
assert client.aio.live.music._api_client.project == project_id
|
||||
assert isinstance(client.aio.live._api_client, api_client.BaseApiClient)
|
||||
|
||||
|
||||
def test_websocket_base_url():
|
||||
base_url = 'https://test.com'
|
||||
api_client = gl_client.BaseApiClient(
|
||||
api_key='google_api_key',
|
||||
http_options={'base_url': base_url},
|
||||
)
|
||||
assert api_client._websocket_base_url() == 'wss://test.com'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_session_send_weighted_prompts(
|
||||
mock_websocket, vertexai
|
||||
):
|
||||
session = live_music.AsyncMusicSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
if vertexai:
|
||||
with pytest.raises(NotImplementedError):
|
||||
await session.set_weighted_prompts(prompts=[types.WeightedPrompt(text='Jazz', weight=1)])
|
||||
return
|
||||
await session.set_weighted_prompts(prompts=[types.WeightedPrompt(text='Jazz', weight=1)])
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'clientContent' in sent_data
|
||||
assert sent_data['clientContent']['weightedPrompts'][0]['text'] == 'Jazz'
|
||||
assert sent_data['clientContent']['weightedPrompts'][0]['weight'] == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_session_send_config(
|
||||
mock_websocket, vertexai
|
||||
):
|
||||
session = live_music.AsyncMusicSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
if vertexai:
|
||||
with pytest.raises(NotImplementedError):
|
||||
await session.set_music_generation_config(
|
||||
config=types.LiveMusicGenerationConfig(
|
||||
bpm=140,
|
||||
music_generation_mode=types.MusicGenerationMode.VOCALIZATION,
|
||||
)
|
||||
)
|
||||
return
|
||||
await session.set_music_generation_config(
|
||||
config=types.LiveMusicGenerationConfig(
|
||||
bpm=140,
|
||||
music_generation_mode=types.MusicGenerationMode.VOCALIZATION,
|
||||
)
|
||||
)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'musicGenerationConfig' in sent_data
|
||||
assert sent_data['musicGenerationConfig']['bpm'] == 140
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_session_control_signal_play(
|
||||
mock_websocket, vertexai
|
||||
):
|
||||
session = live_music.AsyncMusicSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
if vertexai:
|
||||
with pytest.raises(NotImplementedError):
|
||||
await session.play()
|
||||
return
|
||||
await session.play()
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'playbackControl' in sent_data
|
||||
assert 'PLAY' in sent_data['playbackControl']
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_session_control_signal_pause(
|
||||
mock_websocket, vertexai
|
||||
):
|
||||
session = live_music.AsyncMusicSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
if vertexai:
|
||||
with pytest.raises(NotImplementedError):
|
||||
await session.pause()
|
||||
return
|
||||
await session.pause()
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'playbackControl' in sent_data
|
||||
assert 'PAUSE' in sent_data['playbackControl']
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_session_control_signal_stop(
|
||||
mock_websocket, vertexai
|
||||
):
|
||||
session = live_music.AsyncMusicSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
if vertexai:
|
||||
with pytest.raises(NotImplementedError):
|
||||
await session.stop()
|
||||
return
|
||||
await session.stop()
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'playbackControl' in sent_data
|
||||
assert 'STOP' in sent_data['playbackControl']
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_session_control_signal_reset_context(
|
||||
mock_websocket, vertexai
|
||||
):
|
||||
session = live_music.AsyncMusicSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
if vertexai:
|
||||
with pytest.raises(NotImplementedError):
|
||||
await session.reset_context()
|
||||
return
|
||||
await session.reset_context()
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'playbackControl' in sent_data
|
||||
assert 'RESET_CONTEXT' in sent_data['playbackControl']
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_session_receive( mock_websocket, vertexai):
|
||||
session = live_music.AsyncMusicSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
if vertexai:
|
||||
with pytest.raises(NotImplementedError):
|
||||
async for _ in session.receive():
|
||||
pass
|
||||
return
|
||||
async for response in session.receive():
|
||||
assert isinstance(response, types.LiveMusicServerMessage)
|
||||
audio_chunk = response.server_content.audio_chunks[0]
|
||||
# Data contains decoded b64 audio
|
||||
assert audio_chunk.data == b'gelbanana'
|
||||
assert audio_chunk.mime_type == 'audio/l16;rate=48000;channels=2'
|
||||
assert audio_chunk.source_metadata.client_content.weighted_prompts[0].text == 'Jazz'
|
||||
assert audio_chunk.source_metadata.client_content.weighted_prompts[0].weight == 1
|
||||
assert audio_chunk.source_metadata.music_generation_config.bpm == 140
|
||||
break
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_session_receive_error(
|
||||
mock_websocket, vertexai
|
||||
):
|
||||
mock_websocket.recv = AsyncMock(return_value='invalid json')
|
||||
session = live_music.AsyncMusicSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
with pytest.raises(ValueError):
|
||||
await session.receive().__anext__()
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_session_close( mock_websocket, vertexai):
|
||||
session = live_music.AsyncMusicSession(
|
||||
mock_api_client(vertexai=vertexai), mock_websocket
|
||||
)
|
||||
await session.close()
|
||||
mock_websocket.close.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_setup_to_api(vertexai):
|
||||
if vertexai:
|
||||
with pytest.raises(NotImplementedError):
|
||||
await get_connect_message(
|
||||
mock_api_client(vertexai=vertexai),
|
||||
model='test_model'
|
||||
)
|
||||
return
|
||||
result = await get_connect_message(
|
||||
mock_api_client(vertexai=vertexai),
|
||||
model='test_model'
|
||||
)
|
||||
expected_result = {'setup': {}}
|
||||
if vertexai:
|
||||
# Vertex is not supported yet
|
||||
assert False
|
||||
else:
|
||||
expected_result['setup']['model'] = 'models/test_model'
|
||||
assert result == expected_result
|
||||
@@ -0,0 +1,163 @@
|
||||
# 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 live response handling."""
|
||||
import json
|
||||
from typing import cast
|
||||
from unittest import mock
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import _api_client as api_client
|
||||
from ... import _common
|
||||
from ... import Client
|
||||
from ... import client as gl_client
|
||||
from ... import live
|
||||
from ... import types
|
||||
|
||||
|
||||
def mock_api_client(vertexai=False):
|
||||
"""Creates a mock BaseApiClient."""
|
||||
mock_client = mock.MagicMock(spec=gl_client.BaseApiClient)
|
||||
if not vertexai:
|
||||
mock_client.api_key = 'TEST_API_KEY'
|
||||
mock_client.location = None
|
||||
mock_client.project = None
|
||||
else:
|
||||
mock_client.api_key = None
|
||||
mock_client.location = 'us-central1'
|
||||
mock_client.project = 'test_project'
|
||||
|
||||
mock_client._host = lambda: 'test_host'
|
||||
mock_client._http_options = types.HttpOptions.model_validate(
|
||||
{'headers': {}}
|
||||
)
|
||||
mock_client.vertexai = vertexai
|
||||
return mock_client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_websocket():
|
||||
"""Provides a mock websocket connection."""
|
||||
# Use live.ClientConnection if that's the specific type hint in AsyncSession
|
||||
websocket = AsyncMock(spec=live.ClientConnection)
|
||||
websocket.send = AsyncMock()
|
||||
# Set default recv value, will be overridden in the test
|
||||
websocket.recv = AsyncMock(return_value='{}')
|
||||
websocket.close = AsyncMock()
|
||||
return websocket
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_receive_server_content(mock_websocket, vertexai):
|
||||
|
||||
raw_response_json = json.dumps({
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 15,
|
||||
"responseTokenCount": 25,
|
||||
"candidatesTokenCount": 50,
|
||||
"totalTokenCount": 200,
|
||||
"responseTokensDetails": [
|
||||
{
|
||||
"tokenCount": 20,
|
||||
"modality": "TEXT",
|
||||
}
|
||||
],
|
||||
"candidatesTokensDetails": [
|
||||
{
|
||||
"tokenCount": 10,
|
||||
"modality": "TEXT",
|
||||
}
|
||||
],
|
||||
},
|
||||
"serverContent": {
|
||||
"modelTurn": {
|
||||
"parts": [{"text": "This is a simple response."}]
|
||||
},
|
||||
"turnComplete": True,
|
||||
"groundingMetadata": {
|
||||
"web_search_queries": ["test query"],
|
||||
"groundingChunks": [{
|
||||
"web": {
|
||||
"domain": "google.com",
|
||||
"title": "Search results",
|
||||
}
|
||||
}]
|
||||
}
|
||||
}
|
||||
})
|
||||
mock_websocket.recv.return_value = raw_response_json
|
||||
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
result = await session._receive()
|
||||
|
||||
# Assert the results
|
||||
assert isinstance(result, types.LiveServerMessage)
|
||||
|
||||
assert (
|
||||
result.server_content.model_turn.parts[0].text
|
||||
== "This is a simple response."
|
||||
)
|
||||
assert result.server_content.turn_complete
|
||||
assert result.server_content.grounding_metadata.web_search_queries == ["test query"]
|
||||
assert result.server_content.grounding_metadata.grounding_chunks[0].web.domain == "google.com"
|
||||
assert result.server_content.grounding_metadata.grounding_chunks[0].web.title == "Search results"
|
||||
# Verify usageMetadata was parsed
|
||||
assert isinstance(result.usage_metadata, types.UsageMetadata)
|
||||
assert result.usage_metadata.prompt_token_count == 15
|
||||
assert result.usage_metadata.total_token_count == 200
|
||||
if not vertexai:
|
||||
assert result.usage_metadata.response_token_count == 25
|
||||
assert result.usage_metadata.response_tokens_details[0].token_count == 20
|
||||
else:
|
||||
# VertexAI maps candidatesTokenCount to responseTokenCount and maps
|
||||
# candidatesTokensDetails to responseTokensDetails.
|
||||
assert result.usage_metadata.response_token_count == 50
|
||||
assert result.usage_metadata.response_tokens_details[0].token_count == 10
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_receive_server_content_with_turn_reason(mock_websocket, vertexai):
|
||||
"""Tests parsing of LiveServerContent with turn_complete_reason and waiting_for_input."""
|
||||
|
||||
raw_response_json = json.dumps({
|
||||
"serverContent": {
|
||||
"modelTurn": {
|
||||
"parts": [{"text": "Please provide more details."}]
|
||||
},
|
||||
"turnComplete": True,
|
||||
"turnCompleteReason": "NEED_MORE_INPUT",
|
||||
"waitingForInput": True
|
||||
}
|
||||
})
|
||||
mock_websocket.recv.return_value = raw_response_json
|
||||
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
result = await session._receive()
|
||||
|
||||
# Assert the results
|
||||
assert isinstance(result, types.LiveServerMessage)
|
||||
assert result.server_content is not None
|
||||
|
||||
assert result.server_content.model_turn.parts[0].text == "Please provide more details."
|
||||
assert result.server_content.turn_complete is True
|
||||
assert result.server_content.turn_complete_reason == types.TurnCompleteReason.NEED_MORE_INPUT
|
||||
assert result.server_content.waiting_for_input is True
|
||||
@@ -0,0 +1,147 @@
|
||||
# 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 live.py."""
|
||||
import base64
|
||||
import json
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
from websockets import client
|
||||
|
||||
from .. import pytest_helper
|
||||
from ... import client as gl_client
|
||||
from ... import live
|
||||
from ... import types
|
||||
|
||||
|
||||
def mock_api_client(vertexai=False):
|
||||
api_client = mock.MagicMock(spec=gl_client.BaseApiClient)
|
||||
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_websocket():
|
||||
websocket = mock.AsyncMock(spec=client.ClientConnection)
|
||||
websocket.send = mock.AsyncMock()
|
||||
websocket.recv = mock.AsyncMock(
|
||||
return_value='{"serverContent": {"turnComplete": true}}'
|
||||
) # Default response
|
||||
websocket.close = mock.AsyncMock()
|
||||
return websocket
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_content_dict(mock_websocket, vertexai):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
content = [{'parts': [{'text': 'test'}]}]
|
||||
|
||||
await session.send_client_content(turns=content)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'client_content' in sent_data
|
||||
|
||||
assert sent_data['client_content']['turns'][0]['parts'][0]['text'] == 'test'
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_content_dict_list(mock_websocket, vertexai):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
content = [{'parts': [{'text': 'test'}]}]
|
||||
|
||||
await session.send_client_content(turns=content)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'client_content' in sent_data
|
||||
|
||||
assert sent_data['client_content']['turns'][0]['parts'][0]['text'] == 'test'
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_content_content(mock_websocket, vertexai):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
content = types.Content.model_validate({'parts': [{'text': 'test'}]})
|
||||
|
||||
await session.send_client_content(turns=content)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'client_content' in sent_data
|
||||
|
||||
assert sent_data['client_content']['turns'][0]['parts'][0]['text'] == 'test'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_content_with_blob(mock_websocket, vertexai):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
content = types.Content.model_validate(
|
||||
{'parts': [{'inline_data': {'data': b'test'}}]}
|
||||
)
|
||||
|
||||
await session.send_client_content(turns=content)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'client_content' in sent_data
|
||||
|
||||
assert pytest_helper.get_value_ignore_key_case(
|
||||
sent_data['client_content']['turns'][0]['parts'][0], 'inline_data') == {
|
||||
'data': base64.b64encode(b'test').decode()
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_client_content_turn_complete_false(
|
||||
mock_websocket, vertexai
|
||||
):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
|
||||
await session.send_client_content(turn_complete=False)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'client_content' in sent_data
|
||||
assert sent_data['client_content']['turnComplete'] == False
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_client_content_empty(
|
||||
mock_websocket, vertexai
|
||||
):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
|
||||
await session.send_client_content()
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'client_content' in sent_data
|
||||
assert sent_data['client_content']['turnComplete'] == True
|
||||
@@ -0,0 +1,268 @@
|
||||
# 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 live.py."""
|
||||
import json
|
||||
import os
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
from websockets import client
|
||||
|
||||
from ... import client as gl_client
|
||||
from ... import live
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
|
||||
IMAGE_FILE_PATH = os.path.abspath(
|
||||
os.path.join(os.path.dirname(__file__), '../data/google.jpg')
|
||||
)
|
||||
|
||||
|
||||
def mock_api_client(vertexai=False):
|
||||
api_client = mock.MagicMock(spec=gl_client.BaseApiClient)
|
||||
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
|
||||
api_client._api_client = api_client
|
||||
return api_client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_websocket():
|
||||
websocket = mock.AsyncMock(spec=client.ClientConnection)
|
||||
websocket.send = mock.AsyncMock()
|
||||
websocket.recv = mock.AsyncMock(
|
||||
return_value='{"serverContent": {"turnComplete": true}}'
|
||||
) # Default response
|
||||
websocket.close = mock.AsyncMock()
|
||||
return websocket
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_media_blob_dict(mock_websocket, vertexai):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
content = {'data': bytes([0, 0, 0, 0, 0, 0]), 'mime_type': 'audio/pcm'}
|
||||
|
||||
await session.send_realtime_input(media=content)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert sent_data['realtime_input']['mediaChunks'][0]['data'] == 'AAAAAAAA'
|
||||
assert pytest_helper.get_value_ignore_key_case(
|
||||
sent_data['realtime_input']['mediaChunks'][0], 'mime_type'
|
||||
) == 'audio/pcm'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_media_blob(mock_websocket, vertexai):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
content = types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='audio/pcm')
|
||||
|
||||
await session.send_realtime_input(media=content)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert sent_data['realtime_input']['mediaChunks'][0]['data'] == 'AAAAAAAA'
|
||||
assert pytest_helper.get_value_ignore_key_case(
|
||||
sent_data['realtime_input']['mediaChunks'][0], 'mime_type'
|
||||
) == 'audio/pcm'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_media_image(mock_websocket, vertexai, image_jpeg):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
|
||||
await session.send_realtime_input(media=image_jpeg)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert pytest_helper.get_value_ignore_key_case(
|
||||
sent_data['realtime_input']['mediaChunks'][0], 'mime_type'
|
||||
) == 'image/jpeg'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
'content',
|
||||
[
|
||||
{'data': bytes([0, 0, 0, 0, 0, 0]), 'mime_type': 'audio/pcm'},
|
||||
types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='audio/pcm'),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_audio(mock_websocket, vertexai, content):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
|
||||
|
||||
await session.send_realtime_input(audio=content)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert sent_data['realtime_input']['audio']['data'] == 'AAAAAAAA'
|
||||
assert pytest_helper.get_value_ignore_key_case(
|
||||
sent_data['realtime_input']['audio'], 'mime_type'
|
||||
) == 'audio/pcm'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_bad_audio_blob(mock_websocket, vertexai):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
|
||||
content = types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='image/png')
|
||||
|
||||
with pytest.raises(ValueError, match='.*Unsupported mime type.*'):
|
||||
await session.send_realtime_input(audio=content)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_bad_video_blob(mock_websocket, vertexai):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
|
||||
content = types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='audio/pcm')
|
||||
|
||||
with pytest.raises(ValueError, match='.*Unsupported mime type.*'):
|
||||
await session.send_realtime_input(video=content)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_audio_stream_end(mock_websocket, vertexai):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
|
||||
await session.send_realtime_input(audio_stream_end=True)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert sent_data['realtime_input']['audioStreamEnd'] == True
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
'content',
|
||||
[
|
||||
{'data': bytes([0, 0, 0, 0, 0, 0]), 'mime_type': 'image/png'},
|
||||
types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='image/png'),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_video(mock_websocket, vertexai, content):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
|
||||
|
||||
await session.send_realtime_input(video=content)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert sent_data['realtime_input']['video']['data'] == 'AAAAAAAA'
|
||||
assert pytest_helper.get_value_ignore_key_case(
|
||||
sent_data['realtime_input']['video'], 'mime_type'
|
||||
) == 'image/png'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_video_image(mock_websocket, vertexai, image_jpeg):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
|
||||
|
||||
await session.send_realtime_input(video=image_jpeg)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert pytest_helper.get_value_ignore_key_case(
|
||||
sent_data['realtime_input']['video'], 'mime_type'
|
||||
) == 'image/jpeg'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_text(mock_websocket, vertexai):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
|
||||
|
||||
await session.send_realtime_input(text='Hello?')
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert sent_data['realtime_input']['text'] == 'Hello?'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.parametrize('activity', [{}, types.ActivityStart()])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_activity_start(mock_websocket, vertexai, activity):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
|
||||
await session.send_realtime_input(activity_start=activity)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert sent_data['realtime_input']['activityStart'] == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.parametrize('activity', [{}, types.ActivityEnd()])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_activity_end(mock_websocket, vertexai, activity):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
|
||||
await session.send_realtime_input(activity_end=activity)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert sent_data['realtime_input']['activityEnd'] == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_multiple_args(mock_websocket, vertexai):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
with pytest.raises(ValueError, match='.*one argument.*'):
|
||||
await session.send_realtime_input(
|
||||
text='Hello?', activity_start=types.ActivityStart()
|
||||
)
|
||||
@@ -0,0 +1,222 @@
|
||||
# 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 live.py."""
|
||||
|
||||
import contextlib
|
||||
import json
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
from websockets import client
|
||||
|
||||
from .. import pytest_helper
|
||||
from ... import client as gl_client
|
||||
from ... import live
|
||||
from ... import types
|
||||
|
||||
|
||||
def exception_if_mldev(vertexai, exception_type: type[Exception]):
|
||||
if vertexai:
|
||||
return contextlib.nullcontext()
|
||||
else:
|
||||
return pytest.raises(exception_type)
|
||||
|
||||
|
||||
def mock_api_client(vertexai=False):
|
||||
api_client = mock.MagicMock(spec=gl_client.BaseApiClient)
|
||||
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
|
||||
api_client._api_client = api_client
|
||||
return api_client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_websocket():
|
||||
websocket = mock.AsyncMock(spec=client.ClientConnection)
|
||||
websocket.send = mock.AsyncMock()
|
||||
websocket.recv = mock.AsyncMock(
|
||||
return_value='{"serverContent": {"turnComplete": true}}'
|
||||
) # Default response
|
||||
websocket.close = mock.AsyncMock()
|
||||
return websocket
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_function_response_dict(mock_websocket, vertexai):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(
|
||||
api_client=api_client, websocket=mock_websocket
|
||||
)
|
||||
|
||||
input = {
|
||||
'name': 'get_current_weather',
|
||||
'response': {'temperature': 14.5, 'unit': 'C'},
|
||||
}
|
||||
|
||||
if not vertexai:
|
||||
input['id'] = 'some-id'
|
||||
|
||||
await session.send_tool_response(function_responses=input)
|
||||
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'tool_response' in sent_data
|
||||
|
||||
assert (
|
||||
sent_data['tool_response']['functionResponses'][0]['name']
|
||||
== 'get_current_weather'
|
||||
)
|
||||
assert (
|
||||
sent_data['tool_response']['functionResponses'][0]['response'][
|
||||
'temperature'
|
||||
]
|
||||
== 14.5
|
||||
)
|
||||
assert (
|
||||
sent_data['tool_response']['functionResponses'][0]['response']['unit']
|
||||
== 'C'
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_function_response(mock_websocket, vertexai):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
|
||||
input = types.FunctionResponse(
|
||||
name='get_current_weather',
|
||||
response={
|
||||
'temperature': 14.5,
|
||||
'unit': 'C',
|
||||
'user_name': 'test_user_name',
|
||||
'userEmail': 'test_user_email',
|
||||
},
|
||||
)
|
||||
if not vertexai:
|
||||
input.id = 'some-id'
|
||||
|
||||
await session.send_tool_response(function_responses=input)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'tool_response' in sent_data
|
||||
|
||||
assert (
|
||||
sent_data['tool_response']['functionResponses'][0]['name']
|
||||
== 'get_current_weather'
|
||||
)
|
||||
assert (
|
||||
sent_data['tool_response']['functionResponses'][0]['response']
|
||||
== input.response
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_function_response_scheduling(mock_websocket, vertexai):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
|
||||
|
||||
input = types.FunctionResponse(
|
||||
name='get_current_weather',
|
||||
response={'temperature': 14.5, 'unit': 'C'},
|
||||
will_continue=True,
|
||||
scheduling=types.FunctionResponseScheduling.SILENT,
|
||||
)
|
||||
if not vertexai:
|
||||
input.id = 'some-id'
|
||||
|
||||
await session.send_tool_response(function_responses=input)
|
||||
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'tool_response' in sent_data
|
||||
|
||||
assert pytest_helper.get_value_ignore_key_case(
|
||||
sent_data['tool_response']['functionResponses'][0], 'will_continue'
|
||||
)
|
||||
assert (
|
||||
sent_data['tool_response']['functionResponses'][0]['scheduling']
|
||||
== 'SILENT'
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_function_response_list(mock_websocket, vertexai):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
|
||||
input1 = {
|
||||
'name': 'get_current_weather',
|
||||
'response': {'temperature': 14.5, 'unit': 'C'},
|
||||
}
|
||||
input2 = {
|
||||
'name': 'get_current_weather',
|
||||
'response': {'temperature': 99.9, 'unit': 'C'},
|
||||
}
|
||||
|
||||
if not vertexai:
|
||||
input1['id'] = '1'
|
||||
input2['id'] = '2'
|
||||
|
||||
await session.send_tool_response(function_responses=[input1, input2])
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'tool_response' in sent_data
|
||||
|
||||
assert len(sent_data['tool_response']['functionResponses']) == 2
|
||||
assert (
|
||||
sent_data['tool_response']['functionResponses'][0]['response'][
|
||||
'temperature'
|
||||
]
|
||||
== 14.5
|
||||
)
|
||||
assert (
|
||||
sent_data['tool_response']['functionResponses'][1]['response'][
|
||||
'temperature'
|
||||
]
|
||||
== 99.9
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_id(mock_websocket, vertexai):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(
|
||||
api_client=api_client, websocket=mock_websocket
|
||||
)
|
||||
|
||||
input1 = {
|
||||
'name': 'get_current_weather',
|
||||
'response': {'temperature': 14.5, 'unit': 'C'},
|
||||
'id': '1',
|
||||
}
|
||||
input2 = {
|
||||
'name': 'get_current_weather',
|
||||
'response': {'temperature': 99.9, 'unit': 'C'},
|
||||
}
|
||||
|
||||
if not vertexai:
|
||||
with pytest.raises(ValueError, match=".*must have.*"):
|
||||
await session.send_tool_response(function_responses=[input1, input2])
|
||||
@@ -0,0 +1,17 @@
|
||||
# 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 local_tokenizer module."""
|
||||
@@ -0,0 +1,343 @@
|
||||
# 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 unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from sentencepiece import sentencepiece_model_pb2
|
||||
|
||||
from ... import local_tokenizer
|
||||
from ... import types
|
||||
|
||||
|
||||
class TestLocalTokenizer(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
# This setup will be used by all tests
|
||||
self.mock_load_model_proto = patch(
|
||||
'genai._local_tokenizer_loader.load_model_proto'
|
||||
).start()
|
||||
self.mock_get_sentencepiece = patch(
|
||||
'genai._local_tokenizer_loader.get_sentencepiece'
|
||||
).start()
|
||||
|
||||
self.mock_load_model_proto.return_value = MagicMock()
|
||||
self.mock_tokenizer = MagicMock()
|
||||
self.mock_get_sentencepiece.return_value = self.mock_tokenizer
|
||||
|
||||
self.tokenizer = local_tokenizer.LocalTokenizer(model_name='gemini-3-pro-preview')
|
||||
|
||||
def tearDown(self):
|
||||
patch.stopall()
|
||||
|
||||
def test_count_tokens_simple_string(self):
|
||||
self.mock_tokenizer.encode.return_value = [[1, 2, 3]]
|
||||
result = self.tokenizer.count_tokens('Hello world')
|
||||
self.assertEqual(result.total_tokens, 3)
|
||||
self.mock_tokenizer.encode.assert_called_once_with(['Hello world'])
|
||||
|
||||
def test_count_tokens_list_of_strings(self):
|
||||
self.mock_tokenizer.encode.return_value = [[1, 2], [3]]
|
||||
result = self.tokenizer.count_tokens(['Hello', 'world'])
|
||||
self.assertEqual(result.total_tokens, 3)
|
||||
self.mock_tokenizer.encode.assert_called_once_with(['Hello', 'world'])
|
||||
|
||||
def test_count_tokens_with_content_object(self):
|
||||
self.mock_tokenizer.encode.return_value = [[1, 2, 3]]
|
||||
content = types.Content(parts=[types.Part(text='Hello world')])
|
||||
result = self.tokenizer.count_tokens(content)
|
||||
self.assertEqual(result.total_tokens, 3)
|
||||
self.mock_tokenizer.encode.assert_called_once_with(['Hello world'])
|
||||
|
||||
def test_count_tokens_with_chat_history(self):
|
||||
self.mock_tokenizer.encode.return_value = [[1, 2], [3, 4, 5]]
|
||||
history = [
|
||||
types.Content(role='user', parts=[types.Part(text='Hello')]),
|
||||
types.Content(role='model', parts=[types.Part(text='Hi there!')]),
|
||||
]
|
||||
result = self.tokenizer.count_tokens(history)
|
||||
self.assertEqual(result.total_tokens, 5)
|
||||
self.mock_tokenizer.encode.assert_called_once_with(['Hello', 'Hi there!'])
|
||||
|
||||
def test_count_tokens_with_tools(self):
|
||||
self.mock_tokenizer.encode.return_value = [
|
||||
[1],
|
||||
[1, 2],
|
||||
[1, 2, 3],
|
||||
[1, 2, 3, 4],
|
||||
[1, 2, 3, 4, 5],
|
||||
[1, 2, 3, 4, 5, 6],
|
||||
]
|
||||
tool = types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(
|
||||
name='get_weather',
|
||||
description='Get the weather for a location',
|
||||
parameters=types.Schema(
|
||||
type=types.Type.OBJECT,
|
||||
properties={
|
||||
'location': types.Schema(
|
||||
type=types.Type.STRING, description='The location'
|
||||
)
|
||||
},
|
||||
required=['location'],
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
config = types.CountTokensConfig(tools=[tool])
|
||||
result = self.tokenizer.count_tokens(
|
||||
'What is the weather in Boston?', config=config
|
||||
)
|
||||
self.assertEqual(result.total_tokens, 21)
|
||||
self.mock_tokenizer.encode.assert_called_once_with([
|
||||
'What is the weather in Boston?',
|
||||
'get_weather',
|
||||
'Get the weather for a location',
|
||||
'location',
|
||||
'location',
|
||||
'The location',
|
||||
])
|
||||
|
||||
def test_count_tokens_with_function_call(self):
|
||||
self.mock_tokenizer.encode.return_value = [[1, 2], [3], [4, 5]]
|
||||
content = types.Content(
|
||||
role='model',
|
||||
parts=[
|
||||
types.Part(
|
||||
function_call=types.FunctionCall(
|
||||
name='get_weather', args={'location': 'Boston'}
|
||||
)
|
||||
)
|
||||
],
|
||||
)
|
||||
result = self.tokenizer.count_tokens(content)
|
||||
self.assertEqual(result.total_tokens, 5)
|
||||
self.mock_tokenizer.encode.assert_called_once_with(
|
||||
['get_weather', 'location', 'Boston']
|
||||
)
|
||||
|
||||
def test_count_tokens_with_function_response(self):
|
||||
self.mock_tokenizer.encode.return_value = [[1, 2], [3], [4, 5]]
|
||||
content = types.Content(
|
||||
role='user',
|
||||
parts=[
|
||||
types.Part(
|
||||
function_response=types.FunctionResponse(
|
||||
name='get_weather', response={'weather': 'sunny'}
|
||||
)
|
||||
)
|
||||
],
|
||||
)
|
||||
result = self.tokenizer.count_tokens(content)
|
||||
self.assertEqual(result.total_tokens, 5)
|
||||
self.mock_tokenizer.encode.assert_called_once_with(
|
||||
['get_weather', 'weather', 'sunny']
|
||||
)
|
||||
|
||||
def test_count_tokens_with_unsupported_content(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self.tokenizer.count_tokens(
|
||||
[
|
||||
types.Content(
|
||||
parts=[
|
||||
types.Part(
|
||||
inline_data=types.Blob(
|
||||
data=b'test', mime_type='image/png'
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
def test_count_tokens_with_system_instruction(self):
|
||||
self.mock_tokenizer.encode.return_value = [[1, 2, 3], [4, 5]]
|
||||
config = types.CountTokensConfig(
|
||||
system_instruction=types.Content(
|
||||
parts=[types.Part(text='You are a helpful assistant.')]
|
||||
)
|
||||
)
|
||||
result = self.tokenizer.count_tokens('Hello', config=config)
|
||||
self.assertEqual(result.total_tokens, 5)
|
||||
self.mock_tokenizer.encode.assert_called_once_with(
|
||||
['Hello', 'You are a helpful assistant.']
|
||||
)
|
||||
|
||||
def test_count_tokens_with_response_schema(self):
|
||||
self.mock_tokenizer.encode.return_value = [
|
||||
[1],
|
||||
[1, 2],
|
||||
[1, 2, 3],
|
||||
[1, 2, 3, 4],
|
||||
[1, 2, 3, 4, 5],
|
||||
]
|
||||
schema = types.Schema(
|
||||
type=types.Type.OBJECT,
|
||||
format='schema_format',
|
||||
description='Recipe schema',
|
||||
enum=['schema_enum1', 'schema_enum2'],
|
||||
properties={
|
||||
'recipe_name': types.Schema(
|
||||
type=types.Type.STRING,
|
||||
description='Name of the recipe',
|
||||
)
|
||||
},
|
||||
items=types.Schema(
|
||||
type=types.Type.STRING,
|
||||
description='Item in the recipe',
|
||||
),
|
||||
example={
|
||||
'recipe_example': types.Schema(
|
||||
type=types.Type.STRING,
|
||||
description='example in the recipe',
|
||||
)
|
||||
},
|
||||
required=['recipe_name'],
|
||||
)
|
||||
config = types.CountTokensConfig(
|
||||
generation_config=types.GenerationConfig(response_schema=schema)
|
||||
)
|
||||
result = self.tokenizer.count_tokens(
|
||||
'Generate a recipe for chocolate chip cookies.', config=config
|
||||
)
|
||||
self.assertEqual(result.total_tokens, 15)
|
||||
self.mock_tokenizer.encode.assert_called_once_with([
|
||||
'Generate a recipe for chocolate chip cookies.',
|
||||
'schema_format',
|
||||
'Recipe schema',
|
||||
'schema_enum1',
|
||||
'schema_enum2',
|
||||
'recipe_name',
|
||||
'Item in the recipe',
|
||||
'recipe_name',
|
||||
'Name of the recipe',
|
||||
'recipe_example',
|
||||
])
|
||||
|
||||
def test_count_tokens_with_unsupported_fields_logs_warning(self):
|
||||
self.mock_tokenizer.encode.return_value = [[1, 2, 3]]
|
||||
content_with_unsupported = types.Content(
|
||||
role='user',
|
||||
parts=[
|
||||
types.Part(text='hello'),
|
||||
# executable_code is not supported by _TextsAccumulator
|
||||
types.Part(
|
||||
executable_code=types.ExecutableCode(
|
||||
language='PYTHON', code='print(1)'
|
||||
)
|
||||
),
|
||||
],
|
||||
)
|
||||
with self.assertLogs('google_genai.local_tokenizer', level='WARNING') as cm:
|
||||
self.tokenizer.count_tokens(content_with_unsupported)
|
||||
self.assertIn(
|
||||
'Content contains unsupported types for token counting', cm.output[0]
|
||||
)
|
||||
|
||||
def test_compute_tokens_simple_string(self):
|
||||
mock_spt = MagicMock()
|
||||
mock_spt.pieces = [
|
||||
MagicMock(id=1, piece='He'),
|
||||
MagicMock(id=2, piece='llo'),
|
||||
MagicMock(id=3, piece=' world'),
|
||||
]
|
||||
self.mock_tokenizer.EncodeAsImmutableProto.return_value = [mock_spt]
|
||||
result = self.tokenizer.compute_tokens('Hello world')
|
||||
self.assertEqual(len(result.tokens_info), 1)
|
||||
self.assertEqual(result.tokens_info[0].token_ids, [1, 2, 3])
|
||||
self.assertEqual(result.tokens_info[0].tokens, [b'He', b'llo', b' world'])
|
||||
self.assertEqual(result.tokens_info[0].role, 'user')
|
||||
self.mock_tokenizer.EncodeAsImmutableProto.assert_called_once_with(
|
||||
['Hello world']
|
||||
)
|
||||
|
||||
def test_compute_tokens_with_chat_history(self):
|
||||
mock_spt1 = MagicMock()
|
||||
mock_spt1.pieces = [MagicMock(id=1, piece='Hello')]
|
||||
mock_spt2 = MagicMock()
|
||||
mock_spt2.pieces = [
|
||||
MagicMock(id=2, piece='Hi'),
|
||||
MagicMock(id=3, piece=' there!'),
|
||||
]
|
||||
self.mock_tokenizer.EncodeAsImmutableProto.return_value = [
|
||||
mock_spt1,
|
||||
mock_spt2,
|
||||
]
|
||||
history = [
|
||||
types.Content(role='user', parts=[types.Part(text='Hello')]),
|
||||
types.Content(role='model', parts=[types.Part(text='Hi there!')]),
|
||||
]
|
||||
result = self.tokenizer.compute_tokens(history)
|
||||
self.assertEqual(len(result.tokens_info), 2)
|
||||
self.assertEqual(result.tokens_info[0].token_ids, [1])
|
||||
self.assertEqual(result.tokens_info[0].tokens, [b'Hello'])
|
||||
self.assertEqual(result.tokens_info[0].role, 'user')
|
||||
self.assertEqual(result.tokens_info[1].token_ids, [2, 3])
|
||||
self.assertEqual(result.tokens_info[1].tokens, [b'Hi', b' there!'])
|
||||
self.assertEqual(result.tokens_info[1].role, 'model')
|
||||
self.mock_tokenizer.EncodeAsImmutableProto.assert_called_once_with(
|
||||
['Hello', 'Hi there!']
|
||||
)
|
||||
|
||||
def test_compute_tokens_with_byte_tokens(self):
|
||||
mock_spt = MagicMock()
|
||||
mock_spt.pieces = [
|
||||
MagicMock(id=1, piece='<0x48>'),
|
||||
MagicMock(id=2, piece='ello'),
|
||||
]
|
||||
self.mock_tokenizer.EncodeAsImmutableProto.return_value = [mock_spt]
|
||||
self.tokenizer._model_proto = sentencepiece_model_pb2.ModelProto(
|
||||
pieces=[
|
||||
sentencepiece_model_pb2.ModelProto.SentencePiece(),
|
||||
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
||||
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.BYTE
|
||||
),
|
||||
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
||||
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.NORMAL
|
||||
),
|
||||
]
|
||||
)
|
||||
result = self.tokenizer.compute_tokens('Hello')
|
||||
self.assertEqual(len(result.tokens_info), 1)
|
||||
self.assertEqual(result.tokens_info[0].token_ids, [1, 2])
|
||||
self.assertEqual(result.tokens_info[0].tokens, [b'H', b'ello'])
|
||||
self.mock_tokenizer.EncodeAsImmutableProto.assert_called_once_with(
|
||||
['Hello']
|
||||
)
|
||||
|
||||
|
||||
class TestParseHexByte(unittest.TestCase):
|
||||
|
||||
def test_valid_hex(self):
|
||||
self.assertEqual(local_tokenizer._parse_hex_byte('<0x41>'), 65)
|
||||
self.assertEqual(local_tokenizer._parse_hex_byte('<0xFF>'), 255)
|
||||
self.assertEqual(local_tokenizer._parse_hex_byte('<0x00>'), 0)
|
||||
|
||||
def test_invalid_length(self):
|
||||
with self.assertRaisesRegex(ValueError, 'Invalid byte length'):
|
||||
local_tokenizer._parse_hex_byte('<0x41')
|
||||
with self.assertRaisesRegex(ValueError, 'Invalid byte length'):
|
||||
local_tokenizer._parse_hex_byte('<0x411>')
|
||||
|
||||
def test_invalid_format(self):
|
||||
with self.assertRaisesRegex(ValueError, 'Invalid byte format'):
|
||||
local_tokenizer._parse_hex_byte(' 0x41>')
|
||||
with self.assertRaisesRegex(ValueError, 'Invalid byte format'):
|
||||
local_tokenizer._parse_hex_byte('<0x41 ')
|
||||
|
||||
def test_invalid_hex_value(self):
|
||||
with self.assertRaisesRegex(ValueError, 'Invalid hex value'):
|
||||
local_tokenizer._parse_hex_byte('<0xFG>')
|
||||
@@ -0,0 +1,235 @@
|
||||
# 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 unittest
|
||||
from unittest.mock import MagicMock, mock_open, patch
|
||||
|
||||
import sentencepiece as spm
|
||||
from sentencepiece import sentencepiece_model_pb2
|
||||
|
||||
from ... import _local_tokenizer_loader as loader
|
||||
|
||||
# A minimal valid sentencepiece model proto
|
||||
FAKE_MODEL_CONTENT = sentencepiece_model_pb2.ModelProto(
|
||||
pieces=[
|
||||
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
||||
piece="<unk>",
|
||||
score=0,
|
||||
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.UNKNOWN,
|
||||
),
|
||||
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
||||
piece="<s>",
|
||||
score=0,
|
||||
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.CONTROL,
|
||||
),
|
||||
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
||||
piece="</s>",
|
||||
score=0,
|
||||
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.CONTROL,
|
||||
),
|
||||
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
||||
piece="a",
|
||||
score=0,
|
||||
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.NORMAL,
|
||||
),
|
||||
]
|
||||
).SerializeToString()
|
||||
|
||||
GEMMA2_HASH = "61a7b147390c64585d6c3543dd6fc636906c9af3865a5548f27f31aee1d4c8e2"
|
||||
|
||||
|
||||
class TestGetTokenizerName(unittest.TestCase):
|
||||
|
||||
def test_get_tokenizer_name_success(self):
|
||||
self.assertEqual(loader.get_tokenizer_name("gemini-2.5-pro"), "gemma3")
|
||||
self.assertEqual(
|
||||
loader.get_tokenizer_name("gemini-2.5-pro-preview-06-05"), "gemma3"
|
||||
)
|
||||
|
||||
def test_get_tokenizer_name_unsupported(self):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, "Model unsupported-model is not supported"
|
||||
):
|
||||
loader.get_tokenizer_name("unsupported-model")
|
||||
|
||||
|
||||
@patch("genai._local_tokenizer_loader.os.rename")
|
||||
@patch("genai._local_tokenizer_loader.os.makedirs")
|
||||
@patch("genai._local_tokenizer_loader.os.remove")
|
||||
@patch("genai._local_tokenizer_loader.open", new_callable=mock_open)
|
||||
@patch("genai._local_tokenizer_loader.os.path.exists")
|
||||
@patch("genai._local_tokenizer_loader.requests.get")
|
||||
@patch("genai._local_tokenizer_loader.hashlib.sha256")
|
||||
class TestLoaderFunctions(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
# Clear caches before each test
|
||||
loader.load_model_proto.cache_clear()
|
||||
loader.get_sentencepiece.cache_clear()
|
||||
# Patch tempfile.gettempdir to control cache location
|
||||
self.tempdir_patcher = patch(
|
||||
"tempfile.gettempdir", return_value="/tmp/fake_temp_dir"
|
||||
)
|
||||
self.mock_tempdir = self.tempdir_patcher.start()
|
||||
|
||||
def tearDown(self):
|
||||
self.tempdir_patcher.stop()
|
||||
|
||||
def _setup_get_mock(self, mock_get):
|
||||
mock_response = MagicMock()
|
||||
mock_response.content = FAKE_MODEL_CONTENT
|
||||
mock_response.raise_for_status.return_value = None
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
def test_load_model_proto_from_url(
|
||||
self,
|
||||
mock_sha256,
|
||||
mock_get,
|
||||
mock_exists,
|
||||
mock_open_func,
|
||||
mock_remove,
|
||||
mock_makedirs,
|
||||
mock_rename,
|
||||
):
|
||||
mock_exists.return_value = False # Don't use cache
|
||||
self._setup_get_mock(mock_get)
|
||||
mock_sha256.return_value.hexdigest.return_value = GEMMA2_HASH
|
||||
|
||||
proto = loader.load_model_proto("gemma2")
|
||||
|
||||
self.assertIsInstance(proto, sentencepiece_model_pb2.ModelProto)
|
||||
self.assertEqual(len(proto.pieces), 4)
|
||||
mock_get.assert_called_once()
|
||||
mock_makedirs.assert_called_once()
|
||||
mock_open_func.assert_called()
|
||||
mock_rename.assert_called_once()
|
||||
|
||||
def test_load_model_proto_from_cache(
|
||||
self,
|
||||
mock_sha256,
|
||||
mock_get,
|
||||
mock_exists,
|
||||
mock_open_func,
|
||||
mock_remove,
|
||||
mock_makedirs,
|
||||
mock_rename,
|
||||
):
|
||||
mock_exists.return_value = True # Use cache
|
||||
mock_open_func.return_value.read.return_value = FAKE_MODEL_CONTENT
|
||||
mock_sha256.return_value.hexdigest.return_value = GEMMA2_HASH
|
||||
|
||||
proto = loader.load_model_proto("gemma2")
|
||||
|
||||
self.assertIsInstance(proto, sentencepiece_model_pb2.ModelProto)
|
||||
mock_get.assert_not_called()
|
||||
|
||||
def test_load_model_proto_corrupted_cache(
|
||||
self,
|
||||
mock_sha256,
|
||||
mock_get,
|
||||
mock_exists,
|
||||
mock_open_func,
|
||||
mock_remove,
|
||||
mock_makedirs,
|
||||
mock_rename,
|
||||
):
|
||||
mock_exists.return_value = True # Use cache initially
|
||||
self._setup_get_mock(mock_get)
|
||||
mock_open_func.return_value.__enter__.return_value.read.return_value = (
|
||||
b"corrupted"
|
||||
)
|
||||
|
||||
# First hash for corrupted cache, second for good download
|
||||
mock_sha256.side_effect = [
|
||||
MagicMock(hexdigest=MagicMock(return_value="wrong_hash")),
|
||||
MagicMock(hexdigest=MagicMock(return_value=GEMMA2_HASH)),
|
||||
]
|
||||
|
||||
proto = loader.load_model_proto("gemma2")
|
||||
|
||||
self.assertIsInstance(proto, sentencepiece_model_pb2.ModelProto)
|
||||
mock_remove.assert_called_once()
|
||||
mock_get.assert_called_once()
|
||||
|
||||
def test_load_model_proto_bad_hash_from_url(
|
||||
self,
|
||||
mock_sha256,
|
||||
mock_get,
|
||||
mock_exists,
|
||||
mock_open_func,
|
||||
mock_remove,
|
||||
mock_makedirs,
|
||||
mock_rename,
|
||||
):
|
||||
mock_exists.return_value = False
|
||||
self._setup_get_mock(mock_get)
|
||||
mock_sha256.return_value.hexdigest.return_value = "wrong_hash"
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, "Downloaded model file is corrupted"
|
||||
):
|
||||
loader.load_model_proto("gemma2")
|
||||
|
||||
def test_load_model_proto_unsupported(self, *args):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, "Tokenizer unsupported is not supported"
|
||||
):
|
||||
loader.load_model_proto("unsupported")
|
||||
|
||||
def test_get_sentencepiece_success(
|
||||
self,
|
||||
mock_sha256,
|
||||
mock_get,
|
||||
mock_exists,
|
||||
mock_open_func,
|
||||
mock_remove,
|
||||
mock_makedirs,
|
||||
mock_rename,
|
||||
):
|
||||
mock_exists.return_value = False
|
||||
self._setup_get_mock(mock_get)
|
||||
mock_sha256.return_value.hexdigest.return_value = GEMMA2_HASH
|
||||
|
||||
processor = loader.get_sentencepiece("gemma2")
|
||||
|
||||
self.assertIsInstance(processor, spm.SentencePieceProcessor)
|
||||
mock_get.assert_called_once()
|
||||
|
||||
def test_get_sentencepiece_unsupported(self, *args):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, "Tokenizer unsupported is not supported"
|
||||
):
|
||||
loader.get_sentencepiece("unsupported")
|
||||
|
||||
def test_get_sentencepiece_caching(
|
||||
self,
|
||||
mock_sha256,
|
||||
mock_get,
|
||||
mock_exists,
|
||||
mock_open_func,
|
||||
mock_remove,
|
||||
mock_makedirs,
|
||||
mock_rename,
|
||||
):
|
||||
mock_exists.return_value = False
|
||||
self._setup_get_mock(mock_get)
|
||||
mock_sha256.return_value.hexdigest.return_value = GEMMA2_HASH
|
||||
|
||||
# Call twice
|
||||
loader.get_sentencepiece("gemma2")
|
||||
loader.get_sentencepiece("gemma2")
|
||||
|
||||
# Should only be loaded once due to lru_cache
|
||||
mock_get.assert_called_once()
|
||||
@@ -0,0 +1,17 @@
|
||||
# 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 _mcp_utils module."""
|
||||
@@ -0,0 +1,89 @@
|
||||
# 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 typing
|
||||
from typing import Any
|
||||
|
||||
from ... import _mcp_utils
|
||||
from ... import types
|
||||
|
||||
if typing.TYPE_CHECKING:
|
||||
from mcp.types import Tool as McpTool
|
||||
from mcp import ClientSession as McpClientSession
|
||||
else:
|
||||
McpTool: typing.Type = Any
|
||||
McpClientSession: typing.Type = Any
|
||||
try:
|
||||
from mcp.types import Tool as McpTool
|
||||
from mcp import ClientSession as McpClientSession
|
||||
except ImportError:
|
||||
McpTool = None
|
||||
McpClientSession = None
|
||||
|
||||
|
||||
def test_mcp_tools():
|
||||
"""Test whether the list of tools contains any MCP tools."""
|
||||
if McpTool is None:
|
||||
return
|
||||
mcp_tools = [
|
||||
McpTool(
|
||||
name='tool',
|
||||
description='tool-description',
|
||||
inputSchema={
|
||||
'type': 'OBJECT',
|
||||
'properties': {
|
||||
'key1': {'type': 'STRING'},
|
||||
'key2': {'type': 'NUMBER'},
|
||||
},
|
||||
},
|
||||
),
|
||||
]
|
||||
assert _mcp_utils.has_mcp_tool_usage(mcp_tools)
|
||||
|
||||
|
||||
def test_mcp_client_session():
|
||||
"""Test whether the list of tools contains any MCP tools."""
|
||||
|
||||
class MockMcpClientSession(McpClientSession):
|
||||
|
||||
def __init__(self):
|
||||
self._read_stream = None
|
||||
self._write_stream = None
|
||||
|
||||
mcp_tools = [
|
||||
MockMcpClientSession(),
|
||||
]
|
||||
assert _mcp_utils.has_mcp_tool_usage(mcp_tools)
|
||||
|
||||
|
||||
def test_no_mcp_tools():
|
||||
if McpClientSession is None:
|
||||
return
|
||||
"""Test whether the list of tools contains any MCP tools."""
|
||||
gemini_tools = [
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(
|
||||
name='tool',
|
||||
description='tool-description',
|
||||
parameters=types.Schema(
|
||||
type='OBJECT',
|
||||
properties={},
|
||||
),
|
||||
),
|
||||
],
|
||||
),
|
||||
]
|
||||
assert not _mcp_utils.has_mcp_tool_usage(gemini_tools)
|
||||
@@ -0,0 +1,191 @@
|
||||
# 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 ... import _mcp_utils
|
||||
from ... import types
|
||||
|
||||
try:
|
||||
from mcp import types as mcp_types
|
||||
except ImportError as e:
|
||||
import sys
|
||||
|
||||
if sys.version_info < (3, 10):
|
||||
raise ImportError(
|
||||
'MCP Tool requires Python 3.10 or above. Please upgrade your Python'
|
||||
' version.'
|
||||
) from e
|
||||
else:
|
||||
raise e
|
||||
|
||||
|
||||
def test_empty_mcp_tools_list():
|
||||
"""Test conversion of empty MCP tools list to Gemini tools list."""
|
||||
result = _mcp_utils.mcp_to_gemini_tools([])
|
||||
|
||||
assert result == []
|
||||
|
||||
|
||||
def test_unknown_field_conversion():
|
||||
"""Test conversion of MCP tools with unknown fields to Gemini tools."""
|
||||
mcp_tools = [
|
||||
mcp_types.Tool(
|
||||
name='tool',
|
||||
description='tool-description',
|
||||
inputSchema={
|
||||
'type': 'object',
|
||||
'properties': {},
|
||||
'unknown_field': 'unknownField',
|
||||
'unknown_object': {},
|
||||
},
|
||||
),
|
||||
]
|
||||
result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
|
||||
assert result == [
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(
|
||||
name='tool',
|
||||
description='tool-description',
|
||||
parameters=types.Schema(
|
||||
type='OBJECT',
|
||||
properties={},
|
||||
),
|
||||
),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_items_conversion():
|
||||
"""Test conversion of MCP tools with items to Gemini tools."""
|
||||
mcp_tools = [
|
||||
mcp_types.Tool(
|
||||
name='tool',
|
||||
description='tool-description',
|
||||
inputSchema={
|
||||
'type': 'array',
|
||||
'items': {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'key1': {
|
||||
'type': 'string',
|
||||
},
|
||||
'key2': {
|
||||
'type': 'number',
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
),
|
||||
]
|
||||
result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
|
||||
assert result == [
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
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'),
|
||||
},
|
||||
),
|
||||
),
|
||||
),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_any_of_conversion():
|
||||
"""Test conversion of MCP tools with any_of to Gemini tools."""
|
||||
mcp_tools = [
|
||||
mcp_types.Tool(
|
||||
name='tool',
|
||||
description='tool-description',
|
||||
inputSchema={
|
||||
'type': 'object',
|
||||
'any_of': [
|
||||
{
|
||||
'type': 'string',
|
||||
},
|
||||
{
|
||||
'type': 'number',
|
||||
},
|
||||
],
|
||||
},
|
||||
),
|
||||
]
|
||||
result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
|
||||
assert result == [
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(
|
||||
name='tool',
|
||||
description='tool-description',
|
||||
parameters=types.Schema(
|
||||
type='OBJECT',
|
||||
any_of=[
|
||||
types.Schema(type='STRING'),
|
||||
types.Schema(type='NUMBER'),
|
||||
],
|
||||
),
|
||||
),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_properties_conversion():
|
||||
"""Test conversion of MCP tools with properties to Gemini tools."""
|
||||
mcp_tools = [
|
||||
mcp_types.Tool(
|
||||
name='tool',
|
||||
description='tool-description',
|
||||
inputSchema={
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'key1': {
|
||||
'type': 'string',
|
||||
},
|
||||
'key2': {
|
||||
'type': 'number',
|
||||
},
|
||||
},
|
||||
},
|
||||
),
|
||||
]
|
||||
result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
|
||||
assert result == [
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(
|
||||
name='tool',
|
||||
description='tool-description',
|
||||
parameters=types.Schema(
|
||||
type='OBJECT',
|
||||
properties={
|
||||
'key1': types.Schema(type='STRING'),
|
||||
'key2': types.Schema(type='NUMBER'),
|
||||
},
|
||||
),
|
||||
),
|
||||
],
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1,201 @@
|
||||
# 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 _asyncio
|
||||
import pytest
|
||||
from ... import _extra_utils
|
||||
from ... import types
|
||||
from ..._adapters import McpToGenAiToolAdapter
|
||||
|
||||
try:
|
||||
from mcp import types as mcp_types
|
||||
from mcp import ClientSession as McpClientSession
|
||||
except ImportError as e:
|
||||
import sys
|
||||
|
||||
if sys.version_info < (3, 10):
|
||||
raise ImportError(
|
||||
'MCP Tool requires Python 3.10 or above. Please upgrade your Python'
|
||||
' version.'
|
||||
) from e
|
||||
else:
|
||||
raise e
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_empty_config_dict():
|
||||
"""Test conversion of empty GenerateContentConfigDict to parsed config."""
|
||||
config = {}
|
||||
parsed_config, mcp_to_genai_tool_adapters = (
|
||||
await _extra_utils.parse_config_for_mcp_sessions(config)
|
||||
)
|
||||
assert parsed_config == None
|
||||
assert not mcp_to_genai_tool_adapters
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_empty_config_object():
|
||||
"""Test conversion of empty GenerateContentConfig to parsed config."""
|
||||
config = types.GenerateContentConfig()
|
||||
parsed_config, mcp_to_genai_tool_adapters = (
|
||||
await _extra_utils.parse_config_for_mcp_sessions(config)
|
||||
)
|
||||
assert config is not parsed_config # config is not modified
|
||||
assert not mcp_to_genai_tool_adapters
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_config_object_with_tools():
|
||||
"""Test conversion of GenerateContentConfig with tools to parsed config."""
|
||||
|
||||
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='get_weather_2',
|
||||
description='Different tool to get the weather.',
|
||||
inputSchema={
|
||||
'type': 'object',
|
||||
'properties': {'location': {'type': 'string'}},
|
||||
},
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
mock_session_instance = MockMcpClientSession()
|
||||
config = types.GenerateContentConfig(tools=[mock_session_instance])
|
||||
parsed_config, mcp_to_genai_tool_adapters = (
|
||||
await _extra_utils.parse_config_for_mcp_sessions(config)
|
||||
)
|
||||
assert len(config.tools) == 1
|
||||
assert config.tools[0] is mock_session_instance
|
||||
assert config is not parsed_config # config is not modified
|
||||
assert len(mcp_to_genai_tool_adapters) == 2
|
||||
assert mcp_to_genai_tool_adapters.keys() == {
|
||||
'get_weather',
|
||||
'get_weather_2',
|
||||
}
|
||||
assert isinstance(
|
||||
mcp_to_genai_tool_adapters['get_weather'], McpToGenAiToolAdapter
|
||||
)
|
||||
assert isinstance(
|
||||
mcp_to_genai_tool_adapters['get_weather_2'], McpToGenAiToolAdapter
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_config_object_with_tools_complex_type():
|
||||
"""Test conversion of GenerateContentConfig with tools to parsed config."""
|
||||
|
||||
class MockMcpClientSession(McpClientSession):
|
||||
|
||||
def __init__(self):
|
||||
self._read_stream = None
|
||||
self._write_stream = None
|
||||
self._future = _asyncio.Future() # This object cannot be pickled.
|
||||
|
||||
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='get_weather_2',
|
||||
description='Different tool to get the weather.',
|
||||
inputSchema={
|
||||
'type': 'object',
|
||||
'properties': {'location': {'type': 'string'}},
|
||||
},
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
mock_session_instance = MockMcpClientSession()
|
||||
config = types.GenerateContentConfig(tools=[mock_session_instance])
|
||||
parsed_config, mcp_to_genai_tool_adapters = (
|
||||
await _extra_utils.parse_config_for_mcp_sessions(config)
|
||||
)
|
||||
assert len(config.tools) == 1
|
||||
assert config.tools[0] is mock_session_instance
|
||||
assert config is not parsed_config # config is not modified
|
||||
assert len(mcp_to_genai_tool_adapters) == 2
|
||||
assert mcp_to_genai_tool_adapters.keys() == {
|
||||
'get_weather',
|
||||
'get_weather_2',
|
||||
}
|
||||
assert isinstance(
|
||||
mcp_to_genai_tool_adapters['get_weather'], McpToGenAiToolAdapter
|
||||
)
|
||||
assert isinstance(
|
||||
mcp_to_genai_tool_adapters['get_weather_2'], McpToGenAiToolAdapter
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_config_object_with_non_mcp_tools():
|
||||
"""Test conversion of GenerateContentConfig with regular tools to parsed config."""
|
||||
|
||||
config = types.GenerateContentConfig(
|
||||
tools=[
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
{'name': 'tool-1', 'description': 'tool-1-description'}
|
||||
]
|
||||
),
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
{'name': 'tool-2', 'description': 'tool-2-description'}
|
||||
]
|
||||
),
|
||||
]
|
||||
)
|
||||
parsed_config, mcp_to_genai_tool_adapters = (
|
||||
await _extra_utils.parse_config_for_mcp_sessions(config)
|
||||
)
|
||||
assert len(config.tools) == 2
|
||||
assert config is not parsed_config # config is not modified
|
||||
assert mcp_to_genai_tool_adapters == {}
|
||||
assert parsed_config.tools == [
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
{'name': 'tool-1', 'description': 'tool-1-description'}
|
||||
]
|
||||
),
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
{'name': 'tool-2', 'description': 'tool-2-description'}
|
||||
]
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1,130 @@
|
||||
# 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 re
|
||||
import pytest
|
||||
from ... import _extra_utils
|
||||
from ... import types
|
||||
from ..._adapters import McpToGenAiToolAdapter
|
||||
|
||||
try:
|
||||
from mcp import types as mcp_types
|
||||
from mcp import ClientSession as McpClientSession
|
||||
except ImportError as e:
|
||||
import sys
|
||||
|
||||
if sys.version_info < (3, 10):
|
||||
raise ImportError(
|
||||
'MCP Tool requires Python 3.10 or above. Please upgrade your Python'
|
||||
' version.'
|
||||
) from e
|
||||
else:
|
||||
raise e
|
||||
|
||||
|
||||
def test_parse_empty_config_dict():
|
||||
"""Test conversion of empty GenerateContentConfigDict to parsed config."""
|
||||
config = {}
|
||||
parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
|
||||
assert parsed_config == None
|
||||
|
||||
|
||||
def test_parse_empty_config_object():
|
||||
"""Test conversion of empty GenerateContentConfig to parsed config."""
|
||||
config = types.GenerateContentConfig()
|
||||
parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
|
||||
assert config is not parsed_config # config is not modified
|
||||
|
||||
|
||||
def test_parse_config_object_with_tools():
|
||||
"""Test conversion of GenerateContentConfig with tools to parsed config."""
|
||||
|
||||
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='get_weather_2',
|
||||
description='Different tool to get the weather.',
|
||||
inputSchema={
|
||||
'type': 'object',
|
||||
'properties': {'location': {'type': 'string'}},
|
||||
},
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
mock_session_instance = MockMcpClientSession()
|
||||
config = types.GenerateContentConfig(tools=[mock_session_instance])
|
||||
parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
|
||||
assert config.http_options is None
|
||||
assert config is not parsed_config # config is not modified
|
||||
assert re.match(
|
||||
r'mcp_used/\d+\.\d+\.\d+',
|
||||
parsed_config.http_options.headers['x-goog-api-client'],
|
||||
)
|
||||
|
||||
|
||||
def test_parse_config_object_with_tools_and_existing_headers():
|
||||
"""Test conversion of GenerateContentConfig with tools and existing headers to parsed config."""
|
||||
|
||||
config = types.GenerateContentConfig(
|
||||
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='get_weather_2',
|
||||
description='Different tool to get the weather.',
|
||||
inputSchema={
|
||||
'type': 'object',
|
||||
'properties': {'location': {'type': 'string'}},
|
||||
},
|
||||
),
|
||||
],
|
||||
http_options=types.HttpOptions(
|
||||
headers={
|
||||
'x-goog-api-client': 'google-genai-sdk/1.0.0 gl-python/1.0.0'
|
||||
}
|
||||
),
|
||||
)
|
||||
parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
|
||||
assert not re.match(
|
||||
r'mcp_used/\d+\.\d+\.\d+',
|
||||
config.http_options.headers['x-goog-api-client'],
|
||||
)
|
||||
assert config is not parsed_config # config is not modified
|
||||
assert re.match(
|
||||
r'google-genai-sdk/1.0.0 gl-python/1.0.0 mcp_used/\d+\.\d+\.\d+',
|
||||
parsed_config.http_options.headers['x-goog-api-client'],
|
||||
)
|
||||
@@ -0,0 +1,72 @@
|
||||
# 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 importlib.metadata import version
|
||||
import re
|
||||
import typing
|
||||
from typing import Any
|
||||
|
||||
from ... import _mcp_utils
|
||||
from ... import types
|
||||
|
||||
_is_mcp_imported = False
|
||||
if typing.TYPE_CHECKING:
|
||||
import mcp
|
||||
|
||||
_is_mcp_imported = True
|
||||
else:
|
||||
try:
|
||||
import mcp
|
||||
|
||||
_is_mcp_imported = True
|
||||
except ImportError:
|
||||
_is_mcp_imported = False
|
||||
|
||||
|
||||
def test_set_mcp_usage_header_from_empty_dict():
|
||||
if not _is_mcp_imported:
|
||||
return
|
||||
"""Test whether the MCP usage header is set correctly from an empty dict."""
|
||||
headers = {}
|
||||
_mcp_utils.set_mcp_usage_header(headers)
|
||||
assert re.match(r'mcp_used/\d+\.\d+\.\d+', headers['x-goog-api-client'])
|
||||
|
||||
|
||||
def test_set_mcp_usage_header_with_existing_header():
|
||||
if not _is_mcp_imported:
|
||||
return
|
||||
"""Test whether the MCP usage header is set correctly from an existing header."""
|
||||
headers = {'x-goog-api-client': 'google-genai-sdk/1.0.0 gl-python/1.0.0'}
|
||||
_mcp_utils.set_mcp_usage_header(headers)
|
||||
assert re.match(
|
||||
r'google-genai-sdk/1.0.0 gl-python/1.0.0 mcp_used/\d+\.\d+\.\d+',
|
||||
headers['x-goog-api-client'],
|
||||
)
|
||||
|
||||
|
||||
def test_set_mcp_usage_header_with_existing_mcp_header():
|
||||
if not _is_mcp_imported:
|
||||
return
|
||||
"""Test whether the MCP usage header is set correctly from an existing MCP header."""
|
||||
headers = {
|
||||
'x-goog-api-client': (
|
||||
'google-genai-sdk/1.0.0 gl-python/1.0.0 mcp_used/1.0.0'
|
||||
)
|
||||
}
|
||||
_mcp_utils.set_mcp_usage_header(headers)
|
||||
assert re.match(
|
||||
r'google-genai-sdk/1.0.0 gl-python/1.0.0 mcp_used/\d+\.\d+\.\d+',
|
||||
headers['x-goog-api-client'],
|
||||
)
|
||||
@@ -0,0 +1,17 @@
|
||||
# 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 models module."""
|
||||
@@ -0,0 +1,8 @@
|
||||
VERTEX_HTTP_OPTIONS = {
|
||||
'api_version': 'v1beta1',
|
||||
'base_url': 'https://us-central1-aiplatform.googleapis.com/',
|
||||
}
|
||||
MLDEV_HTTP_OPTIONS = {
|
||||
'api_version': 'v1beta',
|
||||
'base_url': 'https://generativelanguage.googleapis.com/',
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
# 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 pytest
|
||||
from ... import _transformers as t
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
from . import constants
|
||||
|
||||
_COMPUTE_TOKENS_PARAMS = types._ComputeTokensParameters(
|
||||
model='gemini-2.5-flash',
|
||||
contents=[t.t_content('Tell me a story in 300 words.')],
|
||||
)
|
||||
_COMPUTE_TOKENS_PARAMS_VERTEX_CUSTOM_URL = types._ComputeTokensParameters(
|
||||
model='gemini-2.5-flash',
|
||||
contents=[t.t_content('Tell me a story in 300 words.')],
|
||||
config={'http_options': constants.VERTEX_HTTP_OPTIONS},
|
||||
)
|
||||
_COMPUTE_TOKENS_PARAMS_MLDEV_CUSTOM_URL = types._ComputeTokensParameters(
|
||||
model='gemini-2.5-flash',
|
||||
contents=[t.t_content('Tell me a story in 300 words.')],
|
||||
config={'http_options': constants.MLDEV_HTTP_OPTIONS},
|
||||
)
|
||||
_UNICODE_STRING = '这是一条unicode测试🤪❤★'
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_compute_tokens',
|
||||
exception_if_mldev='only supported in',
|
||||
parameters=types._ComputeTokensParameters(
|
||||
model='gemini-2.5-flash',
|
||||
contents=[t.t_content('Tell me a story in 300 words.')],
|
||||
),
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_compute_tokens_vertex_custom_url',
|
||||
parameters=_COMPUTE_TOKENS_PARAMS_VERTEX_CUSTOM_URL,
|
||||
exception_if_mldev='only supported in',
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_compute_tokens_mldev_custom_url',
|
||||
parameters=_COMPUTE_TOKENS_PARAMS_MLDEV_CUSTOM_URL,
|
||||
exception_if_vertex='404',
|
||||
exception_if_mldev='only supported in',
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_compute_tokens_unicode',
|
||||
exception_if_mldev='only supported in',
|
||||
parameters=types._ComputeTokensParameters(
|
||||
model='gemini-2.5-flash', contents=[t.t_content(_UNICODE_STRING)]
|
||||
),
|
||||
),
|
||||
]
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='models.compute_tokens',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
def test_token_bytes_deserialization(client):
|
||||
if client._api_client.vertexai:
|
||||
response = client.models.compute_tokens(
|
||||
model=_COMPUTE_TOKENS_PARAMS.model,
|
||||
contents=_UNICODE_STRING,
|
||||
)
|
||||
decoded_tokens = b''.join(response.tokens_info[0].tokens)
|
||||
assert (
|
||||
decoded_tokens
|
||||
== b'\xe8\xbf\x99\xe6\x98\xaf\xe4\xb8\x80\xe6\x9d\xa1unicode\xe6\xb5\x8b\xe8\xaf\x95\xf0\x9f\xa4\xaa\xe2\x9d\xa4\xe2\x98\x85'
|
||||
)
|
||||
assert decoded_tokens.decode('utf-8') == _UNICODE_STRING
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async(client):
|
||||
if client._api_client.vertexai:
|
||||
response = await client.aio.models.compute_tokens(
|
||||
model=_COMPUTE_TOKENS_PARAMS.model,
|
||||
contents=_COMPUTE_TOKENS_PARAMS.contents,
|
||||
)
|
||||
assert response
|
||||
else:
|
||||
with pytest.raises(Exception):
|
||||
await client.aio.models.compute_tokens(
|
||||
model=_COMPUTE_TOKENS_PARAMS.model,
|
||||
contents=_COMPUTE_TOKENS_PARAMS.contents,
|
||||
)
|
||||
|
||||
|
||||
def test_different_model_names(client):
|
||||
if client._api_client.vertexai:
|
||||
response1 = client.models.compute_tokens(
|
||||
model='gemini-2.5-flash', contents=_COMPUTE_TOKENS_PARAMS.contents
|
||||
)
|
||||
assert response1
|
||||
response3 = client.models.compute_tokens(
|
||||
model='publishers/google/models/gemini-2.5-flash',
|
||||
contents=_COMPUTE_TOKENS_PARAMS.contents,
|
||||
)
|
||||
assert response3
|
||||
response4 = client.models.compute_tokens(
|
||||
model='projects/vertexsdk/locations/us-central1/publishers/google/models/gemini-2.5-flash',
|
||||
contents=_COMPUTE_TOKENS_PARAMS.contents,
|
||||
)
|
||||
assert response4
|
||||
@@ -0,0 +1,159 @@
|
||||
# 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 copy
|
||||
import pytest
|
||||
from ... import _transformers as t
|
||||
from ... import errors
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
from . import constants
|
||||
|
||||
_COUNT_TOKENS_PARAMS = types._CountTokensParameters(
|
||||
model='gemini-2.5-flash',
|
||||
contents=[t.t_content('Tell me a story in 300 words.')],
|
||||
)
|
||||
|
||||
_COUNT_TOKENS_PARAMS_WITH_SYSTEM_INSTRUCTION = copy.deepcopy(
|
||||
_COUNT_TOKENS_PARAMS
|
||||
)
|
||||
_COUNT_TOKENS_PARAMS_WITH_SYSTEM_INSTRUCTION.config = {
|
||||
'system_instruction': t.t_content('you are a chatbot.')
|
||||
}
|
||||
|
||||
_COUNT_TOKENS_PARAMS_WITH_TOOLS = copy.deepcopy(_COUNT_TOKENS_PARAMS)
|
||||
_COUNT_TOKENS_PARAMS_WITH_TOOLS.config = {
|
||||
'tools': [{'google_search_retrieval': {}}]
|
||||
}
|
||||
|
||||
_COUNT_TOKENS_PARAMS_WITH_GENERATION_CONFIG = copy.deepcopy(
|
||||
_COUNT_TOKENS_PARAMS
|
||||
)
|
||||
_COUNT_TOKENS_PARAMS_WITH_GENERATION_CONFIG.config = {
|
||||
'generation_config': {'max_output_tokens': 50}
|
||||
}
|
||||
|
||||
_COUNT_TOKENS_PARAMS_VERTEX_CUSTOM_URL = copy.deepcopy(_COUNT_TOKENS_PARAMS)
|
||||
_COUNT_TOKENS_PARAMS_VERTEX_CUSTOM_URL.config = {
|
||||
'http_options': constants.VERTEX_HTTP_OPTIONS
|
||||
}
|
||||
_COUNT_TOKENS_PARAMS_MLDEV_CUSTOM_URL = copy.deepcopy(_COUNT_TOKENS_PARAMS)
|
||||
_COUNT_TOKENS_PARAMS_MLDEV_CUSTOM_URL.config = {
|
||||
'http_options': constants.MLDEV_HTTP_OPTIONS
|
||||
}
|
||||
|
||||
|
||||
# TODO(b/378952792): MLDev count_tokens needs to merge contents and model
|
||||
# param into generateContentRequest field.
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_count_tokens',
|
||||
parameters=_COUNT_TOKENS_PARAMS,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_count_tokens_vertex_custom_url',
|
||||
parameters=_COUNT_TOKENS_PARAMS_VERTEX_CUSTOM_URL,
|
||||
exception_if_mldev='404',
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_count_tokens_mldev_custom_url',
|
||||
parameters=_COUNT_TOKENS_PARAMS_MLDEV_CUSTOM_URL,
|
||||
exception_if_vertex='404',
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_count_tokens_with_system_instruction',
|
||||
exception_if_mldev='not supported',
|
||||
parameters=_COUNT_TOKENS_PARAMS_WITH_SYSTEM_INSTRUCTION,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_count_tokens_with_tools',
|
||||
exception_if_mldev='not supported',
|
||||
parameters=_COUNT_TOKENS_PARAMS_WITH_TOOLS,
|
||||
),
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_count_tokens_with_generation_config',
|
||||
exception_if_mldev='not supported',
|
||||
parameters=_COUNT_TOKENS_PARAMS_WITH_GENERATION_CONFIG,
|
||||
),
|
||||
]
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='models.count_tokens',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async(client):
|
||||
response = await client.aio.models.count_tokens(
|
||||
model=_COUNT_TOKENS_PARAMS.model, contents=_COUNT_TOKENS_PARAMS.contents
|
||||
)
|
||||
assert response
|
||||
|
||||
|
||||
def test_different_model_names(client):
|
||||
if client._api_client.vertexai:
|
||||
response1 = client.models.count_tokens(
|
||||
model='gemini-2.5-flash', contents=_COUNT_TOKENS_PARAMS.contents
|
||||
)
|
||||
assert response1
|
||||
response3 = client.models.count_tokens(
|
||||
model='publishers/google/models/gemini-2.5-flash',
|
||||
contents=_COUNT_TOKENS_PARAMS.contents,
|
||||
)
|
||||
assert response3
|
||||
response4 = client.models.count_tokens(
|
||||
model='projects/vertexsdk/locations/us-central1/publishers/google/models/gemini-2.5-flash',
|
||||
contents=_COUNT_TOKENS_PARAMS.contents,
|
||||
)
|
||||
assert response4
|
||||
else:
|
||||
response1 = client.models.count_tokens(
|
||||
model='gemini-2.5-flash', contents=_COUNT_TOKENS_PARAMS.contents
|
||||
)
|
||||
assert response1
|
||||
response2 = client.models.count_tokens(
|
||||
model='models/gemini-2.5-flash', contents=_COUNT_TOKENS_PARAMS.contents
|
||||
)
|
||||
assert response2
|
||||
|
||||
|
||||
def test_extra_body(client):
|
||||
config = {
|
||||
'http_options': {
|
||||
'extra_body': {
|
||||
'systemInstruction': {
|
||||
'parts': [{'text': 'you are a chatbot.'}],
|
||||
'role': 'user',
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if client._api_client.vertexai:
|
||||
response = client.models.count_tokens(
|
||||
model=_COUNT_TOKENS_PARAMS.model,
|
||||
contents=_COUNT_TOKENS_PARAMS.contents,
|
||||
config=config,
|
||||
)
|
||||
assert response.total_tokens
|
||||
else:
|
||||
with pytest.raises(errors.ClientError):
|
||||
client.models.count_tokens(
|
||||
model=_COUNT_TOKENS_PARAMS.model,
|
||||
contents=_COUNT_TOKENS_PARAMS.contents,
|
||||
config=config,
|
||||
)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user