google-genai 1.53.0__py3-none-any.whl → 1.55.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- google/genai/__init__.py +1 -0
- google/genai/_api_client.py +6 -6
- google/genai/_interactions/__init__.py +117 -0
- google/genai/_interactions/_base_client.py +2019 -0
- google/genai/_interactions/_client.py +511 -0
- google/genai/_interactions/_compat.py +234 -0
- google/genai/_interactions/_constants.py +29 -0
- google/genai/_interactions/_exceptions.py +122 -0
- google/genai/_interactions/_files.py +139 -0
- google/genai/_interactions/_models.py +873 -0
- google/genai/_interactions/_qs.py +165 -0
- google/genai/_interactions/_resource.py +58 -0
- google/genai/_interactions/_response.py +847 -0
- google/genai/_interactions/_streaming.py +354 -0
- google/genai/_interactions/_types.py +276 -0
- google/genai/_interactions/_utils/__init__.py +79 -0
- google/genai/_interactions/_utils/_compat.py +61 -0
- google/genai/_interactions/_utils/_datetime_parse.py +151 -0
- google/genai/_interactions/_utils/_logs.py +40 -0
- google/genai/_interactions/_utils/_proxy.py +80 -0
- google/genai/_interactions/_utils/_reflection.py +57 -0
- google/genai/_interactions/_utils/_resources_proxy.py +39 -0
- google/genai/_interactions/_utils/_streams.py +27 -0
- google/genai/_interactions/_utils/_sync.py +73 -0
- google/genai/_interactions/_utils/_transform.py +472 -0
- google/genai/_interactions/_utils/_typing.py +172 -0
- google/genai/_interactions/_utils/_utils.py +437 -0
- google/genai/_interactions/_version.py +18 -0
- google/genai/_interactions/resources/__init__.py +34 -0
- google/genai/_interactions/resources/interactions.py +1350 -0
- google/genai/_interactions/types/__init__.py +107 -0
- google/genai/_interactions/types/allowed_tools.py +33 -0
- google/genai/_interactions/types/allowed_tools_param.py +35 -0
- google/genai/_interactions/types/annotation.py +42 -0
- google/genai/_interactions/types/annotation_param.py +42 -0
- google/genai/_interactions/types/audio_content.py +38 -0
- google/genai/_interactions/types/audio_content_param.py +45 -0
- google/genai/_interactions/types/audio_mime_type.py +25 -0
- google/genai/_interactions/types/audio_mime_type_param.py +27 -0
- google/genai/_interactions/types/code_execution_call_arguments.py +33 -0
- google/genai/_interactions/types/code_execution_call_arguments_param.py +32 -0
- google/genai/_interactions/types/code_execution_call_content.py +37 -0
- google/genai/_interactions/types/code_execution_call_content_param.py +37 -0
- google/genai/_interactions/types/code_execution_result_content.py +42 -0
- google/genai/_interactions/types/code_execution_result_content_param.py +41 -0
- google/genai/_interactions/types/content_delta.py +358 -0
- google/genai/_interactions/types/content_start.py +79 -0
- google/genai/_interactions/types/content_stop.py +35 -0
- google/genai/_interactions/types/deep_research_agent_config.py +33 -0
- google/genai/_interactions/types/deep_research_agent_config_param.py +32 -0
- google/genai/_interactions/types/document_content.py +36 -0
- google/genai/_interactions/types/document_content_param.py +43 -0
- google/genai/_interactions/types/dynamic_agent_config.py +44 -0
- google/genai/_interactions/types/dynamic_agent_config_param.py +33 -0
- google/genai/_interactions/types/error_event.py +46 -0
- google/genai/_interactions/types/file_search_result_content.py +46 -0
- google/genai/_interactions/types/file_search_result_content_param.py +46 -0
- google/genai/_interactions/types/function.py +38 -0
- google/genai/_interactions/types/function_call_content.py +39 -0
- google/genai/_interactions/types/function_call_content_param.py +39 -0
- google/genai/_interactions/types/function_param.py +37 -0
- google/genai/_interactions/types/function_result_content.py +52 -0
- google/genai/_interactions/types/function_result_content_param.py +54 -0
- google/genai/_interactions/types/generation_config.py +57 -0
- google/genai/_interactions/types/generation_config_param.py +59 -0
- google/genai/_interactions/types/google_search_call_arguments.py +29 -0
- google/genai/_interactions/types/google_search_call_arguments_param.py +31 -0
- google/genai/_interactions/types/google_search_call_content.py +37 -0
- google/genai/_interactions/types/google_search_call_content_param.py +37 -0
- google/genai/_interactions/types/google_search_result.py +35 -0
- google/genai/_interactions/types/google_search_result_content.py +43 -0
- google/genai/_interactions/types/google_search_result_content_param.py +44 -0
- google/genai/_interactions/types/google_search_result_param.py +35 -0
- google/genai/_interactions/types/image_content.py +41 -0
- google/genai/_interactions/types/image_content_param.py +48 -0
- google/genai/_interactions/types/image_mime_type.py +23 -0
- google/genai/_interactions/types/image_mime_type_param.py +25 -0
- google/genai/_interactions/types/interaction.py +165 -0
- google/genai/_interactions/types/interaction_create_params.py +212 -0
- google/genai/_interactions/types/interaction_event.py +37 -0
- google/genai/_interactions/types/interaction_get_params.py +46 -0
- google/genai/_interactions/types/interaction_sse_event.py +32 -0
- google/genai/_interactions/types/interaction_status_update.py +37 -0
- google/genai/_interactions/types/mcp_server_tool_call_content.py +42 -0
- google/genai/_interactions/types/mcp_server_tool_call_content_param.py +42 -0
- google/genai/_interactions/types/mcp_server_tool_result_content.py +52 -0
- google/genai/_interactions/types/mcp_server_tool_result_content_param.py +54 -0
- google/genai/_interactions/types/model.py +36 -0
- google/genai/_interactions/types/model_param.py +38 -0
- google/genai/_interactions/types/speech_config.py +35 -0
- google/genai/_interactions/types/speech_config_param.py +35 -0
- google/genai/_interactions/types/text_content.py +37 -0
- google/genai/_interactions/types/text_content_param.py +38 -0
- google/genai/_interactions/types/thinking_level.py +22 -0
- google/genai/_interactions/types/thought_content.py +41 -0
- google/genai/_interactions/types/thought_content_param.py +47 -0
- google/genai/_interactions/types/tool.py +100 -0
- google/genai/_interactions/types/tool_choice.py +26 -0
- google/genai/_interactions/types/tool_choice_config.py +28 -0
- google/genai/_interactions/types/tool_choice_config_param.py +29 -0
- google/genai/_interactions/types/tool_choice_param.py +28 -0
- google/genai/_interactions/types/tool_choice_type.py +22 -0
- google/genai/_interactions/types/tool_param.py +97 -0
- google/genai/_interactions/types/turn.py +76 -0
- google/genai/_interactions/types/turn_param.py +73 -0
- google/genai/_interactions/types/url_context_call_arguments.py +29 -0
- google/genai/_interactions/types/url_context_call_arguments_param.py +31 -0
- google/genai/_interactions/types/url_context_call_content.py +37 -0
- google/genai/_interactions/types/url_context_call_content_param.py +37 -0
- google/genai/_interactions/types/url_context_result.py +33 -0
- google/genai/_interactions/types/url_context_result_content.py +43 -0
- google/genai/_interactions/types/url_context_result_content_param.py +44 -0
- google/genai/_interactions/types/url_context_result_param.py +32 -0
- google/genai/_interactions/types/usage.py +106 -0
- google/genai/_interactions/types/usage_param.py +106 -0
- google/genai/_interactions/types/video_content.py +41 -0
- google/genai/_interactions/types/video_content_param.py +48 -0
- google/genai/_interactions/types/video_mime_type.py +36 -0
- google/genai/_interactions/types/video_mime_type_param.py +38 -0
- google/genai/_live_converters.py +34 -3
- google/genai/_tokens_converters.py +5 -0
- google/genai/batches.py +62 -55
- google/genai/client.py +223 -0
- google/genai/errors.py +16 -1
- google/genai/file_search_stores.py +60 -60
- google/genai/files.py +56 -56
- google/genai/interactions.py +17 -0
- google/genai/live.py +4 -3
- google/genai/models.py +15 -3
- google/genai/tests/__init__.py +21 -0
- google/genai/tests/afc/__init__.py +21 -0
- google/genai/tests/afc/test_convert_if_exist_pydantic_model.py +309 -0
- google/genai/tests/afc/test_convert_number_values_for_function_call_args.py +63 -0
- google/genai/tests/afc/test_find_afc_incompatible_tool_indexes.py +240 -0
- google/genai/tests/afc/test_generate_content_stream_afc.py +530 -0
- google/genai/tests/afc/test_generate_content_stream_afc_thoughts.py +77 -0
- google/genai/tests/afc/test_get_function_map.py +176 -0
- google/genai/tests/afc/test_get_function_response_parts.py +277 -0
- google/genai/tests/afc/test_get_max_remote_calls_for_afc.py +130 -0
- google/genai/tests/afc/test_invoke_function_from_dict_args.py +241 -0
- google/genai/tests/afc/test_raise_error_for_afc_incompatible_config.py +159 -0
- google/genai/tests/afc/test_should_append_afc_history.py +53 -0
- google/genai/tests/afc/test_should_disable_afc.py +214 -0
- google/genai/tests/batches/__init__.py +17 -0
- google/genai/tests/batches/test_cancel.py +77 -0
- google/genai/tests/batches/test_create.py +78 -0
- google/genai/tests/batches/test_create_with_bigquery.py +113 -0
- google/genai/tests/batches/test_create_with_file.py +82 -0
- google/genai/tests/batches/test_create_with_gcs.py +125 -0
- google/genai/tests/batches/test_create_with_inlined_requests.py +255 -0
- google/genai/tests/batches/test_delete.py +86 -0
- google/genai/tests/batches/test_embedding.py +157 -0
- google/genai/tests/batches/test_get.py +78 -0
- google/genai/tests/batches/test_list.py +79 -0
- google/genai/tests/caches/__init__.py +17 -0
- google/genai/tests/caches/constants.py +29 -0
- google/genai/tests/caches/test_create.py +210 -0
- google/genai/tests/caches/test_create_custom_url.py +105 -0
- google/genai/tests/caches/test_delete.py +54 -0
- google/genai/tests/caches/test_delete_custom_url.py +52 -0
- google/genai/tests/caches/test_get.py +94 -0
- google/genai/tests/caches/test_get_custom_url.py +52 -0
- google/genai/tests/caches/test_list.py +68 -0
- google/genai/tests/caches/test_update.py +70 -0
- google/genai/tests/caches/test_update_custom_url.py +58 -0
- google/genai/tests/chats/__init__.py +1 -0
- google/genai/tests/chats/test_get_history.py +597 -0
- google/genai/tests/chats/test_send_message.py +844 -0
- google/genai/tests/chats/test_validate_response.py +90 -0
- google/genai/tests/client/__init__.py +17 -0
- google/genai/tests/client/test_async_stream.py +427 -0
- google/genai/tests/client/test_client_close.py +197 -0
- google/genai/tests/client/test_client_initialization.py +1687 -0
- google/genai/tests/client/test_client_requests.py +355 -0
- google/genai/tests/client/test_custom_client.py +77 -0
- google/genai/tests/client/test_http_options.py +178 -0
- google/genai/tests/client/test_replay_client_equality.py +168 -0
- google/genai/tests/client/test_retries.py +846 -0
- google/genai/tests/client/test_upload_errors.py +136 -0
- google/genai/tests/common/__init__.py +17 -0
- google/genai/tests/common/test_common.py +954 -0
- google/genai/tests/conftest.py +162 -0
- google/genai/tests/documents/__init__.py +17 -0
- google/genai/tests/documents/test_delete.py +51 -0
- google/genai/tests/documents/test_get.py +85 -0
- google/genai/tests/documents/test_list.py +72 -0
- google/genai/tests/errors/__init__.py +1 -0
- google/genai/tests/errors/test_api_error.py +417 -0
- google/genai/tests/file_search_stores/__init__.py +17 -0
- google/genai/tests/file_search_stores/test_create.py +66 -0
- google/genai/tests/file_search_stores/test_delete.py +64 -0
- google/genai/tests/file_search_stores/test_get.py +94 -0
- google/genai/tests/file_search_stores/test_import_file.py +112 -0
- google/genai/tests/file_search_stores/test_list.py +57 -0
- google/genai/tests/file_search_stores/test_upload_to_file_search_store.py +141 -0
- google/genai/tests/files/__init__.py +17 -0
- google/genai/tests/files/test_delete.py +46 -0
- google/genai/tests/files/test_download.py +85 -0
- google/genai/tests/files/test_get.py +46 -0
- google/genai/tests/files/test_list.py +72 -0
- google/genai/tests/files/test_upload.py +255 -0
- google/genai/tests/imports/test_no_optional_imports.py +28 -0
- google/genai/tests/interactions/__init__.py +0 -0
- google/genai/tests/interactions/test_integration.py +80 -0
- google/genai/tests/live/__init__.py +16 -0
- google/genai/tests/live/test_live.py +2177 -0
- google/genai/tests/live/test_live_music.py +362 -0
- google/genai/tests/live/test_live_response.py +163 -0
- google/genai/tests/live/test_send_client_content.py +147 -0
- google/genai/tests/live/test_send_realtime_input.py +268 -0
- google/genai/tests/live/test_send_tool_response.py +222 -0
- google/genai/tests/local_tokenizer/__init__.py +17 -0
- google/genai/tests/local_tokenizer/test_local_tokenizer.py +343 -0
- google/genai/tests/local_tokenizer/test_local_tokenizer_loader.py +235 -0
- google/genai/tests/mcp/__init__.py +17 -0
- google/genai/tests/mcp/test_has_mcp_tool_usage.py +89 -0
- google/genai/tests/mcp/test_mcp_to_gemini_tools.py +191 -0
- google/genai/tests/mcp/test_parse_config_for_mcp_sessions.py +201 -0
- google/genai/tests/mcp/test_parse_config_for_mcp_usage.py +130 -0
- google/genai/tests/mcp/test_set_mcp_usage_header.py +72 -0
- google/genai/tests/models/__init__.py +17 -0
- google/genai/tests/models/constants.py +8 -0
- google/genai/tests/models/test_compute_tokens.py +120 -0
- google/genai/tests/models/test_count_tokens.py +159 -0
- google/genai/tests/models/test_delete.py +107 -0
- google/genai/tests/models/test_edit_image.py +264 -0
- google/genai/tests/models/test_embed_content.py +94 -0
- google/genai/tests/models/test_function_call_streaming.py +442 -0
- google/genai/tests/models/test_generate_content.py +2502 -0
- google/genai/tests/models/test_generate_content_cached_content.py +132 -0
- google/genai/tests/models/test_generate_content_config_zero_value.py +103 -0
- google/genai/tests/models/test_generate_content_from_apikey.py +44 -0
- google/genai/tests/models/test_generate_content_http_options.py +40 -0
- google/genai/tests/models/test_generate_content_image_generation.py +143 -0
- google/genai/tests/models/test_generate_content_mcp.py +343 -0
- google/genai/tests/models/test_generate_content_media_resolution.py +97 -0
- google/genai/tests/models/test_generate_content_model.py +139 -0
- google/genai/tests/models/test_generate_content_part.py +821 -0
- google/genai/tests/models/test_generate_content_thought.py +76 -0
- google/genai/tests/models/test_generate_content_tools.py +1761 -0
- google/genai/tests/models/test_generate_images.py +191 -0
- google/genai/tests/models/test_generate_videos.py +759 -0
- google/genai/tests/models/test_get.py +104 -0
- google/genai/tests/models/test_list.py +233 -0
- google/genai/tests/models/test_recontext_image.py +189 -0
- google/genai/tests/models/test_segment_image.py +148 -0
- google/genai/tests/models/test_update.py +95 -0
- google/genai/tests/models/test_upscale_image.py +157 -0
- google/genai/tests/operations/__init__.py +17 -0
- google/genai/tests/operations/test_get.py +38 -0
- google/genai/tests/public_samples/__init__.py +17 -0
- google/genai/tests/public_samples/test_gemini_text_only.py +34 -0
- google/genai/tests/pytest_helper.py +229 -0
- google/genai/tests/shared/__init__.py +16 -0
- google/genai/tests/shared/batches/__init__.py +14 -0
- google/genai/tests/shared/batches/test_create_delete.py +57 -0
- google/genai/tests/shared/batches/test_create_get_cancel.py +56 -0
- google/genai/tests/shared/batches/test_list.py +40 -0
- google/genai/tests/shared/caches/__init__.py +14 -0
- google/genai/tests/shared/caches/test_create_get_delete.py +67 -0
- google/genai/tests/shared/caches/test_create_update_get.py +71 -0
- google/genai/tests/shared/caches/test_list.py +40 -0
- google/genai/tests/shared/chats/__init__.py +14 -0
- google/genai/tests/shared/chats/test_send_message.py +48 -0
- google/genai/tests/shared/chats/test_send_message_stream.py +50 -0
- google/genai/tests/shared/files/__init__.py +14 -0
- google/genai/tests/shared/files/test_list.py +41 -0
- google/genai/tests/shared/files/test_upload_get_delete.py +54 -0
- google/genai/tests/shared/models/__init__.py +14 -0
- google/genai/tests/shared/models/test_compute_tokens.py +41 -0
- google/genai/tests/shared/models/test_count_tokens.py +40 -0
- google/genai/tests/shared/models/test_edit_image.py +67 -0
- google/genai/tests/shared/models/test_embed.py +40 -0
- google/genai/tests/shared/models/test_generate_content.py +39 -0
- google/genai/tests/shared/models/test_generate_content_stream.py +54 -0
- google/genai/tests/shared/models/test_generate_images.py +40 -0
- google/genai/tests/shared/models/test_generate_videos.py +38 -0
- google/genai/tests/shared/models/test_list.py +37 -0
- google/genai/tests/shared/models/test_recontext_image.py +55 -0
- google/genai/tests/shared/models/test_segment_image.py +52 -0
- google/genai/tests/shared/models/test_upscale_image.py +52 -0
- google/genai/tests/shared/tunings/__init__.py +16 -0
- google/genai/tests/shared/tunings/test_create.py +46 -0
- google/genai/tests/shared/tunings/test_create_get_cancel.py +56 -0
- google/genai/tests/shared/tunings/test_list.py +39 -0
- google/genai/tests/tokens/__init__.py +16 -0
- google/genai/tests/tokens/test_create.py +154 -0
- google/genai/tests/transformers/__init__.py +17 -0
- google/genai/tests/transformers/test_blobs.py +71 -0
- google/genai/tests/transformers/test_bytes.py +15 -0
- google/genai/tests/transformers/test_duck_type.py +96 -0
- google/genai/tests/transformers/test_function_responses.py +72 -0
- google/genai/tests/transformers/test_schema.py +653 -0
- google/genai/tests/transformers/test_t_batch.py +286 -0
- google/genai/tests/transformers/test_t_content.py +160 -0
- google/genai/tests/transformers/test_t_contents.py +398 -0
- google/genai/tests/transformers/test_t_part.py +85 -0
- google/genai/tests/transformers/test_t_parts.py +87 -0
- google/genai/tests/transformers/test_t_tool.py +157 -0
- google/genai/tests/transformers/test_t_tools.py +195 -0
- google/genai/tests/tunings/__init__.py +16 -0
- google/genai/tests/tunings/test_cancel.py +39 -0
- google/genai/tests/tunings/test_end_to_end.py +106 -0
- google/genai/tests/tunings/test_get.py +67 -0
- google/genai/tests/tunings/test_list.py +75 -0
- google/genai/tests/tunings/test_tune.py +268 -0
- google/genai/tests/types/__init__.py +16 -0
- google/genai/tests/types/test_bytes_internal.py +271 -0
- google/genai/tests/types/test_bytes_type.py +152 -0
- google/genai/tests/types/test_future.py +101 -0
- google/genai/tests/types/test_optional_types.py +36 -0
- google/genai/tests/types/test_part_type.py +616 -0
- google/genai/tests/types/test_schema_from_json_schema.py +417 -0
- google/genai/tests/types/test_schema_json_schema.py +468 -0
- google/genai/tests/types/test_types.py +2903 -0
- google/genai/tunings.py +57 -57
- google/genai/types.py +229 -121
- google/genai/version.py +1 -1
- {google_genai-1.53.0.dist-info → google_genai-1.55.0.dist-info}/METADATA +4 -2
- google_genai-1.55.0.dist-info/RECORD +345 -0
- google_genai-1.53.0.dist-info/RECORD +0 -41
- {google_genai-1.53.0.dist-info → google_genai-1.55.0.dist-info}/WHEEL +0 -0
- {google_genai-1.53.0.dist-info → google_genai-1.55.0.dist-info}/licenses/LICENSE +0 -0
- {google_genai-1.53.0.dist-info → google_genai-1.55.0.dist-info}/top_level.txt +0 -0
|
@@ -0,0 +1,343 @@
|
|
|
1
|
+
# Copyright 2025 Google LLC
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
#
|
|
15
|
+
|
|
16
|
+
import unittest
|
|
17
|
+
from unittest.mock import MagicMock, patch
|
|
18
|
+
|
|
19
|
+
from sentencepiece import sentencepiece_model_pb2
|
|
20
|
+
|
|
21
|
+
from ... import local_tokenizer
|
|
22
|
+
from ... import types
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class TestLocalTokenizer(unittest.TestCase):
|
|
26
|
+
|
|
27
|
+
def setUp(self):
|
|
28
|
+
# This setup will be used by all tests
|
|
29
|
+
self.mock_load_model_proto = patch(
|
|
30
|
+
'genai._local_tokenizer_loader.load_model_proto'
|
|
31
|
+
).start()
|
|
32
|
+
self.mock_get_sentencepiece = patch(
|
|
33
|
+
'genai._local_tokenizer_loader.get_sentencepiece'
|
|
34
|
+
).start()
|
|
35
|
+
|
|
36
|
+
self.mock_load_model_proto.return_value = MagicMock()
|
|
37
|
+
self.mock_tokenizer = MagicMock()
|
|
38
|
+
self.mock_get_sentencepiece.return_value = self.mock_tokenizer
|
|
39
|
+
|
|
40
|
+
self.tokenizer = local_tokenizer.LocalTokenizer(model_name='gemini-2.5-flash')
|
|
41
|
+
|
|
42
|
+
def tearDown(self):
|
|
43
|
+
patch.stopall()
|
|
44
|
+
|
|
45
|
+
def test_count_tokens_simple_string(self):
|
|
46
|
+
self.mock_tokenizer.encode.return_value = [[1, 2, 3]]
|
|
47
|
+
result = self.tokenizer.count_tokens('Hello world')
|
|
48
|
+
self.assertEqual(result.total_tokens, 3)
|
|
49
|
+
self.mock_tokenizer.encode.assert_called_once_with(['Hello world'])
|
|
50
|
+
|
|
51
|
+
def test_count_tokens_list_of_strings(self):
|
|
52
|
+
self.mock_tokenizer.encode.return_value = [[1, 2], [3]]
|
|
53
|
+
result = self.tokenizer.count_tokens(['Hello', 'world'])
|
|
54
|
+
self.assertEqual(result.total_tokens, 3)
|
|
55
|
+
self.mock_tokenizer.encode.assert_called_once_with(['Hello', 'world'])
|
|
56
|
+
|
|
57
|
+
def test_count_tokens_with_content_object(self):
|
|
58
|
+
self.mock_tokenizer.encode.return_value = [[1, 2, 3]]
|
|
59
|
+
content = types.Content(parts=[types.Part(text='Hello world')])
|
|
60
|
+
result = self.tokenizer.count_tokens(content)
|
|
61
|
+
self.assertEqual(result.total_tokens, 3)
|
|
62
|
+
self.mock_tokenizer.encode.assert_called_once_with(['Hello world'])
|
|
63
|
+
|
|
64
|
+
def test_count_tokens_with_chat_history(self):
|
|
65
|
+
self.mock_tokenizer.encode.return_value = [[1, 2], [3, 4, 5]]
|
|
66
|
+
history = [
|
|
67
|
+
types.Content(role='user', parts=[types.Part(text='Hello')]),
|
|
68
|
+
types.Content(role='model', parts=[types.Part(text='Hi there!')]),
|
|
69
|
+
]
|
|
70
|
+
result = self.tokenizer.count_tokens(history)
|
|
71
|
+
self.assertEqual(result.total_tokens, 5)
|
|
72
|
+
self.mock_tokenizer.encode.assert_called_once_with(['Hello', 'Hi there!'])
|
|
73
|
+
|
|
74
|
+
def test_count_tokens_with_tools(self):
|
|
75
|
+
self.mock_tokenizer.encode.return_value = [
|
|
76
|
+
[1],
|
|
77
|
+
[1, 2],
|
|
78
|
+
[1, 2, 3],
|
|
79
|
+
[1, 2, 3, 4],
|
|
80
|
+
[1, 2, 3, 4, 5],
|
|
81
|
+
[1, 2, 3, 4, 5, 6],
|
|
82
|
+
]
|
|
83
|
+
tool = types.Tool(
|
|
84
|
+
function_declarations=[
|
|
85
|
+
types.FunctionDeclaration(
|
|
86
|
+
name='get_weather',
|
|
87
|
+
description='Get the weather for a location',
|
|
88
|
+
parameters=types.Schema(
|
|
89
|
+
type=types.Type.OBJECT,
|
|
90
|
+
properties={
|
|
91
|
+
'location': types.Schema(
|
|
92
|
+
type=types.Type.STRING, description='The location'
|
|
93
|
+
)
|
|
94
|
+
},
|
|
95
|
+
required=['location'],
|
|
96
|
+
),
|
|
97
|
+
)
|
|
98
|
+
]
|
|
99
|
+
)
|
|
100
|
+
config = types.CountTokensConfig(tools=[tool])
|
|
101
|
+
result = self.tokenizer.count_tokens(
|
|
102
|
+
'What is the weather in Boston?', config=config
|
|
103
|
+
)
|
|
104
|
+
self.assertEqual(result.total_tokens, 21)
|
|
105
|
+
self.mock_tokenizer.encode.assert_called_once_with([
|
|
106
|
+
'What is the weather in Boston?',
|
|
107
|
+
'get_weather',
|
|
108
|
+
'Get the weather for a location',
|
|
109
|
+
'location',
|
|
110
|
+
'location',
|
|
111
|
+
'The location',
|
|
112
|
+
])
|
|
113
|
+
|
|
114
|
+
def test_count_tokens_with_function_call(self):
|
|
115
|
+
self.mock_tokenizer.encode.return_value = [[1, 2], [3], [4, 5]]
|
|
116
|
+
content = types.Content(
|
|
117
|
+
role='model',
|
|
118
|
+
parts=[
|
|
119
|
+
types.Part(
|
|
120
|
+
function_call=types.FunctionCall(
|
|
121
|
+
name='get_weather', args={'location': 'Boston'}
|
|
122
|
+
)
|
|
123
|
+
)
|
|
124
|
+
],
|
|
125
|
+
)
|
|
126
|
+
result = self.tokenizer.count_tokens(content)
|
|
127
|
+
self.assertEqual(result.total_tokens, 5)
|
|
128
|
+
self.mock_tokenizer.encode.assert_called_once_with(
|
|
129
|
+
['get_weather', 'location', 'Boston']
|
|
130
|
+
)
|
|
131
|
+
|
|
132
|
+
def test_count_tokens_with_function_response(self):
|
|
133
|
+
self.mock_tokenizer.encode.return_value = [[1, 2], [3], [4, 5]]
|
|
134
|
+
content = types.Content(
|
|
135
|
+
role='user',
|
|
136
|
+
parts=[
|
|
137
|
+
types.Part(
|
|
138
|
+
function_response=types.FunctionResponse(
|
|
139
|
+
name='get_weather', response={'weather': 'sunny'}
|
|
140
|
+
)
|
|
141
|
+
)
|
|
142
|
+
],
|
|
143
|
+
)
|
|
144
|
+
result = self.tokenizer.count_tokens(content)
|
|
145
|
+
self.assertEqual(result.total_tokens, 5)
|
|
146
|
+
self.mock_tokenizer.encode.assert_called_once_with(
|
|
147
|
+
['get_weather', 'weather', 'sunny']
|
|
148
|
+
)
|
|
149
|
+
|
|
150
|
+
def test_count_tokens_with_unsupported_content(self):
|
|
151
|
+
with self.assertRaises(ValueError):
|
|
152
|
+
self.tokenizer.count_tokens(
|
|
153
|
+
[
|
|
154
|
+
types.Content(
|
|
155
|
+
parts=[
|
|
156
|
+
types.Part(
|
|
157
|
+
inline_data=types.Blob(
|
|
158
|
+
data=b'test', mime_type='image/png'
|
|
159
|
+
)
|
|
160
|
+
)
|
|
161
|
+
]
|
|
162
|
+
)
|
|
163
|
+
]
|
|
164
|
+
)
|
|
165
|
+
|
|
166
|
+
def test_count_tokens_with_system_instruction(self):
|
|
167
|
+
self.mock_tokenizer.encode.return_value = [[1, 2, 3], [4, 5]]
|
|
168
|
+
config = types.CountTokensConfig(
|
|
169
|
+
system_instruction=types.Content(
|
|
170
|
+
parts=[types.Part(text='You are a helpful assistant.')]
|
|
171
|
+
)
|
|
172
|
+
)
|
|
173
|
+
result = self.tokenizer.count_tokens('Hello', config=config)
|
|
174
|
+
self.assertEqual(result.total_tokens, 5)
|
|
175
|
+
self.mock_tokenizer.encode.assert_called_once_with(
|
|
176
|
+
['Hello', 'You are a helpful assistant.']
|
|
177
|
+
)
|
|
178
|
+
|
|
179
|
+
def test_count_tokens_with_response_schema(self):
|
|
180
|
+
self.mock_tokenizer.encode.return_value = [
|
|
181
|
+
[1],
|
|
182
|
+
[1, 2],
|
|
183
|
+
[1, 2, 3],
|
|
184
|
+
[1, 2, 3, 4],
|
|
185
|
+
[1, 2, 3, 4, 5],
|
|
186
|
+
]
|
|
187
|
+
schema = types.Schema(
|
|
188
|
+
type=types.Type.OBJECT,
|
|
189
|
+
format='schema_format',
|
|
190
|
+
description='Recipe schema',
|
|
191
|
+
enum=['schema_enum1', 'schema_enum2'],
|
|
192
|
+
properties={
|
|
193
|
+
'recipe_name': types.Schema(
|
|
194
|
+
type=types.Type.STRING,
|
|
195
|
+
description='Name of the recipe',
|
|
196
|
+
)
|
|
197
|
+
},
|
|
198
|
+
items=types.Schema(
|
|
199
|
+
type=types.Type.STRING,
|
|
200
|
+
description='Item in the recipe',
|
|
201
|
+
),
|
|
202
|
+
example={
|
|
203
|
+
'recipe_example': types.Schema(
|
|
204
|
+
type=types.Type.STRING,
|
|
205
|
+
description='example in the recipe',
|
|
206
|
+
)
|
|
207
|
+
},
|
|
208
|
+
required=['recipe_name'],
|
|
209
|
+
)
|
|
210
|
+
config = types.CountTokensConfig(
|
|
211
|
+
generation_config=types.GenerationConfig(response_schema=schema)
|
|
212
|
+
)
|
|
213
|
+
result = self.tokenizer.count_tokens(
|
|
214
|
+
'Generate a recipe for chocolate chip cookies.', config=config
|
|
215
|
+
)
|
|
216
|
+
self.assertEqual(result.total_tokens, 15)
|
|
217
|
+
self.mock_tokenizer.encode.assert_called_once_with([
|
|
218
|
+
'Generate a recipe for chocolate chip cookies.',
|
|
219
|
+
'schema_format',
|
|
220
|
+
'Recipe schema',
|
|
221
|
+
'schema_enum1',
|
|
222
|
+
'schema_enum2',
|
|
223
|
+
'recipe_name',
|
|
224
|
+
'Item in the recipe',
|
|
225
|
+
'recipe_name',
|
|
226
|
+
'Name of the recipe',
|
|
227
|
+
'recipe_example',
|
|
228
|
+
])
|
|
229
|
+
|
|
230
|
+
def test_count_tokens_with_unsupported_fields_logs_warning(self):
|
|
231
|
+
self.mock_tokenizer.encode.return_value = [[1, 2, 3]]
|
|
232
|
+
content_with_unsupported = types.Content(
|
|
233
|
+
role='user',
|
|
234
|
+
parts=[
|
|
235
|
+
types.Part(text='hello'),
|
|
236
|
+
# executable_code is not supported by _TextsAccumulator
|
|
237
|
+
types.Part(
|
|
238
|
+
executable_code=types.ExecutableCode(
|
|
239
|
+
language='PYTHON', code='print(1)'
|
|
240
|
+
)
|
|
241
|
+
),
|
|
242
|
+
],
|
|
243
|
+
)
|
|
244
|
+
with self.assertLogs('google_genai.local_tokenizer', level='WARNING') as cm:
|
|
245
|
+
self.tokenizer.count_tokens(content_with_unsupported)
|
|
246
|
+
self.assertIn(
|
|
247
|
+
'Content contains unsupported types for token counting', cm.output[0]
|
|
248
|
+
)
|
|
249
|
+
|
|
250
|
+
def test_compute_tokens_simple_string(self):
|
|
251
|
+
mock_spt = MagicMock()
|
|
252
|
+
mock_spt.pieces = [
|
|
253
|
+
MagicMock(id=1, piece='He'),
|
|
254
|
+
MagicMock(id=2, piece='llo'),
|
|
255
|
+
MagicMock(id=3, piece=' world'),
|
|
256
|
+
]
|
|
257
|
+
self.mock_tokenizer.EncodeAsImmutableProto.return_value = [mock_spt]
|
|
258
|
+
result = self.tokenizer.compute_tokens('Hello world')
|
|
259
|
+
self.assertEqual(len(result.tokens_info), 1)
|
|
260
|
+
self.assertEqual(result.tokens_info[0].token_ids, [1, 2, 3])
|
|
261
|
+
self.assertEqual(result.tokens_info[0].tokens, [b'He', b'llo', b' world'])
|
|
262
|
+
self.assertEqual(result.tokens_info[0].role, 'user')
|
|
263
|
+
self.mock_tokenizer.EncodeAsImmutableProto.assert_called_once_with(
|
|
264
|
+
['Hello world']
|
|
265
|
+
)
|
|
266
|
+
|
|
267
|
+
def test_compute_tokens_with_chat_history(self):
|
|
268
|
+
mock_spt1 = MagicMock()
|
|
269
|
+
mock_spt1.pieces = [MagicMock(id=1, piece='Hello')]
|
|
270
|
+
mock_spt2 = MagicMock()
|
|
271
|
+
mock_spt2.pieces = [
|
|
272
|
+
MagicMock(id=2, piece='Hi'),
|
|
273
|
+
MagicMock(id=3, piece=' there!'),
|
|
274
|
+
]
|
|
275
|
+
self.mock_tokenizer.EncodeAsImmutableProto.return_value = [
|
|
276
|
+
mock_spt1,
|
|
277
|
+
mock_spt2,
|
|
278
|
+
]
|
|
279
|
+
history = [
|
|
280
|
+
types.Content(role='user', parts=[types.Part(text='Hello')]),
|
|
281
|
+
types.Content(role='model', parts=[types.Part(text='Hi there!')]),
|
|
282
|
+
]
|
|
283
|
+
result = self.tokenizer.compute_tokens(history)
|
|
284
|
+
self.assertEqual(len(result.tokens_info), 2)
|
|
285
|
+
self.assertEqual(result.tokens_info[0].token_ids, [1])
|
|
286
|
+
self.assertEqual(result.tokens_info[0].tokens, [b'Hello'])
|
|
287
|
+
self.assertEqual(result.tokens_info[0].role, 'user')
|
|
288
|
+
self.assertEqual(result.tokens_info[1].token_ids, [2, 3])
|
|
289
|
+
self.assertEqual(result.tokens_info[1].tokens, [b'Hi', b' there!'])
|
|
290
|
+
self.assertEqual(result.tokens_info[1].role, 'model')
|
|
291
|
+
self.mock_tokenizer.EncodeAsImmutableProto.assert_called_once_with(
|
|
292
|
+
['Hello', 'Hi there!']
|
|
293
|
+
)
|
|
294
|
+
|
|
295
|
+
def test_compute_tokens_with_byte_tokens(self):
|
|
296
|
+
mock_spt = MagicMock()
|
|
297
|
+
mock_spt.pieces = [
|
|
298
|
+
MagicMock(id=1, piece='<0x48>'),
|
|
299
|
+
MagicMock(id=2, piece='ello'),
|
|
300
|
+
]
|
|
301
|
+
self.mock_tokenizer.EncodeAsImmutableProto.return_value = [mock_spt]
|
|
302
|
+
self.tokenizer._model_proto = sentencepiece_model_pb2.ModelProto(
|
|
303
|
+
pieces=[
|
|
304
|
+
sentencepiece_model_pb2.ModelProto.SentencePiece(),
|
|
305
|
+
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
|
306
|
+
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.BYTE
|
|
307
|
+
),
|
|
308
|
+
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
|
309
|
+
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.NORMAL
|
|
310
|
+
),
|
|
311
|
+
]
|
|
312
|
+
)
|
|
313
|
+
result = self.tokenizer.compute_tokens('Hello')
|
|
314
|
+
self.assertEqual(len(result.tokens_info), 1)
|
|
315
|
+
self.assertEqual(result.tokens_info[0].token_ids, [1, 2])
|
|
316
|
+
self.assertEqual(result.tokens_info[0].tokens, [b'H', b'ello'])
|
|
317
|
+
self.mock_tokenizer.EncodeAsImmutableProto.assert_called_once_with(
|
|
318
|
+
['Hello']
|
|
319
|
+
)
|
|
320
|
+
|
|
321
|
+
|
|
322
|
+
class TestParseHexByte(unittest.TestCase):
|
|
323
|
+
|
|
324
|
+
def test_valid_hex(self):
|
|
325
|
+
self.assertEqual(local_tokenizer._parse_hex_byte('<0x41>'), 65)
|
|
326
|
+
self.assertEqual(local_tokenizer._parse_hex_byte('<0xFF>'), 255)
|
|
327
|
+
self.assertEqual(local_tokenizer._parse_hex_byte('<0x00>'), 0)
|
|
328
|
+
|
|
329
|
+
def test_invalid_length(self):
|
|
330
|
+
with self.assertRaisesRegex(ValueError, 'Invalid byte length'):
|
|
331
|
+
local_tokenizer._parse_hex_byte('<0x41')
|
|
332
|
+
with self.assertRaisesRegex(ValueError, 'Invalid byte length'):
|
|
333
|
+
local_tokenizer._parse_hex_byte('<0x411>')
|
|
334
|
+
|
|
335
|
+
def test_invalid_format(self):
|
|
336
|
+
with self.assertRaisesRegex(ValueError, 'Invalid byte format'):
|
|
337
|
+
local_tokenizer._parse_hex_byte(' 0x41>')
|
|
338
|
+
with self.assertRaisesRegex(ValueError, 'Invalid byte format'):
|
|
339
|
+
local_tokenizer._parse_hex_byte('<0x41 ')
|
|
340
|
+
|
|
341
|
+
def test_invalid_hex_value(self):
|
|
342
|
+
with self.assertRaisesRegex(ValueError, 'Invalid hex value'):
|
|
343
|
+
local_tokenizer._parse_hex_byte('<0xFG>')
|
|
@@ -0,0 +1,235 @@
|
|
|
1
|
+
# Copyright 2025 Google LLC
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
#
|
|
15
|
+
|
|
16
|
+
import unittest
|
|
17
|
+
from unittest.mock import MagicMock, mock_open, patch
|
|
18
|
+
|
|
19
|
+
import sentencepiece as spm
|
|
20
|
+
from sentencepiece import sentencepiece_model_pb2
|
|
21
|
+
|
|
22
|
+
from ... import _local_tokenizer_loader as loader
|
|
23
|
+
|
|
24
|
+
# A minimal valid sentencepiece model proto
|
|
25
|
+
FAKE_MODEL_CONTENT = sentencepiece_model_pb2.ModelProto(
|
|
26
|
+
pieces=[
|
|
27
|
+
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
|
28
|
+
piece="<unk>",
|
|
29
|
+
score=0,
|
|
30
|
+
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.UNKNOWN,
|
|
31
|
+
),
|
|
32
|
+
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
|
33
|
+
piece="<s>",
|
|
34
|
+
score=0,
|
|
35
|
+
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.CONTROL,
|
|
36
|
+
),
|
|
37
|
+
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
|
38
|
+
piece="</s>",
|
|
39
|
+
score=0,
|
|
40
|
+
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.CONTROL,
|
|
41
|
+
),
|
|
42
|
+
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
|
43
|
+
piece="a",
|
|
44
|
+
score=0,
|
|
45
|
+
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.NORMAL,
|
|
46
|
+
),
|
|
47
|
+
]
|
|
48
|
+
).SerializeToString()
|
|
49
|
+
|
|
50
|
+
GEMMA2_HASH = "61a7b147390c64585d6c3543dd6fc636906c9af3865a5548f27f31aee1d4c8e2"
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
class TestGetTokenizerName(unittest.TestCase):
|
|
54
|
+
|
|
55
|
+
def test_get_tokenizer_name_success(self):
|
|
56
|
+
self.assertEqual(loader.get_tokenizer_name("gemini-2.5-pro"), "gemma3")
|
|
57
|
+
self.assertEqual(
|
|
58
|
+
loader.get_tokenizer_name("gemini-2.5-pro-preview-06-05"), "gemma3"
|
|
59
|
+
)
|
|
60
|
+
|
|
61
|
+
def test_get_tokenizer_name_unsupported(self):
|
|
62
|
+
with self.assertRaisesRegex(
|
|
63
|
+
ValueError, "Model unsupported-model is not supported"
|
|
64
|
+
):
|
|
65
|
+
loader.get_tokenizer_name("unsupported-model")
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
@patch("genai._local_tokenizer_loader.os.rename")
|
|
69
|
+
@patch("genai._local_tokenizer_loader.os.makedirs")
|
|
70
|
+
@patch("genai._local_tokenizer_loader.os.remove")
|
|
71
|
+
@patch("genai._local_tokenizer_loader.open", new_callable=mock_open)
|
|
72
|
+
@patch("genai._local_tokenizer_loader.os.path.exists")
|
|
73
|
+
@patch("genai._local_tokenizer_loader.requests.get")
|
|
74
|
+
@patch("genai._local_tokenizer_loader.hashlib.sha256")
|
|
75
|
+
class TestLoaderFunctions(unittest.TestCase):
|
|
76
|
+
|
|
77
|
+
def setUp(self):
|
|
78
|
+
# Clear caches before each test
|
|
79
|
+
loader.load_model_proto.cache_clear()
|
|
80
|
+
loader.get_sentencepiece.cache_clear()
|
|
81
|
+
# Patch tempfile.gettempdir to control cache location
|
|
82
|
+
self.tempdir_patcher = patch(
|
|
83
|
+
"tempfile.gettempdir", return_value="/tmp/fake_temp_dir"
|
|
84
|
+
)
|
|
85
|
+
self.mock_tempdir = self.tempdir_patcher.start()
|
|
86
|
+
|
|
87
|
+
def tearDown(self):
|
|
88
|
+
self.tempdir_patcher.stop()
|
|
89
|
+
|
|
90
|
+
def _setup_get_mock(self, mock_get):
|
|
91
|
+
mock_response = MagicMock()
|
|
92
|
+
mock_response.content = FAKE_MODEL_CONTENT
|
|
93
|
+
mock_response.raise_for_status.return_value = None
|
|
94
|
+
mock_get.return_value = mock_response
|
|
95
|
+
|
|
96
|
+
def test_load_model_proto_from_url(
|
|
97
|
+
self,
|
|
98
|
+
mock_sha256,
|
|
99
|
+
mock_get,
|
|
100
|
+
mock_exists,
|
|
101
|
+
mock_open_func,
|
|
102
|
+
mock_remove,
|
|
103
|
+
mock_makedirs,
|
|
104
|
+
mock_rename,
|
|
105
|
+
):
|
|
106
|
+
mock_exists.return_value = False # Don't use cache
|
|
107
|
+
self._setup_get_mock(mock_get)
|
|
108
|
+
mock_sha256.return_value.hexdigest.return_value = GEMMA2_HASH
|
|
109
|
+
|
|
110
|
+
proto = loader.load_model_proto("gemma2")
|
|
111
|
+
|
|
112
|
+
self.assertIsInstance(proto, sentencepiece_model_pb2.ModelProto)
|
|
113
|
+
self.assertEqual(len(proto.pieces), 4)
|
|
114
|
+
mock_get.assert_called_once()
|
|
115
|
+
mock_makedirs.assert_called_once()
|
|
116
|
+
mock_open_func.assert_called()
|
|
117
|
+
mock_rename.assert_called_once()
|
|
118
|
+
|
|
119
|
+
def test_load_model_proto_from_cache(
|
|
120
|
+
self,
|
|
121
|
+
mock_sha256,
|
|
122
|
+
mock_get,
|
|
123
|
+
mock_exists,
|
|
124
|
+
mock_open_func,
|
|
125
|
+
mock_remove,
|
|
126
|
+
mock_makedirs,
|
|
127
|
+
mock_rename,
|
|
128
|
+
):
|
|
129
|
+
mock_exists.return_value = True # Use cache
|
|
130
|
+
mock_open_func.return_value.read.return_value = FAKE_MODEL_CONTENT
|
|
131
|
+
mock_sha256.return_value.hexdigest.return_value = GEMMA2_HASH
|
|
132
|
+
|
|
133
|
+
proto = loader.load_model_proto("gemma2")
|
|
134
|
+
|
|
135
|
+
self.assertIsInstance(proto, sentencepiece_model_pb2.ModelProto)
|
|
136
|
+
mock_get.assert_not_called()
|
|
137
|
+
|
|
138
|
+
def test_load_model_proto_corrupted_cache(
|
|
139
|
+
self,
|
|
140
|
+
mock_sha256,
|
|
141
|
+
mock_get,
|
|
142
|
+
mock_exists,
|
|
143
|
+
mock_open_func,
|
|
144
|
+
mock_remove,
|
|
145
|
+
mock_makedirs,
|
|
146
|
+
mock_rename,
|
|
147
|
+
):
|
|
148
|
+
mock_exists.return_value = True # Use cache initially
|
|
149
|
+
self._setup_get_mock(mock_get)
|
|
150
|
+
mock_open_func.return_value.__enter__.return_value.read.return_value = (
|
|
151
|
+
b"corrupted"
|
|
152
|
+
)
|
|
153
|
+
|
|
154
|
+
# First hash for corrupted cache, second for good download
|
|
155
|
+
mock_sha256.side_effect = [
|
|
156
|
+
MagicMock(hexdigest=MagicMock(return_value="wrong_hash")),
|
|
157
|
+
MagicMock(hexdigest=MagicMock(return_value=GEMMA2_HASH)),
|
|
158
|
+
]
|
|
159
|
+
|
|
160
|
+
proto = loader.load_model_proto("gemma2")
|
|
161
|
+
|
|
162
|
+
self.assertIsInstance(proto, sentencepiece_model_pb2.ModelProto)
|
|
163
|
+
mock_remove.assert_called_once()
|
|
164
|
+
mock_get.assert_called_once()
|
|
165
|
+
|
|
166
|
+
def test_load_model_proto_bad_hash_from_url(
|
|
167
|
+
self,
|
|
168
|
+
mock_sha256,
|
|
169
|
+
mock_get,
|
|
170
|
+
mock_exists,
|
|
171
|
+
mock_open_func,
|
|
172
|
+
mock_remove,
|
|
173
|
+
mock_makedirs,
|
|
174
|
+
mock_rename,
|
|
175
|
+
):
|
|
176
|
+
mock_exists.return_value = False
|
|
177
|
+
self._setup_get_mock(mock_get)
|
|
178
|
+
mock_sha256.return_value.hexdigest.return_value = "wrong_hash"
|
|
179
|
+
|
|
180
|
+
with self.assertRaisesRegex(
|
|
181
|
+
ValueError, "Downloaded model file is corrupted"
|
|
182
|
+
):
|
|
183
|
+
loader.load_model_proto("gemma2")
|
|
184
|
+
|
|
185
|
+
def test_load_model_proto_unsupported(self, *args):
|
|
186
|
+
with self.assertRaisesRegex(
|
|
187
|
+
ValueError, "Tokenizer unsupported is not supported"
|
|
188
|
+
):
|
|
189
|
+
loader.load_model_proto("unsupported")
|
|
190
|
+
|
|
191
|
+
def test_get_sentencepiece_success(
|
|
192
|
+
self,
|
|
193
|
+
mock_sha256,
|
|
194
|
+
mock_get,
|
|
195
|
+
mock_exists,
|
|
196
|
+
mock_open_func,
|
|
197
|
+
mock_remove,
|
|
198
|
+
mock_makedirs,
|
|
199
|
+
mock_rename,
|
|
200
|
+
):
|
|
201
|
+
mock_exists.return_value = False
|
|
202
|
+
self._setup_get_mock(mock_get)
|
|
203
|
+
mock_sha256.return_value.hexdigest.return_value = GEMMA2_HASH
|
|
204
|
+
|
|
205
|
+
processor = loader.get_sentencepiece("gemma2")
|
|
206
|
+
|
|
207
|
+
self.assertIsInstance(processor, spm.SentencePieceProcessor)
|
|
208
|
+
mock_get.assert_called_once()
|
|
209
|
+
|
|
210
|
+
def test_get_sentencepiece_unsupported(self, *args):
|
|
211
|
+
with self.assertRaisesRegex(
|
|
212
|
+
ValueError, "Tokenizer unsupported is not supported"
|
|
213
|
+
):
|
|
214
|
+
loader.get_sentencepiece("unsupported")
|
|
215
|
+
|
|
216
|
+
def test_get_sentencepiece_caching(
|
|
217
|
+
self,
|
|
218
|
+
mock_sha256,
|
|
219
|
+
mock_get,
|
|
220
|
+
mock_exists,
|
|
221
|
+
mock_open_func,
|
|
222
|
+
mock_remove,
|
|
223
|
+
mock_makedirs,
|
|
224
|
+
mock_rename,
|
|
225
|
+
):
|
|
226
|
+
mock_exists.return_value = False
|
|
227
|
+
self._setup_get_mock(mock_get)
|
|
228
|
+
mock_sha256.return_value.hexdigest.return_value = GEMMA2_HASH
|
|
229
|
+
|
|
230
|
+
# Call twice
|
|
231
|
+
loader.get_sentencepiece("gemma2")
|
|
232
|
+
loader.get_sentencepiece("gemma2")
|
|
233
|
+
|
|
234
|
+
# Should only be loaded once due to lru_cache
|
|
235
|
+
mock_get.assert_called_once()
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
# Copyright 2025 Google LLC
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
#
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
"""Tests for the Google GenAI SDK's _mcp_utils module."""
|
|
@@ -0,0 +1,89 @@
|
|
|
1
|
+
# Copyright 2025 Google LLC
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
#
|
|
15
|
+
|
|
16
|
+
import typing
|
|
17
|
+
from typing import Any
|
|
18
|
+
|
|
19
|
+
from ... import _mcp_utils
|
|
20
|
+
from ... import types
|
|
21
|
+
|
|
22
|
+
if typing.TYPE_CHECKING:
|
|
23
|
+
from mcp.types import Tool as McpTool
|
|
24
|
+
from mcp import ClientSession as McpClientSession
|
|
25
|
+
else:
|
|
26
|
+
McpTool: typing.Type = Any
|
|
27
|
+
McpClientSession: typing.Type = Any
|
|
28
|
+
try:
|
|
29
|
+
from mcp.types import Tool as McpTool
|
|
30
|
+
from mcp import ClientSession as McpClientSession
|
|
31
|
+
except ImportError:
|
|
32
|
+
McpTool = None
|
|
33
|
+
McpClientSession = None
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def test_mcp_tools():
|
|
37
|
+
"""Test whether the list of tools contains any MCP tools."""
|
|
38
|
+
if McpTool is None:
|
|
39
|
+
return
|
|
40
|
+
mcp_tools = [
|
|
41
|
+
McpTool(
|
|
42
|
+
name='tool',
|
|
43
|
+
description='tool-description',
|
|
44
|
+
inputSchema={
|
|
45
|
+
'type': 'OBJECT',
|
|
46
|
+
'properties': {
|
|
47
|
+
'key1': {'type': 'STRING'},
|
|
48
|
+
'key2': {'type': 'NUMBER'},
|
|
49
|
+
},
|
|
50
|
+
},
|
|
51
|
+
),
|
|
52
|
+
]
|
|
53
|
+
assert _mcp_utils.has_mcp_tool_usage(mcp_tools)
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def test_mcp_client_session():
|
|
57
|
+
"""Test whether the list of tools contains any MCP tools."""
|
|
58
|
+
|
|
59
|
+
class MockMcpClientSession(McpClientSession):
|
|
60
|
+
|
|
61
|
+
def __init__(self):
|
|
62
|
+
self._read_stream = None
|
|
63
|
+
self._write_stream = None
|
|
64
|
+
|
|
65
|
+
mcp_tools = [
|
|
66
|
+
MockMcpClientSession(),
|
|
67
|
+
]
|
|
68
|
+
assert _mcp_utils.has_mcp_tool_usage(mcp_tools)
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
def test_no_mcp_tools():
|
|
72
|
+
if McpClientSession is None:
|
|
73
|
+
return
|
|
74
|
+
"""Test whether the list of tools contains any MCP tools."""
|
|
75
|
+
gemini_tools = [
|
|
76
|
+
types.Tool(
|
|
77
|
+
function_declarations=[
|
|
78
|
+
types.FunctionDeclaration(
|
|
79
|
+
name='tool',
|
|
80
|
+
description='tool-description',
|
|
81
|
+
parameters=types.Schema(
|
|
82
|
+
type='OBJECT',
|
|
83
|
+
properties={},
|
|
84
|
+
),
|
|
85
|
+
),
|
|
86
|
+
],
|
|
87
|
+
),
|
|
88
|
+
]
|
|
89
|
+
assert not _mcp_utils.has_mcp_tool_usage(gemini_tools)
|