NeoAnd/adk-python
1
0
Fork 0
mirror of https://github.com/EvolutionAPI/adk-python.git synced 2025-07-13 15:14:50 -06:00
Code Issues Actions Packages Projects Releases Wiki Activity

chore: Add agent engine telemetry

Browse Source 
PiperOrigin-RevId: 765346182
This commit is contained in:
Shawn Yang 2025-05-30 14:49:14 -07:00 committed by Copybara-Service
parent 45ef668435
commit 3930a4b989
2 changed files with 39 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/google/adk/models/google_llm.py
View File

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

34
tests/unittests/models/test_google_llm.py
View File

@ -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)