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>
This commit is contained in:
Eduard van Valkenburg
2026-04-16 21:38:50 +02:00
committed by GitHub
Unverified
parent c14beedb3a
commit 90a633967c
14 changed files with 478 additions and 65 deletions
+2 -1
View File
@@ -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
+19 -2
View File
@@ -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
@@ -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__",
@@ -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,
+4 -2
View File
@@ -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`
@@ -0,0 +1 @@
# Copyright (c) Microsoft. All rights reserved.
@@ -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.
"""
+18 -3
View File
@@ -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.
"""
@@ -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.
"""
@@ -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.
"""
@@ -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).
"""
@@ -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)