google-genai 1.7.0__py3-none-any.whl → 1.53.0__py3-none-any.whl

google/genai/_local_tokenizer_loader.py

@@ -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.
+#
+
+import dataclasses
+import functools
+import hashlib
+import os
+import tempfile
+from typing import Optional, cast
+import uuid
+
+import requests  # type: ignore
+import sentencepiece as spm
+from sentencepiece import sentencepiece_model_pb2
+
+
+# Source of truth: https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models
+_GEMINI_MODELS_TO_TOKENIZER_NAMES = {
+    "gemini-2.5-pro": "gemma3",
+    "gemini-2.5-flash": "gemma3",
+    "gemini-2.5-flash-lite": "gemma3",
+    "gemini-2.0-flash": "gemma3",
+    "gemini-2.0-flash-lite": "gemma3",
+}
+_GEMINI_STABLE_MODELS_TO_TOKENIZER_NAMES = {
+    "gemini-2.5-pro-preview-06-05": "gemma3",
+    "gemini-2.5-pro-preview-05-06": "gemma3",
+    "gemini-2.5-pro-exp-03-25": "gemma3",
+    "gemini-live-2.5-flash": "gemma3",
+    "gemini-2.5-flash-preview-05-20": "gemma3",
+    "gemini-2.5-flash-preview-04-17": "gemma3",
+    "gemini-2.5-flash-lite-preview-06-17": "gemma3",
+    "gemini-2.0-flash-001": "gemma3",
+    "gemini-2.0-flash-lite-001": "gemma3",
+}
+
+
+@dataclasses.dataclass(frozen=True)
+class _TokenizerConfig:
+  model_url: str
+  model_hash: str
+
+
+# TODO: update gemma3 tokenizer
+_TOKENIZERS = {
+    "gemma2": _TokenizerConfig(
+        model_url="https://raw.githubusercontent.com/google/gemma_pytorch/33b652c465537c6158f9a472ea5700e5e770ad3f/tokenizer/tokenizer.model",
+        model_hash=(
+            "61a7b147390c64585d6c3543dd6fc636906c9af3865a5548f27f31aee1d4c8e2"
+        ),
+    ),
+    "gemma3": _TokenizerConfig(
+        model_url="https://raw.githubusercontent.com/google/gemma_pytorch/014acb7ac4563a5f77c76d7ff98f31b568c16508/tokenizer/gemma3_cleaned_262144_v2.spiece.model",
+        model_hash=(
+            "1299c11d7cf632ef3b4e11937501358ada021bbdf7c47638d13c0ee982f2e79c"
+        ),
+    ),
+}
+
+
+def _load_file(file_url_path: str) -> bytes:
+  """Loads file bytes from the given file url path."""
+  resp = requests.get(file_url_path)
+  resp.raise_for_status()
+  return cast(bytes, resp.content)
+
+
+def _is_valid_model(*, model_data: bytes, expected_hash: str) -> bool:
+  """Returns true if the content is valid by checking the hash."""
+  if not expected_hash:
+    raise ValueError("expected_hash is required")
+  return hashlib.sha256(model_data).hexdigest() == expected_hash
+
+
+def _maybe_remove_file(file_path: str) -> None:
+  """Removes the file if exists."""
+  if not os.path.exists(file_path):
+    return
+  try:
+    os.remove(file_path)
+  except OSError:
+    # Don't raise if we cannot remove file.
+    pass
+
+
+def _maybe_load_from_cache(
+    *, file_path: str, expected_hash: str
+) -> Optional[bytes]:
+  """Loads the content from the cache path."""
+  if not os.path.exists(file_path):
+    return None
+  with open(file_path, "rb") as f:
+    content = f.read()
+  if _is_valid_model(model_data=content, expected_hash=expected_hash):
+    return content
+
+  # Cached file corrupted.
+  _maybe_remove_file(file_path)
+  return None
+
+
+def _maybe_save_to_cache(
+    *, cache_dir: str, cache_path: str, content: bytes
+) -> None:
+  """Saves the content to the cache path."""
+  try:
+    os.makedirs(cache_dir, exist_ok=True)
+    tmp_path = cache_dir + "." + str(uuid.uuid4()) + ".tmp"
+    with open(tmp_path, "wb") as f:
+      f.write(content)
+    os.rename(tmp_path, cache_path)
+  except OSError:
+    # Don't raise if we cannot write file.
+    pass
+
+
+def _load_from_url(*, file_url: str, expected_hash: str) -> bytes:
+  """Loads model bytes from the given file url."""
+  content = _load_file(file_url)
+  if not _is_valid_model(model_data=content, expected_hash=expected_hash):
+    actual_hash = hashlib.sha256(content).hexdigest()
+    raise ValueError(
+        "Downloaded model file is corrupted."
+        f" Expected hash {expected_hash}. Got file hash {actual_hash}."
+    )
+  return content
+
+
+def _load(*, file_url: str, expected_hash: str) -> bytes:
+  """Loads model bytes from the given file url.
+
+  1. If the find local cached file for the given url and the cached file hash
+     matches the expected hash, the cached file is returned.
+  2. If local cached file is not found or the hash does not match, the file is
+     downloaded from the given url. And write to local cache and return the
+     file bytes.
+  3. If the file downloaded from the given url does not match the expected
+     hash, raise ValueError.
+
+  Args:
+      file_url: The url of the file to load.
+      expected_hash: The expected hash of the file.
+
+  Returns:
+      The file bytes.
+  """
+  model_dir = os.path.join(tempfile.gettempdir(), "vertexai_tokenizer_model")
+  filename = hashlib.sha1(file_url.encode()).hexdigest()
+  model_path = os.path.join(model_dir, filename)
+
+  model_data = _maybe_load_from_cache(
+      file_path=model_path, expected_hash=expected_hash
+  )
+  if not model_data:
+    model_data = _load_from_url(file_url=file_url, expected_hash=expected_hash)
+
+  _maybe_save_to_cache(
+      cache_dir=model_dir, cache_path=model_path, content=model_data
+  )
+  return model_data
+
+
+def _load_model_proto_bytes(tokenizer_name: str) -> bytes:
+  """Loads model proto bytes from the given tokenizer name."""
+  if tokenizer_name not in _TOKENIZERS:
+    raise ValueError(
+        f"Tokenizer {tokenizer_name} is not supported."
+        f"Supported tokenizers: {list(_TOKENIZERS.keys())}"
+    )
+  return _load(
+      file_url=_TOKENIZERS[tokenizer_name].model_url,
+      expected_hash=_TOKENIZERS[tokenizer_name].model_hash,
+  )
+
+
+@functools.lru_cache()
+def load_model_proto(
+    tokenizer_name: str,
+) -> sentencepiece_model_pb2.ModelProto:
+  """Loads model proto from the given tokenizer name."""
+  model_proto = sentencepiece_model_pb2.ModelProto()
+  model_proto.ParseFromString(_load_model_proto_bytes(tokenizer_name))
+  return model_proto
+
+
+def get_tokenizer_name(model_name: str) -> str:
+  """Gets the tokenizer name for the given model name."""
+  if model_name in _GEMINI_MODELS_TO_TOKENIZER_NAMES.keys():
+    return _GEMINI_MODELS_TO_TOKENIZER_NAMES[model_name]
+  if model_name in _GEMINI_STABLE_MODELS_TO_TOKENIZER_NAMES.keys():
+    return _GEMINI_STABLE_MODELS_TO_TOKENIZER_NAMES[model_name]
+  raise ValueError(
+      f"Model {model_name} is not supported. Supported models: {', '.join(_GEMINI_MODELS_TO_TOKENIZER_NAMES.keys())}, {', '.join(_GEMINI_STABLE_MODELS_TO_TOKENIZER_NAMES.keys())}.\n"  # pylint: disable=line-too-long
+  )
+
+
+@functools.lru_cache()
+def get_sentencepiece(tokenizer_name: str) -> spm.SentencePieceProcessor:
+  """Loads sentencepiece tokenizer from the given tokenizer name."""
+  processor = spm.SentencePieceProcessor()
+  processor.LoadFromSerializedProto(_load_model_proto_bytes(tokenizer_name))
+  return processor

