Merge remote-tracking branch 'upstream/main'

Author: discopops

pyproject.toml

@@ -45,6 +45,7 @@ dependencies = [
   "google-cloud-storage>=2.18.0, <4.0.0",                    # For GCS Artifact service
   "google-genai>=1.56.0, <2.0.0",                            # Google GenAI SDK
   "graphviz>=0.20.2, <1.0.0",                                # Graphviz for graph rendering
+  "httpx>=0.27.0, <1.0.0",                                   # HTTP client library
   "jsonschema>=4.23.0, <5.0.0",                              # Agent Builder config validation
   "mcp>=1.23.0, <2.0.0",                                     # For MCP Toolset
   "opentelemetry-api>=1.37.0, <=1.37.0",                     # OpenTelemetry - limit upper version for sdk and api to not risk breaking changes from unstable _logs package.
@@ -163,7 +164,7 @@ extensions = [
   "toolbox-adk>=0.5.7, <0.6.0",                 # For tools.toolbox_toolset.ToolboxToolset
 ]
 
-otel-gcp = ["opentelemetry-instrumentation-google-genai>=0.3b0, <1.0.0"]
+otel-gcp = ["opentelemetry-instrumentation-google-genai>=0.6b0, <1.0.0"]
 
 toolbox = ["toolbox-adk>=0.5.7, <0.6.0"]
 

src/google/adk/models/anthropic_llm.py

@@ -36,6 +36,7 @@
 from pydantic import BaseModel
 from typing_extensions import override
 
+from ..utils._google_client_headers import get_tracking_headers
 from .base_llm import BaseLlm
 from .llm_response import LlmResponse
 
@@ -345,4 +346,5 @@ def _anthropic_client(self) -> AsyncAnthropicVertex:
     return AsyncAnthropicVertex(
         project_id=os.environ["GOOGLE_CLOUD_PROJECT"],
         region=os.environ["GOOGLE_CLOUD_LOCATION"],
+        default_headers=get_tracking_headers(),
     )

src/google/adk/models/google_llm.py

@@ -30,7 +30,8 @@
 from google.genai.errors import ClientError
 from typing_extensions import override
 
-from ..utils._client_labels_utils import get_client_labels
+from ..utils._google_client_headers import get_tracking_headers
+from ..utils._google_client_headers import merge_tracking_headers
 from ..utils.context_utils import Aclosing
 from ..utils.streaming_utils import StreamingResponseAggregator
 from ..utils.variant_utils import GoogleLLMVariant
@@ -316,13 +317,7 @@ def _api_backend(self) -> GoogleLLMVariant:
     )
 
   def _tracking_headers(self) -> dict[str, str]:
-    labels = get_client_labels()
-    header_value = ' '.join(labels)
-    tracking_headers = {
-        'x-goog-api-client': header_value,
-        'user-agent': header_value,
-    }
-    return tracking_headers
+    return get_tracking_headers()
 
   @cached_property
   def _live_api_version(self) -> str:
@@ -362,8 +357,10 @@ async def connect(self, llm_request: LlmRequest) -> BaseLlmConnection:
     ):
       if not llm_request.live_connect_config.http_options.headers:
         llm_request.live_connect_config.http_options.headers = {}
