Skip to content

Commit d79b22c

Browse files
committed
feat(apigee): add userinfo.email scope for tokeninfo user identification
ApigeeLlm now explicitly requests the userinfo.email OAuth scope alongside cloud-platform when creating credentials. This enables Apigee Gateway to identify callers via Google's tokeninfo API when using Service Account key authentication. Fixes #4721
1 parent 31b005c commit d79b22c

2 files changed

Lines changed: 47 additions & 0 deletions

File tree

src/google/adk/models/apigee_llm.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
from typing import Optional
3030
from typing import TYPE_CHECKING
3131

32+
import google.auth
3233
from google.adk import version as adk_version
3334
from google.genai import types
3435
import httpx
@@ -52,6 +53,11 @@
5253
_PROJECT_ENV_VARIABLE_NAME = 'GOOGLE_CLOUD_PROJECT'
5354
_LOCATION_ENV_VARIABLE_NAME = 'GOOGLE_CLOUD_LOCATION'
5455

56+
_APIGEE_SCOPES = [
57+
'https://www.googleapis.com/auth/cloud-platform',
58+
'https://www.googleapis.com/auth/userinfo.email',
59+
]
60+
5561
_CUSTOM_METADATA_FIELDS = (
5662
'id',
5763
'created',
@@ -232,13 +238,16 @@ def api_client(self) -> Client:
232238
**kwargs_for_http_options,
233239
)
234240

241+
credentials, _ = google.auth.default(scopes=_APIGEE_SCOPES)
242+
235243
kwargs_for_client = {}
236244
kwargs_for_client['vertexai'] = self._isvertexai
237245
if self._isvertexai:
238246
kwargs_for_client['project'] = self._project
239247
kwargs_for_client['location'] = self._location
240248

241249
return Client(
250+
credentials=credentials,
242251
http_options=http_options,
243252
**kwargs_for_client,
244253
)

tests/unittests/models/test_apigee_llm.py

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
from unittest import mock
1919
from unittest.mock import AsyncMock
2020

21+
from google.adk.models.apigee_llm import _APIGEE_SCOPES
2122
from google.adk.models.apigee_llm import ApigeeLlm
2223
from google.adk.models.apigee_llm import CompletionsHTTPClient
2324
from google.adk.models.llm_request import LlmRequest
@@ -649,3 +650,40 @@ def test_parse_response_usage_metadata():
649650
assert llm_response.usage_metadata.candidates_token_count == 5
650651
assert llm_response.usage_metadata.total_token_count == 15
651652
assert llm_response.usage_metadata.thoughts_token_count == 4
653+
654+
655+
@pytest.mark.asyncio
656+
@mock.patch('google.genai.Client')
657+
@mock.patch('google.adk.models.apigee_llm.google.auth.default')
658+
async def test_api_client_requests_userinfo_email_scope(
659+
mock_auth_default, mock_client_constructor, llm_request
660+
):
661+
"""Tests that api_client requests userinfo.email scope for Apigee Gateway tokeninfo."""
662+
mock_credentials = mock.Mock()
663+
mock_auth_default.return_value = (mock_credentials, 'test-project')
664+
665+
mock_client_instance = mock.Mock()
666+
mock_client_instance.aio.models.generate_content = AsyncMock(
667+
return_value=types.GenerateContentResponse(
668+
candidates=[
669+
types.Candidate(
670+
content=Content(
671+
parts=[Part.from_text(text='Test response')],
672+
role='model',
673+
)
674+
)
675+
]
676+
)
677+
)
678+
mock_client_constructor.return_value = mock_client_instance
679+
680+
apigee_llm = ApigeeLlm(
681+
model=APIGEE_GEMINI_MODEL_ID,
682+
proxy_url=PROXY_URL,
683+
)
684+
_ = [resp async for resp in apigee_llm.generate_content_async(llm_request)]
685+
686+
mock_auth_default.assert_called_once_with(scopes=_APIGEE_SCOPES)
687+
688+
_, kwargs = mock_client_constructor.call_args
689+
assert kwargs['credentials'] is mock_credentials

0 commit comments

Comments
 (0)