google/genai/_mcp_utils.py

@@ -0,0 +1,117 @@
+# 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.
+#
+
+"""Utils for working with MCP tools."""
+
+from importlib.metadata import PackageNotFoundError, version
+import typing
+from typing import Any
+
+from . import _common
+from . import types
+
+if typing.TYPE_CHECKING:
+  from mcp.types import Tool as McpTool
+  from mcp import ClientSession as McpClientSession
+else:
+  McpClientSession: typing.Type = Any
+  McpTool: typing.Type = Any
+  try:
+    from mcp.types import Tool as McpTool
+    from mcp import ClientSession as McpClientSession
+  except ImportError:
+    McpTool = None
+    McpClientSession = None
+
+
+def mcp_to_gemini_tool(tool: McpTool) -> types.Tool:
+  """Translates an MCP tool to a Google GenAI tool."""
+  return types.Tool(
+      function_declarations=[{
+          "name": tool.name,
+          "description": tool.description,
+          "parameters": types.Schema.from_json_schema(
+              json_schema=types.JSONSchema(
+                  **_filter_to_supported_schema(tool.inputSchema)
+              )
+          ),
+      }]
+  )
+
+
+def mcp_to_gemini_tools(tools: list[McpTool]) -> list[types.Tool]:
+  """Translates a list of MCP tools to a list of Google GenAI tools."""
+  return [mcp_to_gemini_tool(tool) for tool in tools]
+
+
+def has_mcp_tool_usage(tools: types.ToolListUnion) -> bool:
+  """Checks whether the list of tools contains any MCP tools or sessions."""
+  if McpClientSession is None:
+    return False
+  for tool in tools:
+    if isinstance(tool, McpTool) or isinstance(tool, McpClientSession):
+      return True
+  return False
+
+
+def has_mcp_session_usage(tools: types.ToolListUnion) -> bool:
+  """Checks whether the list of tools contains any MCP sessions."""
+  if McpClientSession is None:
+    return False
+  for tool in tools:
+    if isinstance(tool, McpClientSession):
+      return True
+  return False
+
+
+def set_mcp_usage_header(headers: dict[str, str]) -> None:
+  """Sets the MCP version label in the Google API client header."""
+  if McpClientSession is None:
+    return
+  try:
+    version_label = version("mcp")
+  except PackageNotFoundError:
+    version_label = "0.0.0"
+  existing_header = headers.get("x-goog-api-client", "")
+  headers["x-goog-api-client"] = (
+      existing_header + f" mcp_used/{version_label}"
+  ).lstrip()
+
+
+def _filter_to_supported_schema(
+    schema: _common.StringDict,
+) -> _common.StringDict:
+  """Filters the schema to only include fields that are supported by JSONSchema."""
+  supported_fields: set[str] = set(types.JSONSchema.model_fields.keys())
+  schema_field_names: tuple[str] = ("items",)  # 'additional_properties' to come
+  list_schema_field_names: tuple[str] = (
+      "any_of",  # 'one_of', 'all_of', 'not' to come
+  )
+  dict_schema_field_names: tuple[str] = ("properties",)  # 'defs' to come
+  for field_name, field_value in schema.items():
+    if field_name in schema_field_names:
+      schema[field_name] = _filter_to_supported_schema(field_value)
+    elif field_name in list_schema_field_names:
+      schema[field_name] = [
+          _filter_to_supported_schema(value) for value in field_value
+      ]
+    elif field_name in dict_schema_field_names:
+      schema[field_name] = {
+          key: _filter_to_supported_schema(value)
+          for key, value in field_value.items()
+      }
+  return {
+      key: value for key, value in schema.items() if key in supported_fields
+  }