chore: Add agent engine telemetry

PiperOrigin-RevId: 765346182

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 = {

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)