chore: Adding Eval Client label to model calls made during evals

Co-authored-by: Ankur Sharma <ankusharma@google.com>
PiperOrigin-RevId: 838857867

tests/unittests/models/test_google_llm.py

@@ -22,8 +22,6 @@ from google.adk import version as adk_version
 from google.adk.agents.context_cache_config import ContextCacheConfig
 from google.adk.models.cache_metadata import CacheMetadata
 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 _build_function_declaration_log
 from google.adk.models.google_llm import _build_request_log
 from google.adk.models.google_llm import _RESOURCE_EXHAUSTED_POSSIBLE_FIX_MESSAGE
@@ -31,6 +29,8 @@ from google.adk.models.google_llm import _ResourceExhaustedError
 from google.adk.models.google_llm import Gemini
 from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
+from google.adk.utils._client_labels_utils import _AGENT_ENGINE_TELEMETRY_ENV_VARIABLE_NAME
+from google.adk.utils._client_labels_utils import _AGENT_ENGINE_TELEMETRY_TAG
 from google.adk.utils.variant_utils import GoogleLLMVariant
 from google.genai import types
 from google.genai.errors import ClientError
@@ -142,13 +142,6 @@ def llm_request_with_computer_use():
   )
 
 
-@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) == 4
@@ -193,12 +186,15 @@ 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"
+def test_client_version_header_with_agent_engine(monkeypatch):
+  monkeypatch.setenv(
+      _AGENT_ENGINE_TELEMETRY_ENV_VARIABLE_NAME, "my_test_project"
+  )
   model = Gemini(model="gemini-1.5-flash")
   client = model.api_client
 
-  # Check that ADK version with telemetry tag and Python version are present in headers
+  # Check that ADK version with telemetry tag and Python version are present in
+  # headers
   adk_version_with_telemetry = (
       f"google-adk/{adk_version.__version__}+{_AGENT_ENGINE_TELEMETRY_TAG}"
   )
@@ -473,8 +469,9 @@ async def test_generate_content_async_with_custom_headers(
   """Test that tracking headers are updated when custom headers are provided."""
   # Add custom headers to the request config
   custom_headers = {"custom-header": "custom-value"}
-  for key in gemini_llm._tracking_headers:
-    custom_headers[key] = "custom " + gemini_llm._tracking_headers[key]
+  tracking_headers = gemini_llm._tracking_headers()
+  for key in tracking_headers:
+    custom_headers[key] = "custom " + tracking_headers[key]
   llm_request.config.http_options = types.HttpOptions(headers=custom_headers)
 
   with mock.patch.object(gemini_llm, "api_client") as mock_client:
@@ -497,8 +494,9 @@ async def test_generate_content_async_with_custom_headers(
     config_arg = call_args.kwargs["config"]
 
     for key, value in config_arg.http_options.headers.items():
-      if key in gemini_llm._tracking_headers:
-        assert value == gemini_llm._tracking_headers[key] + " custom"
+      tracking_headers = gemini_llm._tracking_headers()
+      if key in tracking_headers:
+        assert value == tracking_headers[key] + " custom"
       else:
         assert value == custom_headers[key]
 
@@ -547,7 +545,7 @@ async def test_generate_content_async_stream_with_custom_headers(
     config_arg = call_args.kwargs["config"]
 
     expected_headers = custom_headers.copy()
-    expected_headers.update(gemini_llm._tracking_headers)
+    expected_headers.update(gemini_llm._tracking_headers())
     assert config_arg.http_options.headers == expected_headers
 
     assert len(responses) == 2
@@ -601,7 +599,7 @@ async def test_generate_content_async_patches_tracking_headers(
     assert final_config.http_options is not None
     assert (
         final_config.http_options.headers["x-goog-api-client"]
-        == gemini_llm._tracking_headers["x-goog-api-client"]
+        == gemini_llm._tracking_headers()["x-goog-api-client"]
     )
 
     assert len(responses) == 2 if stream else 1
@@ -635,7 +633,7 @@ def test_live_api_client_properties(gemini_llm):
     assert http_options.api_version == "v1beta1"
 
     # Check that tracking headers are included
-    tracking_headers = gemini_llm._tracking_headers
+    tracking_headers = gemini_llm._tracking_headers()
     for key, value in tracking_headers.items():
       assert key in http_options.headers
       assert value in http_options.headers[key]
@@ -673,7 +671,7 @@ async def test_connect_with_custom_headers(gemini_llm, llm_request):
 
       # Verify that tracking headers were merged with custom headers
       expected_headers = custom_headers.copy()
-      expected_headers.update(gemini_llm._tracking_headers)
+      expected_headers.update(gemini_llm._tracking_headers())
       assert config_arg.http_options.headers == expected_headers
 
       # Verify that API version was set