-      llm_request.live_connect_config.http_options.headers.update(
-          self._tracking_headers()
+      llm_request.live_connect_config.http_options.headers = (
+          self._merge_tracking_headers(
+              llm_request.live_connect_config.http_options.headers
+          )
       )
       llm_request.live_connect_config.http_options.api_version = (
           self._live_api_version
@@ -456,20 +453,7 @@ async def _preprocess_request(self, llm_request: LlmRequest) -> None:
 
   def _merge_tracking_headers(self, headers: dict[str, str]) -> dict[str, str]:
     """Merge tracking headers to the given headers."""
-    headers = headers or {}
-    for key, tracking_header_value in self._tracking_headers().items():
-      custom_value = headers.get(key, None)
-      if not custom_value:
-        headers[key] = tracking_header_value
-        continue
-
-      # Merge tracking headers with existing headers and avoid duplicates.
-      value_parts = tracking_header_value.split(' ')
-      for custom_value_part in custom_value.split(' '):
-        if custom_value_part not in value_parts:
-          value_parts.append(custom_value_part)
-      headers[key] = ' '.join(value_parts)
-    return headers
+    return merge_tracking_headers(headers)
 
 
 def _build_function_declaration_log(

src/google/adk/models/lite_llm.py

@@ -51,6 +51,7 @@
 from pydantic import Field
 from typing_extensions import override
 
+from ..utils._google_client_headers import merge_tracking_headers
 from .base_llm import BaseLlm
 from .llm_request import LlmRequest
 from .llm_response import LlmResponse
@@ -1699,6 +1700,18 @@ def _build_request_log(req: LlmRequest) -> str:
 """
 
 
+def _is_litellm_vertex_model(model_string: str) -> bool:
+  """Check if the model is a Vertex AI model accessed via LiteLLM.
+
+  Args:
+    model_string: A LiteLLM model string (e.g., "vertex_ai/gemini-2.5-flash")
+
+  Returns:
+    True if it's a Vertex AI model accessed via LiteLLM, False otherwise
+  """
+  return model_string.startswith("vertex_ai/")
+
+
 def _is_litellm_gemini_model(model_string: str) -> bool:
   """Check if the model is a Gemini model accessed via LiteLLM.
 
@@ -1867,6 +1880,14 @@ async def generate_content_async(
     }
     completion_args.update(self._additional_args)
 
+    # merge headers
+    if _is_litellm_vertex_model(effective_model) or _is_litellm_gemini_model(
+        effective_model
+    ):
+      completion_args["headers"] = merge_tracking_headers(
+          completion_args.get("headers")
+      )
+
     if generation_params:
       completion_args.update(generation_params)
 

src/google/adk/telemetry/tracing.py

@@ -35,6 +35,7 @@
 from google.genai import types
 from google.genai.models import Models
 from opentelemetry import _logs
+from opentelemetry import context as otel_context
 from opentelemetry import trace
 from opentelemetry._logs import LogRecord
 from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_AGENT_DESCRIPTION
@@ -444,7 +445,8 @@ def use_generate_content_span(
       _is_gemini_agent(invocation_context.agent)
       and _instrumented_with_opentelemetry_instrumentation_google_genai()
   ):
-    yield None
+    with _use_extra_generate_content_attributes(common_attributes):
+      yield
   else:
     with _use_native_generate_content_span(
         llm_request=llm_request,
@@ -492,6 +494,33 @@ def _instrumented_with_opentelemetry_instrumentation_google_genai() -> bool:
   return False
 
 
+@contextmanager
+def _use_extra_generate_content_attributes(
+    extra_attributes: Mapping[str, AttributeValue],
+):
+  try:
+    from opentelemetry.instrumentation.google_genai import GENERATE_CONTENT_EXTRA_ATTRIBUTES_CONTEXT_KEY
+  except (ImportError, AttributeError):
+    logger.warning(
+        'opentelemetry-instrumentor-google-genai is installed but has'
+        ' insufficient version,'
+        + ' so some tracing dependent features may not work properly.'
+        + ' Please upgrade to version to 0.6b0 or above.'
+    )
+    yield
+    return
+
+  tok = otel_context.attach(
+      otel_context.set_value(
+          GENERATE_CONTENT_EXTRA_ATTRIBUTES_CONTEXT_KEY, extra_attributes
+      )
+  )
+  try:
+    yield
+  finally:
+    otel_context.detach(tok)
+
+
 def _is_gemini_agent(agent: BaseAgent) -> bool:
   from ..agents.llm_agent import LlmAgent
 

src/google/adk/tools/openapi_tool/auth/auth_helpers.py

@@ -28,9 +28,9 @@
 from fastapi.openapi.models import OAuth2
 from fastapi.openapi.models import OpenIdConnect
 from fastapi.openapi.models import Schema
+import httpx
 from pydantic import BaseModel
 from pydantic import ValidationError
-import requests
 
 from ....auth.auth_credential import AuthCredential
 from ....auth.auth_credential import AuthCredentialTypes
@@ -289,14 +289,14 @@ def openid_url_to_scheme_credential(
   Raises:
       ValueError: If the OpenID URL is invalid, fetching fails, or required
         fields are missing.
-      requests.exceptions.RequestException:  If there's an error during the
+      httpx.HTTPStatusError or httpx.RequestError: If there's an error during the
           HTTP request.
   """
   try:
-    response = requests.get(openid_url, timeout=10)
+    response = httpx.get(openid_url, timeout=10)
     response.raise_for_status()
     config_dict = response.json()
-  except requests.exceptions.RequestException as e:
+  except httpx.RequestError as e:
     raise ValueError(
         f"Failed to fetch OpenID configuration from {openid_url}: {e}"
     ) from e

src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py

@@ -28,7 +28,7 @@
 from fastapi.openapi.models import Operation
 from fastapi.openapi.models import Schema
 from google.genai.types import FunctionDeclaration
-import requests
+import httpx
 from typing_extensions import override
 
 from ....agents.readonly_context import ReadonlyContext
@@ -312,7 +312,7 @@ def _prepare_request_params(
 
     Returns:
         A dictionary containing the  request parameters for the API call. This
-        initializes a requests.request() call.
+        initializes an httpx.AsyncClient.request() call.
 
     Example:
         self._prepare_request_params({"input_id": "test-id"})
@@ -497,17 +497,7 @@ async def call(
       if provider_headers:
         request_params.setdefault("headers", {}).update(provider_headers)
 
-    # Log the API request
-    self._logger.debug(
-        "API Request: %s %s",
-        request_params.get("method", "").upper(),
-        request_params.get("url", ""),
-    )
-    self._logger.debug("API Request params: %s", request_params.get("params"))
-    if "json" in request_params:
-      self._logger.debug("API Request body: %s", request_params.get("json"))
-
-    response = requests.request(**request_params)
+    response = await _request(**request_params)
 
     # Log the API response
     self._logger.debug(
@@ -519,11 +509,9 @@ async def call(
 
     # Parse API response
     try:
-      response.raise_for_status()  # Raise HTTPError for bad responses
-      result = response.json()  # Try to decode JSON
-      self._logger.debug("API Response body: %s", result)
-      return result
-    except requests.exceptions.HTTPError:
+      response.raise_for_status()  # Raise HTTPStatusError for bad responses
+      return response.json()  # Try to decode JSON
+    except httpx.HTTPStatusError:
       error_details = response.content.decode("utf-8")
       self._logger.warning(
           "API call failed for tool %s: Status %d - %s",
@@ -556,3 +544,10 @@ def __repr__(self):
         f' auth_scheme="{self.auth_scheme}",'
         f' auth_credential="{self.auth_credential}")'
     )
+
+
+async def _request(**request_params) -> httpx.Response:
+  async with httpx.AsyncClient(
+      verify=request_params.pop("verify", True)
+  ) as client:
+    return await client.request(**request_params)

src/google/adk/utils/_google_client_headers.py

@@ -0,0 +1,56 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from ._client_labels_utils import get_client_labels
+
+
+def get_tracking_headers() -> dict[str, str]:
+  """Returns a dictionary of HTTP headers for tracking API requests.
+
+  These headers are used to identify HTTP calls made by ADK towards
+   Vertex AI LLM APIs.
+  """
+  labels = get_client_labels()
+  header_value = " ".join(labels)
+  return {
+      "x-goog-api-client": header_value,
+      "user-agent": header_value,
+  }
+
+
+def merge_tracking_headers(headers: dict[str, str] | None) -> dict[str, str]:
+  """Merge tracking headers to the given headers.
+
+  Args:
+    headers: headers to merge tracking headers into.
+
+  Returns:
+    A dictionary of HTTP headers with tracking headers merged.
+  """
+  new_headers = (headers or {}).copy()
+  for key, tracking_header_value in get_tracking_headers().items():
+    custom_value = new_headers.get(key, None)
+    if not custom_value:
+      new_headers[key] = tracking_header_value
+      continue
+
+    # Merge tracking headers with existing headers and avoid duplicates.
+    value_parts = tracking_header_value.split(" ")
+    for custom_value_part in custom_value.split(" "):
+      if custom_value_part not in value_parts:
+        value_parts.append(custom_value_part)
+    new_headers[key] = " ".join(value_parts)
+  return new_headers

tests/unittests/models/test_anthropic_llm.py

@@ -391,6 +391,31 @@ async def mock_coro():
       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