feat: support mTLS and GOOGLE_API_USE_MTLS_ENDPOINT for GDA client

PiperOrigin-RevId: 941478713

Author: google-genai-bot

src/google/adk/integrations/bigquery/data_insights_tool.py

@@ -13,17 +13,13 @@
 # limitations under the License.
 from __future__ import annotations
 
-import json
 from typing import Any
 from typing import Dict
 from typing import List
 
 from google.adk.tools import _gda_stream_util
 from google.auth.credentials import Credentials
-from google.cloud import bigquery
-import requests
 
-from . import client
 from .config import BigQueryToolConfig
 
 _GDA_CLIENT_ID = "GOOGLE_ADK"
@@ -117,48 +113,39 @@ def ask_data_insights(
   """
   try:
     location = "global"
-    if not credentials.token:
-      error_message = (
-          "Error: The provided credentials object does not have a valid access"
-          " token.\n\nThis is often because the credentials need to be"
-          " refreshed or require specific API scopes. Please ensure the"
-          " credentials are prepared correctly before calling this"
-          " function.\n\nThere may be other underlying causes as well."
-      )
-      return {
-          "status": "ERROR",
-          "error_details": "ask_data_insights requires a valid access token.",
+    session, endpoint = _gda_stream_util.get_gda_session(credentials)
+    with session:
+      headers = {
+          "Content-Type": "application/json",
+          "X-Goog-API-Client": _GDA_CLIENT_ID,
       }
-    headers = {
-        "Authorization": f"Bearer {credentials.token}",
-        "Content-Type": "application/json",
-        "X-Goog-API-Client": _GDA_CLIENT_ID,
-    }
-    ca_url = f"https://geminidataanalytics.googleapis.com/v1beta/projects/{project_id}/locations/{location}:chat"
+      ca_url = (
+          f"{endpoint}/v1beta/projects/{project_id}/locations/{location}:chat"
+      )
 
-    instructions = """**INSTRUCTIONS - FOLLOW THESE RULES:**
+      instructions = """**INSTRUCTIONS - FOLLOW THESE RULES:**
     1.  **CONTENT:** Your answer should present the supporting data and then provide a conclusion based on that data, including relevant details and observations where possible.
     2.  **ANALYSIS DEPTH:** Your analysis must go beyond surface-level observations. Crucially, you must prioritize metrics that measure impact or outcomes over metrics that simply measure volume or raw counts. For open-ended questions, explore the topic from multiple perspectives to provide a holistic view.
     3.  **OUTPUT FORMAT:** Your entire response MUST be in plain text format ONLY.
     4.  **NO CHARTS:** You are STRICTLY FORBIDDEN from generating any charts, graphs, images, or any other form of visualization.
     """
 
-    ca_payload = {
-        "project": f"projects/{project_id}",
-        "messages": [{"userMessage": {"text": user_query_with_context}}],
-        "inlineContext": {
-            "datasourceReferences": {
-                "bq": {"tableReferences": table_references}
-            },
-            "systemInstruction": instructions,
-            "options": {"chart": {"image": {"noImage": {}}}},
-        },
-        "clientIdEnum": _GDA_CLIENT_ID,
-    }
+      ca_payload = {
+          "project": f"projects/{project_id}",
+          "messages": [{"userMessage": {"text": user_query_with_context}}],
+          "inlineContext": {
+              "datasourceReferences": {
+                  "bq": {"tableReferences": table_references}
+              },
+              "systemInstruction": instructions,
+              "options": {"chart": {"image": {"noImage": {}}}},
+          },
+          "clientIdEnum": _GDA_CLIENT_ID,
+      }
 
-    resp = _gda_stream_util.get_stream(
-        ca_url, ca_payload, headers, settings.max_query_result_rows
-    )
+      resp = _gda_stream_util.get_stream(
+          session, ca_url, ca_payload, headers, settings.max_query_result_rows
+      )
   except Exception as ex:  # pylint: disable=broad-except
     return {
         "status": "ERROR",

src/google/adk/tools/_gda_stream_util.py

@@ -16,69 +16,116 @@
 import json
 from typing import Any
 
+from google.auth.transport import mtls
+from google.auth.transport import requests as auth_requests
 import requests
 
+from google import auth
+
+from ..utils import _mtls_utils
+
+_GDA_DEFAULT_TEMPLATE = "https://geminidataanalytics.googleapis.com"
+_GDA_MTLS_TEMPLATE = "https://geminidataanalytics.mtls.googleapis.com"
+
+
+def get_gda_endpoint() -> str:
+  """Returns the GDA API endpoint based on mTLS configuration."""
+  return _mtls_utils.get_api_endpoint(
+      location="",
+      default_template=_GDA_DEFAULT_TEMPLATE,
+      mtls_template=_GDA_MTLS_TEMPLATE,
+  )
+
+
+def get_gda_session(
+    credentials: auth.credentials.Credentials,
+) -> tuple[requests.Session, str]:
+  """Creates an AuthorizedSession and returns it with the correct endpoint.
+
+  Args:
+      credentials: The credentials to use for the request.
+
+  Returns:
+      A tuple containing the authorized requests Session and the GDA endpoint.
+
+  Raises:
+      ValueError: If the mTLS endpoint is selected but the client certificate
+        is disabled.
+  """
+  session = auth_requests.AuthorizedSession(credentials=credentials)  # type: ignore[no-untyped-call]
+  endpoint = get_gda_endpoint()
+
+  if endpoint == _GDA_MTLS_TEMPLATE:
+    if not mtls.has_default_client_cert_source():  # type: ignore[no-untyped-call]
+      raise ValueError(
+          "mTLS endpoint is selected, but client certificate is not"
+          " provisioned."
+      )
+    session.configure_mtls_channel()  # type: ignore[no-untyped-call]
+
+  return session, endpoint
+
 
 def get_stream(
+    session: requests.Session,
     url: str,
     ca_payload: dict[str, Any],
     headers: dict[str, str],
     max_query_result_rows: int,
 ) -> list[dict[str, Any]]:
   """Sends a JSON request to a streaming API and returns a list of messages."""
-  with requests.Session() as s:
-    accumulator = ""
-    messages = []
-    data_msg_idx = -1
-
-    with s.post(url, json=ca_payload, headers=headers, stream=True) as resp:
-      resp.raise_for_status()
-      for line in resp.iter_lines():
-        if not line:
-          continue
-
-        decoded_line = line.decode("utf-8")
-
-        if decoded_line == "[{":
-          accumulator = "{"
-        elif decoded_line == "}]":
-          accumulator += "}"
-        elif decoded_line == ",":
-          continue
-        else:
-          accumulator += decoded_line
-
-        try:
-          data_json = json.loads(accumulator)
-        except ValueError:
-          continue
-
-        accumulator = ""
-
-        if not isinstance(data_json, dict):
-          messages.append(data_json)
-          continue
-
-        processed_msg = None
-        data_result = _extract_data_result(data_json)
-        if data_result is not None:
-          processed_msg = _format_data_retrieved(
-              data_result, max_query_result_rows
-          )
-          if data_msg_idx >= 0:
-            messages[data_msg_idx] = {
-                "Data Retrieved": "Intermediate result omitted"
-            }
-          data_msg_idx = len(messages)
-        elif isinstance(data_json.get("systemMessage"), dict):
-          processed_msg = data_json["systemMessage"]
-        else:
-          processed_msg = data_json
-
-        if processed_msg is not None:
-          messages.append(processed_msg)
-
-    return messages
+  accumulator = ""
+  messages = []
+  data_msg_idx = -1
+
+  with session.post(url, json=ca_payload, headers=headers, stream=True) as resp:
+    resp.raise_for_status()
+    for line in resp.iter_lines():
+      if not line:
+        continue
+
+      decoded_line = line.decode("utf-8")
+
+      if decoded_line == "[{":
+        accumulator = "{"
+      elif decoded_line == "}]":
+        accumulator += "}"
+      elif decoded_line == ",":
+        continue
+      else:
+        accumulator += decoded_line
+
+      try:
+        data_json = json.loads(accumulator)
+      except ValueError:
+        continue
+
+      accumulator = ""
+
+      if not isinstance(data_json, dict):
+        messages.append(data_json)
+        continue
+
+      processed_msg = None
+      data_result = _extract_data_result(data_json)
+      if data_result is not None:
+        processed_msg = _format_data_retrieved(
+            data_result, max_query_result_rows
+        )
+        if data_msg_idx >= 0:
+          messages[data_msg_idx] = {
+              "Data Retrieved": "Intermediate result omitted"
+          }
+        data_msg_idx = len(messages)
+      elif isinstance(data_json.get("systemMessage"), dict):
+        processed_msg = data_json["systemMessage"]
+      else:
+        processed_msg = data_json
+
+      if processed_msg is not None:
+        messages.append(processed_msg)
+
+  return messages
 
 
 def _extract_data_result(msg: dict[str, Any]) -> dict[str, Any] | None: