NeoAnd/adk-python
1
0
Fork 0
mirror of https://github.com/EvolutionAPI/adk-python.git synced 2025-12-18 19:32:21 -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
Download Patch File Download Diff File Expand all files Collapse all files

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)