Skip to content

Commit 3551171

Browse files
markmcdmarkmcd
authored
Python: Add User-Agent to google-genai (#13703)
### Motivation and Context Currently, the Python connectors for Google AI and Vertex AI don't pass the standard semantic-kernel user agent string. This PR adds the user agent header so that the Google API calls properly identify Semantic Kernel. This aligns with the dotnet implementation, and our (Google DeepMind's) [best practices](https://ai.google.dev/gemini-api/docs/partner-integration). ### Description Sends the `User-Agent` header conditionally in `GoogleAIChatCompletion`, `GoogleAITextCompletion`, and `GoogleAITextEmbedding`, matching the OpenAI implementation as much as possible (gated behind telemetry flag, using `SEMANTIC_KERNEL_USER_AGENT`). ### Contribution Checklist - [x] The code builds clean without any errors or warnings - [x] The PR follows the [SK Contribution Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md) and the [pre-submission formatting script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts) raises no violations - [x] All unit tests pass, and I have added new tests where possible - [x] I didn't break anyone 😄.
1 parent 872d29e commit 3551171

5 files changed

Lines changed: 211 additions & 6 deletions

File tree

‎python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_base.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,12 @@
33
from abc import ABC
44
from typing import ClassVar
55

6-
from google.genai import Client
6+
from google.genai import Client, types
77

88
from semantic_kernel.connectors.ai.google.google_ai.google_ai_settings import GoogleAISettings
9+
from semantic_kernel.const import USER_AGENT
910
from semantic_kernel.kernel_pydantic import KernelBaseModel
11+
from semantic_kernel.utils.telemetry.user_agent import IS_TELEMETRY_ENABLED, SEMANTIC_KERNEL_USER_AGENT
1012

1113

1214
class GoogleAIBase(KernelBaseModel, ABC):
@@ -17,3 +19,14 @@ class GoogleAIBase(KernelBaseModel, ABC):
1719
service_settings: GoogleAISettings
1820

1921
client: Client | None = None
22+
23+
def _get_http_options(self) -> types.HttpOptionsDict | None:
24+
"""Get the HTTP options for the Google AI client.
25+
26+
Returns:
27+
The HTTP options dictionary, or None if telemetry is disabled.
28+
"""
29+
if not IS_TELEMETRY_ENABLED:
30+
return None
31+
32+
return types.HttpOptionsDict(headers={USER_AGENT: SEMANTIC_KERNEL_USER_AGENT})

‎python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_chat_completion.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -171,10 +171,14 @@ async def _generate_content(client: Client) -> GenerateContentResponse:
171171
vertexai=True,
172172
project=self.service_settings.cloud_project_id,
173173
location=self.service_settings.cloud_region,
174+
http_options=self._get_http_options(),
174175
) as client:
175176
response: GenerateContentResponse = await _generate_content(client) # type: ignore[no-redef]
176177
else:
177-
with Client(api_key=self.service_settings.api_key.get_secret_value()) as client: # type: ignore[union-attr]
178+
with Client(
179+
api_key=self.service_settings.api_key.get_secret_value(), # type: ignore[union-attr]
180+
http_options=self._get_http_options(),
181+
) as client:
178182
response: GenerateContentResponse = await _generate_content(client) # type: ignore[no-redef]
179183

