From 90a633967ca60601fc696d335d770f9f05e236e2 Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Thu, 16 Apr 2026 21:38:50 +0200 Subject: [PATCH] Python: Fix Gemini client support for Gemini API and Vertex AI (#5258) * Add Gemini and Vertex AI client support Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Address Gemini PR review feedback Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * removed sample run readme part --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Evan Mattson <35585003+moonbox3@users.noreply.github.com> --- .../packages/core/agent_framework/_tools.py | 3 + python/packages/core/tests/core/test_tools.py | 14 ++ python/packages/gemini/AGENTS.md | 3 +- python/packages/gemini/README.md | 21 +- .../gemini/agent_framework_gemini/__init__.py | 10 +- .../agent_framework_gemini/_chat_client.py | 191 +++++++++++++--- python/packages/gemini/samples/README.md | 6 +- python/packages/gemini/samples/__init__.py | 1 + .../gemini/samples/gemini_advanced.py | 18 +- .../packages/gemini/samples/gemini_basic.py | 21 +- .../samples/gemini_with_code_execution.py | 15 +- .../gemini/samples/gemini_with_google_maps.py | 16 +- .../samples/gemini_with_google_search.py | 15 +- .../gemini/tests/test_gemini_client.py | 209 ++++++++++++++++-- 14 files changed, 478 insertions(+), 65 deletions(-) diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 6cdc74b313..47eefe8da9 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -906,6 +906,9 @@ def _tools_to_dict( # pyright: ignore[reportUnusedFunction] if isinstance(tool_item, FunctionTool): results.append(tool_item.to_json_schema_spec()) continue + if isinstance(tool_item, BaseModel): + results.append(tool_item.model_dump(exclude_none=True)) + continue if isinstance(tool_item, SerializationMixin): results.append(tool_item.to_dict()) continue diff --git a/python/packages/core/tests/core/test_tools.py b/python/packages/core/tests/core/test_tools.py index 143aa95727..91ba663d84 100644 --- a/python/packages/core/tests/core/test_tools.py +++ b/python/packages/core/tests/core/test_tools.py @@ -16,12 +16,26 @@ from agent_framework._middleware import FunctionInvocationContext from agent_framework._tools import ( _parse_annotation, _parse_inputs, + _tools_to_dict, ) from agent_framework.observability import OtelAttr # region FunctionTool and tool decorator tests +def test_tools_to_dict_supports_pydantic_tool_models() -> None: + """Pydantic-based tool specs are serialized without logging parse warnings.""" + + class ProviderTool(BaseModel): + kind: str + enabled: bool = True + note: str | None = None + + result = _tools_to_dict([ProviderTool(kind="google_search")]) + + assert result == [{"kind": "google_search", "enabled": True}] + + def test_tool_decorator(): """Test the tool decorator.""" diff --git a/python/packages/gemini/AGENTS.md b/python/packages/gemini/AGENTS.md index aa12fddf0a..b87406b460 100644 --- a/python/packages/gemini/AGENTS.md +++ b/python/packages/gemini/AGENTS.md @@ -1,6 +1,6 @@ # Gemini Package (agent-framework-gemini) -Integration with Google's Gemini API via the `google-genai` SDK. +Integration with Google's Gemini Developer API and Vertex AI via the `google-genai` SDK. ## Core Classes @@ -8,6 +8,7 @@ Integration with Google's Gemini API via the `google-genai` SDK. - **`GeminiChatClient`** - Full-featured chat client with function invocation, middleware, and telemetry - **`GeminiChatOptions`** - Options TypedDict for Gemini-specific parameters - **`GeminiSettings`** - Settings loaded from environment variables +- **`GoogleGeminiSettings`** - SDK-standard `GOOGLE_*` settings loaded from environment variables - **`ThinkingConfig`** - Configuration for extended thinking ## Gemini-specific Options diff --git a/python/packages/gemini/README.md b/python/packages/gemini/README.md index e72e2f8126..80b7adba73 100644 --- a/python/packages/gemini/README.md +++ b/python/packages/gemini/README.md @@ -12,11 +12,28 @@ The Gemini integration enables Microsoft Agent Framework applications to call Go ## Authentication -Obtain an API key from [Google AI Studio](https://aistudio.google.com/apikey) and set it via environment variable: +The connector supports both `google-genai` authentication modes. + +### Gemini Developer API + +Obtain an API key from [Google AI Studio](https://aistudio.google.com/apikey) and set either the package-prefixed or SDK-standard environment variable: ```bash export GEMINI_API_KEY="your-api-key" -export GEMINI_MODEL="gemini-2.5-flash" +# or: export GOOGLE_API_KEY="your-api-key" +export GEMINI_MODEL="gemini-2.5-flash-lite" +# or: export GOOGLE_MODEL="gemini-2.5-flash-lite" +``` + +### Vertex AI + +Set the standard Vertex AI environment variables used by `google-genai`: + +```bash +export GOOGLE_GENAI_USE_VERTEXAI=true +export GOOGLE_CLOUD_PROJECT="your-project-id" +export GOOGLE_CLOUD_LOCATION="global" +export GOOGLE_MODEL="gemini-2.5-flash-lite" ``` ## Examples diff --git a/python/packages/gemini/agent_framework_gemini/__init__.py b/python/packages/gemini/agent_framework_gemini/__init__.py index 42099ae0b1..7a0d014846 100644 --- a/python/packages/gemini/agent_framework_gemini/__init__.py +++ b/python/packages/gemini/agent_framework_gemini/__init__.py @@ -2,7 +2,14 @@ import importlib.metadata -from ._chat_client import GeminiChatClient, GeminiChatOptions, GeminiSettings, RawGeminiChatClient, ThinkingConfig +from ._chat_client import ( + GeminiChatClient, + GeminiChatOptions, + GeminiSettings, + GoogleGeminiSettings, + RawGeminiChatClient, + ThinkingConfig, +) try: __version__ = importlib.metadata.version(__name__) @@ -13,6 +20,7 @@ __all__ = [ "GeminiChatClient", "GeminiChatOptions", "GeminiSettings", + "GoogleGeminiSettings", "RawGeminiChatClient", "ThinkingConfig", "__version__", diff --git a/python/packages/gemini/agent_framework_gemini/_chat_client.py b/python/packages/gemini/agent_framework_gemini/_chat_client.py index 2f56b3f9a4..b0fa52a676 100644 --- a/python/packages/gemini/agent_framework_gemini/_chat_client.py +++ b/python/packages/gemini/agent_framework_gemini/_chat_client.py @@ -30,6 +30,7 @@ from agent_framework import ( from agent_framework._settings import SecretString, load_settings from agent_framework.observability import ChatTelemetryLayer from google import genai +from google.auth.credentials import Credentials from google.genai import types from pydantic import BaseModel @@ -54,6 +55,7 @@ __all__ = [ "GeminiChatClient", "GeminiChatOptions", "GeminiSettings", + "GoogleGeminiSettings", "RawGeminiChatClient", "ThinkingConfig", ] @@ -161,10 +163,74 @@ class GeminiSettings(TypedDict, total=False): model: str | None +class GoogleGeminiSettings(TypedDict, total=False): + """Google SDK configuration settings loaded from ``GOOGLE_*`` environment variables.""" + + api_key: SecretString | None + model: str | None + genai_use_vertexai: bool | None + cloud_project: str | None + cloud_location: str | None + + # endregion -_GEMINI_SERVICE_URL = "https://generativelanguage.googleapis.com" +_GEMINI_API_BASE_URL = "https://generativelanguage.googleapis.com" +_VERTEX_AI_BASE_URL = "https://aiplatform.googleapis.com" + + +def _resolve_vertexai_mode(client: genai.Client, *, fallback: bool | None = None) -> bool: + """Resolve whether a client targets Vertex AI, preferring the instantiated SDK client state.""" + api_client = getattr(client, "_api_client", None) + vertexai = getattr(api_client, "vertexai", None) + if isinstance(vertexai, bool): + return vertexai + return bool(fallback) + + +def _resolve_service_url(client: genai.Client, *, vertexai: bool) -> str: + """Resolve the base service URL from the instantiated SDK client, with a stable fallback.""" + api_client = getattr(client, "_api_client", None) + http_options = getattr(api_client, "_http_options", None) + base_url = getattr(http_options, "base_url", None) + if isinstance(base_url, str) and base_url: + return base_url.rstrip("/") + return _VERTEX_AI_BASE_URL if vertexai else _GEMINI_API_BASE_URL + + +def _validate_client_auth_configuration( + *, + vertexai: bool | None, + api_key: SecretString | None, + project: str | None, + location: str | None, + credentials: Credentials | None, +) -> None: + """Validate supported auth combinations before instantiating the SDK client.""" + if vertexai is not True: + if api_key is None: + raise ValueError( + "Gemini client requires an API key when Vertex AI is not enabled. " + "Set GOOGLE_API_KEY or GEMINI_API_KEY, or pass api_key explicitly." + ) + return + + if api_key is not None or credentials is not None or (project and location): + return + + if project or location: + raise ValueError( + "Gemini client requires both GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION " + "when Vertex AI is enabled without an API key." + ) + + raise ValueError( + "Gemini client requires Vertex AI credentials or configuration when Vertex AI is enabled. " + "Provide GOOGLE_API_KEY for Vertex AI express mode, pass credentials, or set " + "GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION." + ) + # Keys mapping to a different GenerateContentConfig field name _OPTION_TRANSLATIONS: dict[str, str] = { @@ -210,7 +276,7 @@ class RawGeminiChatClient( BaseChatClient[GeminiChatOptionsT], Generic[GeminiChatOptionsT], ): - """A raw Gemini chat client for the Google Gemini API without function invocation, middleware or telemetry. + """A raw Gemini chat client for Gemini Developer API or Vertex AI. Use this when you want full control over the request pipeline. For instance, to opt out of telemetry, use custom middleware, or compose your own layers. If you want the full-featured @@ -224,6 +290,10 @@ class RawGeminiChatClient( *, api_key: str | None = None, model: str | None = None, + vertexai: bool | None = None, + project: str | None = None, + location: str | None = None, + credentials: Credentials | None = None, env_file_path: str | None = None, env_file_encoding: str | None = None, client: genai.Client | None = None, @@ -232,11 +302,21 @@ class RawGeminiChatClient( """Create a raw Gemini chat client. Args: - api_key: Google AI Studio API key. Falls back to ``GEMINI_API_KEY`` environment variable. - model: Default model identifier. Falls back to ``GEMINI_MODEL`` environment variable. + api_key: Gemini Developer API key. Falls back to environment settings, preferring + ``GOOGLE_API_KEY`` over ``GEMINI_API_KEY``. + model: Default model identifier. Falls back to environment settings, preferring + ``GOOGLE_MODEL`` over ``GEMINI_MODEL``. + vertexai: Whether to use Vertex AI endpoints. Falls back to environment settings, + using ``GOOGLE_GENAI_USE_VERTEXAI`` when not passed explicitly. + project: Google Cloud project ID for Vertex AI. Falls back to environment settings, + using ``GOOGLE_CLOUD_PROJECT`` when not passed explicitly. + location: Vertex AI location. Falls back to environment settings, preferring + using ``GOOGLE_CLOUD_LOCATION`` when not passed explicitly. + credentials: Google Cloud credentials for Vertex AI. When omitted, the SDK can use + Application Default Credentials. env_file_path: Path to a ``.env`` file for credential loading. env_file_encoding: Encoding for the ``.env`` file. - client: Pre-built ``genai.Client`` instance. When provided, ``api_key`` is not required. + client: Pre-built ``genai.Client`` instance. When provided, connector auth settings are not required. additional_properties: Extra properties stored on the client instance. """ settings = load_settings( @@ -247,21 +327,58 @@ class RawGeminiChatClient( env_file_path=env_file_path, env_file_encoding=env_file_encoding, ) + google_settings = load_settings( + GoogleGeminiSettings, + env_prefix="GOOGLE_", + api_key=api_key, + model=model, + genai_use_vertexai=vertexai, + cloud_project=project, + cloud_location=location, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + configured_vertexai = google_settings.get("genai_use_vertexai") if client: self._genai_client = client else: - resolved_key = settings.get("api_key") - if not resolved_key: - raise ValueError( - "Gemini API key is required. Set via api_key parameter or GEMINI_API_KEY environment variable." - ) - self._genai_client = genai.Client( - api_key=resolved_key.get_secret_value(), - http_options={"headers": {"x-goog-api-client": AGENT_FRAMEWORK_USER_AGENT}}, + resolved_key = google_settings.get("api_key") or settings.get("api_key") + resolved_project = google_settings.get("cloud_project") + resolved_location = google_settings.get("cloud_location") + _validate_client_auth_configuration( + vertexai=configured_vertexai, + api_key=resolved_key, + project=resolved_project, + location=resolved_location, + credentials=credentials, ) - self.model = settings.get("model") + client_kwargs: dict[str, Any] = { + "http_options": {"headers": {"x-goog-api-client": AGENT_FRAMEWORK_USER_AGENT}}, + } + if configured_vertexai is not None: + client_kwargs["vertexai"] = configured_vertexai + + if resolved_key is not None and ( + configured_vertexai is not True + or (credentials is None and not (resolved_project and resolved_location)) + ): + client_kwargs["api_key"] = resolved_key.get_secret_value() + + if configured_vertexai is True and resolved_project: + client_kwargs["project"] = resolved_project + + if configured_vertexai is True and resolved_location: + client_kwargs["location"] = resolved_location + if configured_vertexai is True and credentials is not None: + client_kwargs["credentials"] = credentials + + self._genai_client = genai.Client(**client_kwargs) + + self._vertexai = _resolve_vertexai_mode(self._genai_client, fallback=configured_vertexai) + self._service_url = _resolve_service_url(self._genai_client, vertexai=self._vertexai) + self.model = google_settings.get("model") or settings.get("model") super().__init__(additional_properties=additional_properties) @@ -414,12 +531,12 @@ class RawGeminiChatClient( @override def service_url(self) -> str: - """Return the base URL of the Gemini API service. + """Return the base URL of the configured Gemini or Vertex AI service. Returns: - The Gemini API base URL. + The resolved service base URL. """ - return _GEMINI_SERVICE_URL + return self._service_url # region Request preparation @@ -528,15 +645,16 @@ class RawGeminiChatClient( call_id = content.call_id or self._generate_tool_call_id() if content.name: call_id_to_name[call_id] = content.name - parts.append( - types.Part( - function_call=types.FunctionCall( - id=call_id, - name=content.name or "", - args=content.parse_arguments() or {}, - ) - ) + function_call = types.FunctionCall( + id=call_id, + name=content.name or "", + args=content.parse_arguments() or {}, ) + raw_part = content.raw_representation + if isinstance(raw_part, types.Part) and raw_part.function_call is not None: + parts.append(raw_part.model_copy(update={"function_call": function_call}, deep=True)) + else: + parts.append(types.Part(function_call=function_call)) case _: logger.debug("Skipping unsupported content type for Gemini: %s", content.type) return parts @@ -889,7 +1007,7 @@ class GeminiChatClient( RawGeminiChatClient[GeminiChatOptionsT], Generic[GeminiChatOptionsT], ): - """Gemini chat client for the Google Gemini API with function invocation, middleware, and telemetry. + """Gemini chat client for Gemini Developer API or Vertex AI with function invocation, middleware, and telemetry. This is the recommended client for most use cases. It builds on ``RawGeminiChatClient`` and adds: @@ -908,6 +1026,10 @@ class GeminiChatClient( *, api_key: str | None = None, model: str | None = None, + vertexai: bool | None = None, + project: str | None = None, + location: str | None = None, + credentials: Credentials | None = None, env_file_path: str | None = None, env_file_encoding: str | None = None, client: genai.Client | None = None, @@ -918,11 +1040,18 @@ class GeminiChatClient( """Create a Gemini chat client. Args: - api_key: The Google AI Studio API key. Falls back to ``GEMINI_API_KEY`` environment variable. - model: Default model identifier. Falls back to ``GEMINI_MODEL`` environment variable. + api_key: Gemini Developer API key. Falls back to environment settings, preferring + ``GOOGLE_API_KEY`` over ``GEMINI_API_KEY``. + model: Default model identifier. Falls back to environment settings, preferring + ``GOOGLE_MODEL`` over ``GEMINI_MODEL``. + vertexai: Whether to use Vertex AI endpoints. Falls back to ``GOOGLE_GENAI_USE_VERTEXAI``. + project: Google Cloud project ID for Vertex AI. Falls back to ``GOOGLE_CLOUD_PROJECT``. + location: Vertex AI location. Falls back to ``GOOGLE_CLOUD_LOCATION``. + credentials: Google Cloud credentials for Vertex AI. When omitted, the SDK can use + Application Default Credentials. env_file_path: Path to a ``.env`` file for credential loading. env_file_encoding: Encoding for the ``.env`` file. - client: Pre-built ``genai.Client`` instance. When provided, ``api_key`` is not required. + client: Pre-built ``genai.Client`` instance. When provided, connector auth settings are not required. additional_properties: Extra properties stored on the client instance. middleware: Optional middleware chain applied to every call. function_invocation_configuration: Optional configuration for the function invocation loop. @@ -930,6 +1059,10 @@ class GeminiChatClient( super().__init__( api_key=api_key, model=model, + vertexai=vertexai, + project=project, + location=location, + credentials=credentials, env_file_path=env_file_path, env_file_encoding=env_file_encoding, client=client, diff --git a/python/packages/gemini/samples/README.md b/python/packages/gemini/samples/README.md index 28fb05abeb..c1687368b8 100644 --- a/python/packages/gemini/samples/README.md +++ b/python/packages/gemini/samples/README.md @@ -14,5 +14,7 @@ This folder contains examples demonstrating how to use Google Gemini models with ## Environment Variables -- `GEMINI_API_KEY`: Your Google AI Studio API key (get one from [Google AI Studio](https://aistudio.google.com/apikey)) -- `GEMINI_MODEL`: The Gemini model to use (e.g., `gemini-2.5-flash`, `gemini-2.5-pro`) +- `GOOGLE_MODEL` or `GEMINI_MODEL`: The Gemini model to use (for example, + `gemini-2.5-flash-lite` or `gemini-2.5-pro`) +- For Gemini Developer API: `GEMINI_API_KEY` or `GOOGLE_API_KEY` +- For Vertex AI: `GOOGLE_GENAI_USE_VERTEXAI=true`, `GOOGLE_CLOUD_PROJECT`, and `GOOGLE_CLOUD_LOCATION` diff --git a/python/packages/gemini/samples/__init__.py b/python/packages/gemini/samples/__init__.py index e69de29bb2..2a50eae894 100644 --- a/python/packages/gemini/samples/__init__.py +++ b/python/packages/gemini/samples/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Microsoft. All rights reserved. diff --git a/python/packages/gemini/samples/gemini_advanced.py b/python/packages/gemini/samples/gemini_advanced.py index a8afbecbd2..a38a59773f 100644 --- a/python/packages/gemini/samples/gemini_advanced.py +++ b/python/packages/gemini/samples/gemini_advanced.py @@ -4,9 +4,9 @@ Allows the model to reason through complex problems before responding. -Requires the following environment variables to be set: -- GEMINI_API_KEY -- GEMINI_MODEL +Requires ``GOOGLE_MODEL`` or ``GEMINI_MODEL`` and either Gemini Developer API credentials +(``GEMINI_API_KEY`` or ``GOOGLE_API_KEY``) or Vertex AI settings +(``GOOGLE_GENAI_USE_VERTEXAI``, ``GOOGLE_CLOUD_PROJECT``, and ``GOOGLE_CLOUD_LOCATION``). """ import asyncio @@ -23,10 +23,12 @@ async def main() -> None: """Example of extended thinking with a Python version comparison question.""" print("=== Extended thinking ===") + # 1. Configure Gemini extended thinking for a reasoning-heavy request. options: GeminiChatOptions = { "thinking_config": ThinkingConfig(thinking_budget=2048), } + # 2. Create the agent with the Gemini chat client and default thinking options. agent = Agent( client=GeminiChatClient(), name="PythonAgent", @@ -34,6 +36,7 @@ async def main() -> None: default_options=options, ) + # 3. Stream the answer so you can see the final response as it arrives. query = "What new language features were introduced in Python between 3.10 and 3.14?" print(f"User: {query}") print("Agent: ", end="", flush=True) @@ -45,3 +48,12 @@ async def main() -> None: if __name__ == "__main__": asyncio.run(main()) + +""" +Sample output: +=== Extended thinking === +User: What new language features were introduced in Python between 3.10 and 3.14? +Agent: Python 3.11 introduced exception groups and TaskGroup. +Python 3.12 added PEP 695 type parameter syntax. +Python 3.13-3.14 continued improving typing, performance, and developer ergonomics. +""" diff --git a/python/packages/gemini/samples/gemini_basic.py b/python/packages/gemini/samples/gemini_basic.py index af1b5f1076..81e386beda 100644 --- a/python/packages/gemini/samples/gemini_basic.py +++ b/python/packages/gemini/samples/gemini_basic.py @@ -4,9 +4,9 @@ Covers both non-streaming and streaming responses. -Requires the following environment variables to be set: -- GEMINI_API_KEY -- GEMINI_MODEL +Requires ``GOOGLE_MODEL`` or ``GEMINI_MODEL`` and either Gemini Developer API credentials +(``GEMINI_API_KEY`` or ``GOOGLE_API_KEY``) or Vertex AI settings +(``GOOGLE_GENAI_USE_VERTEXAI``, ``GOOGLE_CLOUD_PROJECT``, and ``GOOGLE_CLOUD_LOCATION``). """ import asyncio @@ -35,6 +35,7 @@ async def non_streaming_example() -> None: """Runs the agent and waits for the complete response before printing it.""" print("=== Non-streaming ===") + # 1. Create the agent with the Gemini chat client and local weather tool. agent = Agent( client=GeminiChatClient(), name="WeatherAgent", @@ -42,6 +43,7 @@ async def non_streaming_example() -> None: tools=[get_weather], ) + # 2. Ask the agent for a single weather lookup and print the final response. query = "What's the weather like in Karlsruhe, Germany?" print(f"User: {query}") result = await agent.run(query) @@ -52,6 +54,7 @@ async def streaming_example() -> None: """Runs the agent and prints each chunk as it is received.""" print("=== Streaming ===") + # 1. Create the same agent configuration for a streaming tool-call example. agent = Agent( client=GeminiChatClient(), name="WeatherAgent", @@ -59,6 +62,7 @@ async def streaming_example() -> None: tools=[get_weather], ) + # 2. Ask a multi-location question and stream the model output as it arrives. query = "What's the weather like in Portland and in Paris?" print(f"User: {query}") print("Agent: ", end="", flush=True) @@ -76,3 +80,14 @@ async def main() -> None: if __name__ == "__main__": asyncio.run(main()) + +""" +Sample output: +=== Non-streaming === +User: What's the weather like in Karlsruhe, Germany? +Result: The weather in Karlsruhe, Germany is currently sunny with a high of 16°C. + +=== Streaming === +User: What's the weather like in Portland and in Paris? +Agent: In Portland, it is currently rainy with a high of 11°C. In Paris, it is cloudy with a high of 27°C. +""" diff --git a/python/packages/gemini/samples/gemini_with_code_execution.py b/python/packages/gemini/samples/gemini_with_code_execution.py index e41c63637c..ed4ae5a387 100644 --- a/python/packages/gemini/samples/gemini_with_code_execution.py +++ b/python/packages/gemini/samples/gemini_with_code_execution.py @@ -4,9 +4,9 @@ Allows the model to write and run code in a sandboxed environment to answer questions. -Requires the following environment variables to be set: -- GEMINI_API_KEY -- GEMINI_MODEL +Requires ``GOOGLE_MODEL`` or ``GEMINI_MODEL`` and either Gemini Developer API credentials +(``GEMINI_API_KEY`` or ``GOOGLE_API_KEY``) or Vertex AI settings +(``GOOGLE_GENAI_USE_VERTEXAI``, ``GOOGLE_CLOUD_PROJECT``, and ``GOOGLE_CLOUD_LOCATION``). """ import asyncio @@ -23,6 +23,7 @@ async def main() -> None: """Run the code execution example.""" print("=== Code execution ===") + # 1. Create the agent with Gemini and the built-in code execution tool. agent = Agent( client=GeminiChatClient(), name="CodeAgent", @@ -30,6 +31,7 @@ async def main() -> None: tools=[GeminiChatClient.get_code_interpreter_tool()], ) + # 2. Ask for a computed answer and stream the generated code and final result. query = "What are the first 20 prime numbers? Compute them in code." print(f"User: {query}") print("Agent: ", end="", flush=True) @@ -41,3 +43,10 @@ async def main() -> None: if __name__ == "__main__": asyncio.run(main()) + +""" +Sample output: +=== Code execution === +User: What are the first 20 prime numbers? Compute them in code. +Agent: The first 20 prime numbers are 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, and 71. +""" diff --git a/python/packages/gemini/samples/gemini_with_google_maps.py b/python/packages/gemini/samples/gemini_with_google_maps.py index 8083655b7d..92e7b3e708 100644 --- a/python/packages/gemini/samples/gemini_with_google_maps.py +++ b/python/packages/gemini/samples/gemini_with_google_maps.py @@ -4,9 +4,9 @@ Allows Gemini to retrieve location and mapping information before responding. -Requires the following environment variables to be set: -- GEMINI_API_KEY -- GEMINI_MODEL +Requires ``GOOGLE_MODEL`` or ``GEMINI_MODEL`` and either Gemini Developer API credentials +(``GEMINI_API_KEY`` or ``GOOGLE_API_KEY``) or Vertex AI settings +(``GOOGLE_GENAI_USE_VERTEXAI``, ``GOOGLE_CLOUD_PROJECT``, and ``GOOGLE_CLOUD_LOCATION``). """ import asyncio @@ -23,6 +23,7 @@ async def main() -> None: """Run the Google Maps grounding example.""" print("=== Google Maps grounding ===") + # 1. Create the agent with Gemini and the built-in Google Maps grounding tool. agent = Agent( client=GeminiChatClient(), name="MapsAgent", @@ -30,6 +31,7 @@ async def main() -> None: tools=[GeminiChatClient.get_maps_grounding_tool()], ) + # 2. Ask a location-aware question and stream the grounded answer. query = "What are some highly rated restaurants in the city center of Karlsruhe, Germany?" print(f"User: {query}") print("Agent: ", end="", flush=True) @@ -41,3 +43,11 @@ async def main() -> None: if __name__ == "__main__": asyncio.run(main()) + +""" +Sample output: +=== Google Maps grounding === +User: What are some highly rated restaurants in the city center of Karlsruhe, Germany? +Agent: Here are several highly rated restaurants near Karlsruhe city center, +along with their cuisine styles and approximate walking distance. +""" diff --git a/python/packages/gemini/samples/gemini_with_google_search.py b/python/packages/gemini/samples/gemini_with_google_search.py index 741f4d4d27..9c03119cdf 100644 --- a/python/packages/gemini/samples/gemini_with_google_search.py +++ b/python/packages/gemini/samples/gemini_with_google_search.py @@ -4,9 +4,9 @@ Allows Gemini to retrieve up-to-date information from the web before responding. -Requires the following environment variables to be set: -- GEMINI_API_KEY -- GEMINI_MODEL +Requires ``GOOGLE_MODEL`` or ``GEMINI_MODEL`` and either Gemini Developer API credentials +(``GEMINI_API_KEY`` or ``GOOGLE_API_KEY``) or Vertex AI settings +(``GOOGLE_GENAI_USE_VERTEXAI``, ``GOOGLE_CLOUD_PROJECT``, and ``GOOGLE_CLOUD_LOCATION``). """ import asyncio @@ -23,6 +23,7 @@ async def main() -> None: """Run the Google Search grounding example.""" print("=== Google Search grounding ===") + # 1. Create the agent with Gemini and the built-in Google Search grounding tool. agent = Agent( client=GeminiChatClient(), name="SearchAgent", @@ -30,6 +31,7 @@ async def main() -> None: tools=[GeminiChatClient.get_web_search_tool()], ) + # 2. Ask a current-events style question and stream the grounded answer. query = "What is the latest stable release of the .NET SDK?" print(f"User: {query}") print("Agent: ", end="", flush=True) @@ -41,3 +43,10 @@ async def main() -> None: if __name__ == "__main__": asyncio.run(main()) + +""" +Sample output: +=== Google Search grounding === +User: What is the latest stable release of the .NET SDK? +Agent: As of April 14, 2026, the latest stable release of the .NET SDK is .NET 10.0 (SDK 10.0.201). +""" diff --git a/python/packages/gemini/tests/test_gemini_client.py b/python/packages/gemini/tests/test_gemini_client.py index 07248cb7e5..d5fcf5dbe0 100644 --- a/python/packages/gemini/tests/test_gemini_client.py +++ b/python/packages/gemini/tests/test_gemini_client.py @@ -15,12 +15,28 @@ from pydantic import BaseModel from agent_framework_gemini import GeminiChatClient, GeminiChatOptions, ThinkingConfig -skip_if_no_api_key = pytest.mark.skipif( - not os.getenv("GEMINI_API_KEY"), - reason="GEMINI_API_KEY not set; skipping integration tests.", + +def _has_gemini_integration_credentials() -> bool: + """Return whether integration credentials for either Gemini API or Vertex AI appear to be configured.""" + if os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY"): + return True + + if os.getenv("GOOGLE_GENAI_USE_VERTEXAI", "").lower() in {"true", "1", "yes", "on"}: + return bool( + os.getenv("GOOGLE_CLOUD_PROJECT") + or os.getenv("GOOGLE_APPLICATION_CREDENTIALS") + or os.getenv("GOOGLE_API_KEY") + ) + + return False + + +skip_if_no_credentials = pytest.mark.skipif( + not _has_gemini_integration_credentials(), + reason="Gemini Developer API or Vertex AI credentials not set; skipping integration tests.", ) -_TEST_MODEL = "gemini-2.5-flash" +_TEST_MODEL = os.getenv("GOOGLE_MODEL") or os.getenv("GEMINI_MODEL", "gemini-2.5-flash-lite") # stub helpers @@ -89,6 +105,7 @@ def _make_response( candidate.finish_reason = None response.candidates = [candidate] + response.finish_reason = finish_reason response.model_version = model_version if prompt_tokens is not None or output_tokens is not None: @@ -115,6 +132,8 @@ def _make_gemini_client( ) -> tuple[GeminiChatClient, MagicMock]: """Return a (GeminiChatClient, mock_genai_client) pair.""" mock = mock_client or MagicMock() + mock._api_client.vertexai = False + mock._api_client._http_options.base_url = "https://generativelanguage.googleapis.com/" client = GeminiChatClient(client=mock, model=model) return client, mock @@ -135,12 +154,134 @@ def test_client_created_from_api_key(monkeypatch: pytest.MonkeyPatch) -> None: assert client.model == "gemini-2.5-flash" -def test_missing_api_key_raises_when_no_client_injected(monkeypatch: pytest.MonkeyPatch) -> None: - """Raises ValueError at construction when neither an API key nor a pre-built client is available.""" +def test_client_created_from_google_api_key_env(monkeypatch: pytest.MonkeyPatch) -> None: + """Initialises successfully when the SDK-standard Google API key environment variable is set.""" monkeypatch.delenv("GEMINI_API_KEY", raising=False) monkeypatch.delenv("GEMINI_MODEL", raising=False) + monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False) + monkeypatch.delenv("GOOGLE_CLOUD_PROJECT", raising=False) + monkeypatch.delenv("GOOGLE_CLOUD_LOCATION", raising=False) + monkeypatch.setenv("GOOGLE_API_KEY", "test-key-123") + monkeypatch.setenv("GOOGLE_MODEL", "gemini-2.5-flash-lite") - with pytest.raises(ValueError, match="GEMINI_API_KEY"): + mock_client = MagicMock() + mock_client._api_client.vertexai = False + mock_client._api_client._http_options.base_url = "https://generativelanguage.googleapis.com/" + + with patch("agent_framework_gemini._chat_client.genai.Client") as client_factory: + client_factory.return_value = mock_client + client = GeminiChatClient() + + assert client_factory.call_args.kwargs["api_key"] == "test-key-123" + assert "vertexai" not in client_factory.call_args.kwargs + assert client.model == "gemini-2.5-flash-lite" + assert client.service_url() == "https://generativelanguage.googleapis.com" + + +def test_client_created_from_vertex_ai_env(monkeypatch: pytest.MonkeyPatch) -> None: + """Initialises a Vertex AI client when the SDK-standard Vertex AI environment variables are set.""" + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "true") + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "test-project") + monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "global") + + mock_client = MagicMock() + mock_client._api_client.vertexai = True + mock_client._api_client._http_options.base_url = "https://aiplatform.googleapis.com/" + + with patch("agent_framework_gemini._chat_client.genai.Client", return_value=mock_client) as client_factory: + client = GeminiChatClient() + + assert client_factory.call_args.kwargs["vertexai"] is True + assert client_factory.call_args.kwargs["project"] == "test-project" + assert client_factory.call_args.kwargs["location"] == "global" + assert "api_key" not in client_factory.call_args.kwargs + assert client.service_url() == "https://aiplatform.googleapis.com" + + +def test_google_settings_take_precedence_over_gemini_aliases(monkeypatch: pytest.MonkeyPatch) -> None: + """Prefers SDK-standard ``GOOGLE_*`` settings when both env families are present.""" + monkeypatch.setenv("GEMINI_API_KEY", "gemini-key") + monkeypatch.setenv("GEMINI_MODEL", "gemini-model") + monkeypatch.setenv("GOOGLE_API_KEY", "google-key") + monkeypatch.setenv("GOOGLE_MODEL", "google-model") + monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "true") + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "google-project") + monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "global") + + mock_client = MagicMock() + mock_client._api_client.vertexai = True + mock_client._api_client._http_options.base_url = "https://aiplatform.googleapis.com/" + + with patch("agent_framework_gemini._chat_client.genai.Client", return_value=mock_client) as client_factory: + client = GeminiChatClient() + + assert client_factory.call_args.kwargs["vertexai"] is True + assert client_factory.call_args.kwargs["project"] == "google-project" + assert client_factory.call_args.kwargs["location"] == "global" + assert "api_key" not in client_factory.call_args.kwargs + assert client.model == "google-model" + assert client.service_url() == "https://aiplatform.googleapis.com" + + +def test_missing_api_key_raises_when_no_client_injected(monkeypatch: pytest.MonkeyPatch) -> None: + """Raises ValueError at construction when neither Gemini API nor Vertex AI settings are available.""" + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + monkeypatch.delenv("GEMINI_MODEL", raising=False) + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False) + monkeypatch.delenv("GOOGLE_CLOUD_PROJECT", raising=False) + monkeypatch.delenv("GOOGLE_CLOUD_LOCATION", raising=False) + + with pytest.raises(ValueError, match="requires an API key when Vertex AI is not enabled"): + GeminiChatClient(model="gemini-2.5-flash") + + +def test_vertex_ai_express_mode_uses_api_key(monkeypatch: pytest.MonkeyPatch) -> None: + """Passes the API key in Vertex AI express mode when no project/location pair is configured.""" + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + monkeypatch.delenv("GEMINI_MODEL", raising=False) + monkeypatch.setenv("GOOGLE_API_KEY", "test-key-123") + monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "true") + monkeypatch.delenv("GOOGLE_CLOUD_PROJECT", raising=False) + monkeypatch.delenv("GOOGLE_CLOUD_LOCATION", raising=False) + + mock_client = MagicMock() + mock_client._api_client.vertexai = True + mock_client._api_client._http_options.base_url = "https://aiplatform.googleapis.com/" + + with patch("agent_framework_gemini._chat_client.genai.Client", return_value=mock_client) as client_factory: + client = GeminiChatClient(model="gemini-2.5-flash-lite") + + assert client_factory.call_args.kwargs["vertexai"] is True + assert client_factory.call_args.kwargs["api_key"] == "test-key-123" + assert "project" not in client_factory.call_args.kwargs + assert "location" not in client_factory.call_args.kwargs + assert client.service_url() == "https://aiplatform.googleapis.com" + + +def test_vertex_ai_requires_configuration(monkeypatch: pytest.MonkeyPatch) -> None: + """Raises a deterministic error when Vertex AI is enabled without any auth configuration.""" + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "true") + monkeypatch.delenv("GOOGLE_CLOUD_PROJECT", raising=False) + monkeypatch.delenv("GOOGLE_CLOUD_LOCATION", raising=False) + + with pytest.raises(ValueError, match="requires Vertex AI credentials or configuration"): + GeminiChatClient(model="gemini-2.5-flash") + + +def test_vertex_ai_requires_project_and_location_together(monkeypatch: pytest.MonkeyPatch) -> None: + """Raises a deterministic error when only one Vertex AI location setting is present.""" + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "true") + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "test-project") + monkeypatch.delenv("GOOGLE_CLOUD_LOCATION", raising=False) + + with pytest.raises(ValueError, match="requires both GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION"): GeminiChatClient(model="gemini-2.5-flash") @@ -495,6 +636,30 @@ async def test_thinking_parts_are_silently_skipped() -> None: assert response.messages[0].text == "The answer is 42." +def test_function_call_part_preserves_thought_signature_from_raw_part() -> None: + """Reuses the original Gemini Part so tool loops retain thought_signature metadata.""" + client, _ = _make_gemini_client() + raw_part = types.Part( + function_call=types.FunctionCall(id="call-1", name="get_weather", args={"location": "Paris"}), + thought_signature=b"sig-123", + ) + content = Content.from_function_call( + call_id="call-1", + name="get_weather", + arguments={"location": "Paris"}, + raw_representation=raw_part, + ) + + parts = client._convert_message_contents([content], {}) + + assert len(parts) == 1 + assert parts[0].thought_signature == b"sig-123" + assert parts[0].function_call is not None + assert parts[0].function_call.id == "call-1" + assert parts[0].function_call.name == "get_weather" + assert parts[0].function_call.args == {"location": "Paris"} + + # code execution parts @@ -1283,12 +1448,26 @@ def test_service_url() -> None: assert client.service_url() == "https://generativelanguage.googleapis.com" +def test_service_url_falls_back_when_sdk_base_url_is_unavailable() -> None: + """Falls back to the known service URL when the SDK client does not expose a base URL.""" + gemini_sdk_client = MagicMock() + gemini_sdk_client._api_client.vertexai = False + gemini_client = GeminiChatClient(client=gemini_sdk_client, model="gemini-2.5-flash") + + vertex_sdk_client = MagicMock() + vertex_sdk_client._api_client.vertexai = True + vertex_client = GeminiChatClient(client=vertex_sdk_client, model="gemini-2.5-flash") + + assert gemini_client.service_url() == "https://generativelanguage.googleapis.com" + assert vertex_client.service_url() == "https://aiplatform.googleapis.com" + + # integration tests @pytest.mark.flaky @pytest.mark.integration -@skip_if_no_api_key +@skip_if_no_credentials async def test_integration_basic_chat() -> None: """Basic request/response round-trip returns a non-empty text reply.""" client = GeminiChatClient(model=_TEST_MODEL) @@ -1302,7 +1481,7 @@ async def test_integration_basic_chat() -> None: @pytest.mark.flaky @pytest.mark.integration -@skip_if_no_api_key +@skip_if_no_credentials async def test_integration_streaming() -> None: """Streaming yields multiple chunks that together form a non-empty response.""" client = GeminiChatClient(model=_TEST_MODEL) @@ -1319,7 +1498,7 @@ async def test_integration_streaming() -> None: @pytest.mark.flaky @pytest.mark.integration -@skip_if_no_api_key +@skip_if_no_credentials async def test_integration_structured_output() -> None: """Structured output with a Pydantic response_format returns a parsed value via response.value.""" @@ -1340,7 +1519,7 @@ async def test_integration_structured_output() -> None: @pytest.mark.flaky @pytest.mark.integration -@skip_if_no_api_key +@skip_if_no_credentials async def test_integration_tool_calling() -> None: """Model invokes the registered tool when asked a question that requires it.""" @@ -1363,7 +1542,7 @@ async def test_integration_tool_calling() -> None: @pytest.mark.flaky @pytest.mark.integration -@skip_if_no_api_key +@skip_if_no_credentials async def test_integration_thinking_config() -> None: """Model accepts a thinking budget and returns a non-empty text reply.""" options: GeminiChatOptions = {"thinking_config": ThinkingConfig(thinking_budget=512)} @@ -1380,7 +1559,7 @@ async def test_integration_thinking_config() -> None: @pytest.mark.flaky @pytest.mark.integration -@skip_if_no_api_key +@skip_if_no_credentials async def test_integration_google_search_grounding() -> None: """Google Search grounding returns a non-empty response for a current-events question.""" client = GeminiChatClient(model=_TEST_MODEL) @@ -1396,7 +1575,7 @@ async def test_integration_google_search_grounding() -> None: @pytest.mark.flaky @pytest.mark.integration -@skip_if_no_api_key +@skip_if_no_credentials async def test_integration_google_maps_grounding() -> None: """Google Maps grounding returns a non-empty response for a location-based question.""" client = GeminiChatClient(model=_TEST_MODEL) @@ -1417,7 +1596,7 @@ async def test_integration_google_maps_grounding() -> None: @pytest.mark.flaky @pytest.mark.integration -@skip_if_no_api_key +@skip_if_no_credentials async def test_integration_code_execution() -> None: """Code execution tool produces a non-empty response for a computation request.""" client = GeminiChatClient(model=_TEST_MODEL)