diff --git a/src/google/adk/models/google_llm.py b/src/google/adk/models/google_llm.py index 42dc964..7a3932a 100644 --- a/src/google/adk/models/google_llm.py +++ b/src/google/adk/models/google_llm.py @@ -18,6 +18,7 @@ from __future__ import annotations import contextlib from functools import cached_property import logging +import os import sys from typing import AsyncGenerator from typing import cast @@ -40,6 +41,8 @@ logger = logging.getLogger('google_adk.' + __name__) _NEW_LINE = '\n' _EXCLUDED_PART_FIELD = {'inline_data': {'data'}} +_AGENT_ENGINE_TELEMETRY_TAG = 'remote_reasoning_engine' +_AGENT_ENGINE_TELEMETRY_ENV_VARIABLE_NAME = 'GOOGLE_CLOUD_AGENT_ENGINE_ID' class Gemini(BaseLlm): @@ -181,6 +184,8 @@ class Gemini(BaseLlm): @cached_property def _tracking_headers(self) -> dict[str, str]: framework_label = f'google-adk/{version.__version__}' + if os.environ.get(_AGENT_ENGINE_TELEMETRY_ENV_VARIABLE_NAME): + framework_label = f'{framework_label}+{_AGENT_ENGINE_TELEMETRY_TAG}' language_label = 'gl-python/' + sys.version.split()[0] version_header_value = f'{framework_label} {language_label}' tracking_headers = { diff --git a/tests/unittests/models/test_google_llm.py b/tests/unittests/models/test_google_llm.py index e1deadd..349dd13 100644 --- a/tests/unittests/models/test_google_llm.py +++ b/tests/unittests/models/test_google_llm.py @@ -12,11 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import sys from unittest import mock from google.adk import version as adk_version from google.adk.models.gemini_llm_connection import GeminiLlmConnection +from google.adk.models.google_llm import _AGENT_ENGINE_TELEMETRY_ENV_VARIABLE_NAME +from google.adk.models.google_llm import _AGENT_ENGINE_TELEMETRY_TAG from google.adk.models.google_llm import Gemini from google.adk.models.llm_request import LlmRequest from google.adk.models.llm_response import LlmResponse @@ -60,6 +63,13 @@ def llm_request(): ) +@pytest.fixture +def mock_os_environ(): + initial_env = os.environ.copy() + with mock.patch.dict(os.environ, initial_env, clear=False) as m: + yield m + + def test_supported_models(): models = Gemini.supported_models() assert len(models) == 3 @@ -91,6 +101,30 @@ def test_client_version_header(): ) +def test_client_version_header_with_agent_engine(mock_os_environ): + os.environ[_AGENT_ENGINE_TELEMETRY_ENV_VARIABLE_NAME] = "my_test_project" + model = Gemini(model="gemini-1.5-flash") + client = model.api_client + adk_header_base = f"google-adk/{adk_version.__version__}" + adk_header_with_telemetry = ( + f"{adk_header_base}+{_AGENT_ENGINE_TELEMETRY_TAG}" + f" gl-python/{sys.version.split()[0]}" + ) + genai_header = ( + f"google-genai-sdk/{genai_version.__version__} " + f"gl-python/{sys.version.split()[0]} " + ) + expected_header = genai_header + adk_header_with_telemetry + + assert ( + expected_header + in client._api_client._http_options.headers["x-goog-api-client"] + ) + assert ( + expected_header in client._api_client._http_options.headers["user-agent"] + ) + + def test_maybe_append_user_content(gemini_llm, llm_request): # Test with user content already present gemini_llm._maybe_append_user_content(llm_request)