180184
return [self._create_chat_message_content(response, candidate) for candidate in response.candidates] # type: ignore
@@ -218,14 +222,18 @@ async def _generate_content_stream(client: Client) -> AsyncGenerator[GenerateCon
218222
vertexai=True,
219223
project=self.service_settings.cloud_project_id,
220224
location=self.service_settings.cloud_region,
225+
http_options=self._get_http_options(),
221226
) as client:
222227
async for chunk in _generate_content_stream(client):
223228
yield [
224229
self._create_streaming_chat_message_content(chunk, candidate, function_invoke_attempt)
225230
for candidate in chunk.candidates # type: ignore
226231
]
227232
else:
228-
with Client(api_key=self.service_settings.api_key.get_secret_value()) as client: # type: ignore[union-attr]
233+
with Client(
234+
api_key=self.service_settings.api_key.get_secret_value(), # type: ignore[union-attr]
235+
http_options=self._get_http_options(),
236+
) as client:
229237
async for chunk in _generate_content_stream(client):
230238
yield [
231239
self._create_streaming_chat_message_content(chunk, candidate, function_invoke_attempt)

‎python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_text_completion.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -137,10 +137,14 @@ async def _generate_content(client: Client) -> GenerateContentResponse:
137137
vertexai=True,
138138
project=self.service_settings.cloud_project_id,
139139
location=self.service_settings.cloud_region,
140+
http_options=self._get_http_options(),
140141
) as client:
141142
response: GenerateContentResponse = await _generate_content(client) # type: ignore[no-redef]
142143
else:
143-
with Client(api_key=self.service_settings.api_key.get_secret_value()) as client: # type: ignore[union-attr]
144+
with Client(
145+
api_key=self.service_settings.api_key.get_secret_value(), # type: ignore[union-attr]
146+
http_options=self._get_http_options(),
147+
) as client:
144148
response: GenerateContentResponse = await _generate_content(client) # type: ignore[no-redef]
145149

146150
return [self._create_text_content(response, candidate) for candidate in response.candidates] # type: ignore
@@ -175,11 +179,15 @@ async def _generate_content_stream(client: Client) -> AsyncGenerator[GenerateCon
175179
vertexai=True,
176180
project=self.service_settings.cloud_project_id,
177181
location=self.service_settings.cloud_region,
182+
http_options=self._get_http_options(),
178183
) as client:
179184
async for chunk in _generate_content_stream(client):
180185
yield [self._create_streaming_text_content(chunk, candidate) for candidate in chunk.candidates] # type: ignore
181186
else:
182-
with Client(api_key=self.service_settings.api_key.get_secret_value()) as client: # type: ignore[union-attr]
187+
with Client(
188+
api_key=self.service_settings.api_key.get_secret_value(), # type: ignore[union-attr]
189+
http_options=self._get_http_options(),
190+
) as client:
183191
async for chunk in _generate_content_stream(client):
184192
yield [self._create_streaming_text_content(chunk, candidate) for candidate in chunk.candidates] # type: ignore
185193

‎python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_text_embedding.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,10 +133,14 @@ async def _embed_content(client: Client) -> EmbedContentResponse:
133133
vertexai=True,
134134
project=self.service_settings.cloud_project_id,
135135
location=self.service_settings.cloud_region,
136+
http_options=self._get_http_options(),
136137
) as client:
137138
response: EmbedContentResponse = await _embed_content(client) # type: ignore[no-redef]
138139
else:
139-
with Client(api_key=self.service_settings.api_key.get_secret_value()) as client: # type: ignore[union-attr]
140+
with Client(
141+
api_key=self.service_settings.api_key.get_secret_value(), # type: ignore[union-attr]
142+
http_options=self._get_http_options(),
143+
) as client:
140144
response: EmbedContentResponse = await _embed_content(client) # type: ignore[no-redef]
141145

142146
return [embedding.values for embedding in response.embeddings] # type: ignore

‎python/tests/unit/connectors/ai/google/google_ai/services/test_google_ai_user_agent.py‎

