feat: Mark Vertex calls made from non-gemini models

PiperOrigin-RevId: 848140103

tests/unittests/models/test_anthropic_llm.py

@@ -391,6 +391,31 @@ async def test_anthropic_llm_generate_content_async(
       assert responses[0].content.parts[0].text == "Hello, how can I help you?"
 
 
+def test_claude_vertex_client_uses_tracking_headers():
+  """Tests that Claude vertex client is called with tracking headers."""
+  with mock.patch.object(
+      anthropic_llm, "AsyncAnthropicVertex", autospec=True
+  ) as mock_anthropic_vertex:
+    with mock.patch.dict(
+        os.environ,
+        {
+            "GOOGLE_CLOUD_PROJECT": "test-project",
+            "GOOGLE_CLOUD_LOCATION": "us-central1",
+        },
+    ):
+      instance = Claude(model="claude-3-5-sonnet-v2@20241022")
+      _ = instance._anthropic_client
+      mock_anthropic_vertex.assert_called_once()
+      _, kwargs = mock_anthropic_vertex.call_args
+      assert "default_headers" in kwargs
+      assert "x-goog-api-client" in kwargs["default_headers"]
+      assert "user-agent" in kwargs["default_headers"]
+      assert (
+          f"google-adk/{adk_version.__version__}"
+          in kwargs["default_headers"]["user-agent"]
+      )
+
+
 @pytest.mark.asyncio
 async def test_generate_content_async_with_max_tokens(
     llm_request, generate_content_response, generate_llm_response

tests/unittests/models/test_google_llm.py

@@ -31,6 +31,7 @@ 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._google_client_headers import get_tracking_headers
 from google.adk.utils.variant_utils import GoogleLLMVariant
 from google.genai import types
 from google.genai.errors import ClientError
@@ -469,7 +470,7 @@ 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"}
-  tracking_headers = gemini_llm._tracking_headers()
+  tracking_headers = get_tracking_headers()
   for key in tracking_headers:
     custom_headers[key] = "custom " + tracking_headers[key]
   llm_request.config.http_options = types.HttpOptions(headers=custom_headers)
@@ -494,7 +495,7 @@ 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():
-      tracking_headers = gemini_llm._tracking_headers()
+      tracking_headers = get_tracking_headers()
       if key in tracking_headers:
         assert value == tracking_headers[key] + " custom"
       else:
@@ -545,7 +546,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(get_tracking_headers())
     assert config_arg.http_options.headers == expected_headers
 
     assert len(responses) == 2
@@ -599,7 +600,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"]
+        == get_tracking_headers()["x-goog-api-client"]
     )
 
     assert len(responses) == 2 if stream else 1
@@ -633,7 +634,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 = get_tracking_headers()
     for key, value in tracking_headers.items():
       assert key in http_options.headers
       assert value in http_options.headers[key]
@@ -671,7 +672,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(get_tracking_headers())
       assert config_arg.http_options.headers == expected_headers
 
       # Verify that API version was set

tests/unittests/models/test_litellm.py

@@ -2447,11 +2447,12 @@ def test_model_response_to_chunk(
 async def test_acompletion_additional_args(mock_acompletion, mock_client):
   lite_llm_instance = LiteLlm(
       # valid args
-      model="test_model",
+      model="vertex_ai/test_model",
       llm_client=mock_client,
       api_key="test_key",
       api_base="some://url",
       api_version="2024-09-12",
+      headers={"custom": "header"},  # Add custom header to test merge
       # invalid args (ignored)
       stream=True,
       messages=[{"role": "invalid", "content": "invalid"}],
@@ -2478,13 +2479,43 @@ async def test_acompletion_additional_args(mock_acompletion, mock_client):
 
   _, kwargs = mock_acompletion.call_args
 
-  assert kwargs["model"] == "test_model"
+  assert kwargs["model"] == "vertex_ai/test_model"
   assert kwargs["messages"][0]["role"] == "user"
   assert kwargs["messages"][0]["content"] == "Test prompt"
   assert kwargs["tools"][0]["function"]["name"] == "test_function"
   assert "stream" not in kwargs
   assert "llm_client" not in kwargs
   assert kwargs["api_base"] == "some://url"
+  assert "headers" in kwargs
+  assert kwargs["headers"]["custom"] == "header"
+  assert "x-goog-api-client" in kwargs["headers"]
+  assert "user-agent" in kwargs["headers"]
+
+
+@pytest.mark.asyncio
+async def test_acompletion_additional_args_non_vertex(
+    mock_acompletion, mock_client
+):
+  """Test that tracking headers are not added for non-Vertex AI models."""
+  lite_llm_instance = LiteLlm(
+      model="openai/gpt-4o",
+      llm_client=mock_client,
+      api_key="test_key",
+      headers={"custom": "header"},
+  )
+
+  async for _ in lite_llm_instance.generate_content_async(
+      LLM_REQUEST_WITH_FUNCTION_DECLARATION
+  ):
+    pass
+
+  mock_acompletion.assert_called_once()
+  _, kwargs = mock_acompletion.call_args
+  assert kwargs["model"] == "openai/gpt-4o"
+  assert "headers" in kwargs
+  assert kwargs["headers"]["custom"] == "header"
+  assert "x-goog-api-client" not in kwargs["headers"]
+  assert "user-agent" not in kwargs["headers"]
 
 
 @pytest.mark.asyncio