Lines changed: 172 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,172 @@
1+
# Copyright (c) Microsoft. All rights reserved.
2+
3+
from unittest.mock import AsyncMock, patch
4+
5+
from semantic_kernel.connectors.ai.google.google_ai.google_ai_prompt_execution_settings import (
6+
GoogleAIChatPromptExecutionSettings,
7+
GoogleAIEmbeddingPromptExecutionSettings,
8+
GoogleAITextPromptExecutionSettings,
9+
)
10+
from semantic_kernel.connectors.ai.google.google_ai.services.google_ai_chat_completion import GoogleAIChatCompletion
11+
from semantic_kernel.connectors.ai.google.google_ai.services.google_ai_text_completion import GoogleAITextCompletion
12+
from semantic_kernel.connectors.ai.google.google_ai.services.google_ai_text_embedding import GoogleAITextEmbedding
13+
from semantic_kernel.const import USER_AGENT
14+
from semantic_kernel.contents.chat_history import ChatHistory
15+
16+
17+
async def test_google_ai_chat_completion_user_agent(google_ai_unit_test_env):
18+
"""Test that GoogleAIChatCompletion sends the User-Agent header."""
19+
chat_history = ChatHistory()
20+
chat_history.add_user_message("hi")
21+
22+
with patch(
23+
"semantic_kernel.connectors.ai.google.google_ai.services.google_ai_chat_completion.Client"
24+
) as mock_client:
25+
mock_instance = mock_client.return_value.__enter__.return_value
26+
mock_instance.aio.models.generate_content = AsyncMock()
27+
28+
service = GoogleAIChatCompletion(gemini_model_id="gemini-3-flash-preview", api_key="test-api-key")
29+
30+
await service.get_chat_message_contents(
31+
chat_history=chat_history, settings=GoogleAIChatPromptExecutionSettings()
32+
)
33+
34+
_, kwargs = mock_client.call_args
35+
assert "http_options" in kwargs
36+
assert kwargs["http_options"] is not None
37+
assert "headers" in kwargs["http_options"]
38+
assert USER_AGENT in kwargs["http_options"]["headers"]
39+
assert "semantic-kernel-python" in kwargs["http_options"]["headers"][USER_AGENT]
40+
41+
42+
async def test_google_ai_chat_completion_vertex_ai_user_agent(google_ai_unit_test_env):
43+
"""Test that GoogleAIChatCompletion sends the User-Agent header with Vertex AI."""
44+
chat_history = ChatHistory()
45+
chat_history.add_user_message("hi")
46+
47+
with patch(
48+
"semantic_kernel.connectors.ai.google.google_ai.services.google_ai_chat_completion.Client"
49+
) as mock_client:
50+
mock_instance = mock_client.return_value.__enter__.return_value
51+
mock_instance.aio.models.generate_content = AsyncMock()
52+
53+
service = GoogleAIChatCompletion(
54+
gemini_model_id="gemini-3-flash-preview",
55+
use_vertexai=True,
56+
project_id="test-project",
57+
region="test-region",
58+
)
59+
60+
await service.get_chat_message_contents(
61+
chat_history=chat_history, settings=GoogleAIChatPromptExecutionSettings()
62+
)
63+
64+
_, kwargs = mock_client.call_args
65+
assert "http_options" in kwargs
66+
assert kwargs["http_options"] is not None
67+
assert "headers" in kwargs["http_options"]
68+
assert USER_AGENT in kwargs["http_options"]["headers"]
69+
assert "semantic-kernel-python" in kwargs["http_options"]["headers"][USER_AGENT]
70+
71+
72+
async def test_google_ai_chat_completion_no_telemetry(google_ai_unit_test_env):
73+
"""Test that GoogleAIChatCompletion does not send the User-Agent header when telemetry is disabled."""
74+
chat_history = ChatHistory()
75+
chat_history.add_user_message("hi")
76+
77+
with (
78+
patch("semantic_kernel.connectors.ai.google.google_ai.services.google_ai_base.IS_TELEMETRY_ENABLED", False),
79+
patch(
80+
"semantic_kernel.connectors.ai.google.google_ai.services.google_ai_chat_completion.Client"
81+
) as mock_client,
82+
):
83+
mock_instance = mock_client.return_value.__enter__.return_value
84+
mock_instance.aio.models.generate_content = AsyncMock()
85+
86+
service = GoogleAIChatCompletion(gemini_model_id="gemini-3-flash-preview", api_key="test-api-key")
87+
88+
await service.get_chat_message_contents(
89+
chat_history=chat_history, settings=GoogleAIChatPromptExecutionSettings()
90+
)
91+
92+
_, kwargs = mock_client.call_args
93+
assert "http_options" in kwargs
94+
assert kwargs["http_options"] is None
95+
96+
97+
async def test_google_ai_text_completion_user_agent(google_ai_unit_test_env):
98+
"""Test that GoogleAITextCompletion sends the User-Agent header."""
99+
with patch(
100+
"semantic_kernel.connectors.ai.google.google_ai.services.google_ai_text_completion.Client"
101+
) as mock_client:
102+
mock_instance = mock_client.return_value.__enter__.return_value
103+
mock_instance.aio.models.generate_content = AsyncMock()
104+
105+
service = GoogleAITextCompletion(gemini_model_id="gemini-3-flash-preview", api_key="test-api-key")
106+
107+
await service.get_text_contents(prompt="hi", settings=GoogleAITextPromptExecutionSettings())
108+
109+
_, kwargs = mock_client.call_args
110+
assert "http_options" in kwargs
111+
assert kwargs["http_options"] is not None
112+
assert "headers" in kwargs["http_options"]
113+
assert USER_AGENT in kwargs["http_options"]["headers"]
114+
assert "semantic-kernel-python" in kwargs["http_options"]["headers"][USER_AGENT]
115+
116+
117+
async def test_google_ai_text_completion_no_telemetry(google_ai_unit_test_env):
118+
"""Test that GoogleAITextCompletion does not send the User-Agent header when telemetry is disabled."""
119+
with (
120+
patch("semantic_kernel.connectors.ai.google.google_ai.services.google_ai_base.IS_TELEMETRY_ENABLED", False),
121+
patch(
122+
"semantic_kernel.connectors.ai.google.google_ai.services.google_ai_text_completion.Client"
123+
) as mock_client,
124+
):
125+
mock_instance = mock_client.return_value.__enter__.return_value
126+
mock_instance.aio.models.generate_content = AsyncMock()
127+
128+
service = GoogleAITextCompletion(gemini_model_id="gemini-3-flash-preview", api_key="test-api-key")
129+
130+
await service.get_text_contents(prompt="hi", settings=GoogleAITextPromptExecutionSettings())
131+
132+
_, kwargs = mock_client.call_args
133+
assert "http_options" in kwargs
134+
assert kwargs["http_options"] is None
135+
136+
137+
async def test_google_ai_text_embedding_user_agent(google_ai_unit_test_env):
138+
"""Test that GoogleAITextEmbedding sends the User-Agent header."""
139+
with patch(
140+
"semantic_kernel.connectors.ai.google.google_ai.services.google_ai_text_embedding.Client"
141+
) as mock_client:
142+
mock_instance = mock_client.return_value.__enter__.return_value
143+
mock_instance.aio.models.embed_content = AsyncMock()
144+
145+
service = GoogleAITextEmbedding(embedding_model_id="gemini-embedding-2-preview", api_key="test-api-key")
146+
147+
await service.generate_embeddings(texts=["hi"], settings=GoogleAIEmbeddingPromptExecutionSettings())
148+
149+
_, kwargs = mock_client.call_args
150+
assert "http_options" in kwargs
151+
assert kwargs["http_options"] is not None
152+
assert "headers" in kwargs["http_options"]
153+
assert USER_AGENT in kwargs["http_options"]["headers"]
154+
assert "semantic-kernel-python" in kwargs["http_options"]["headers"][USER_AGENT]
155+
156+
157+
async def test_google_ai_text_embedding_no_telemetry(google_ai_unit_test_env):
158+
"""Test that GoogleAITextEmbedding does not send the User-Agent header when telemetry is disabled."""
159+
with (
160+
patch("semantic_kernel.connectors.ai.google.google_ai.services.google_ai_base.IS_TELEMETRY_ENABLED", False),
161+
patch("semantic_kernel.connectors.ai.google.google_ai.services.google_ai_text_embedding.Client") as mock_client,
162+
):
163+
mock_instance = mock_client.return_value.__enter__.return_value
164+
mock_instance.aio.models.embed_content = AsyncMock()
165+
166+
service = GoogleAITextEmbedding(embedding_model_id="gemini-embedding-2-preview", api_key="test-api-key")
167+
168+
await service.generate_embeddings(texts=["hi"], settings=GoogleAIEmbeddingPromptExecutionSettings())
169+
170+
_, kwargs = mock_client.call_args
171+
assert "http_options" in kwargs
172+
assert kwargs["http_options"] is None

0 commit comments

Comments
 (0)