diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py index 795b706559f..6e40439d607 100644 --- a/src/google/adk/flows/llm_flows/functions.py +++ b/src/google/adk/flows/llm_flows/functions.py @@ -668,7 +668,7 @@ async def handle_function_calls_live( invocation_context: InvocationContext, function_call_event: Event, tools_dict: dict[str, BaseTool], -) -> Event: +) -> Event | None: """Calls the functions and returns the function response event.""" from ...agents.llm_agent import LlmAgent diff --git a/src/google/adk/integrations/bigquery/data_insights_tool.py b/src/google/adk/integrations/bigquery/data_insights_tool.py index 37959c02a3d..7d231f98808 100644 --- a/src/google/adk/integrations/bigquery/data_insights_tool.py +++ b/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", diff --git a/src/google/adk/tools/_gda_stream_util.py b/src/google/adk/tools/_gda_stream_util.py index b8a68631683..67f84a788a1 100644 --- a/src/google/adk/tools/_gda_stream_util.py +++ b/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: diff --git a/src/google/adk/tools/data_agent/data_agent_tool.py b/src/google/adk/tools/data_agent/data_agent_tool.py index 007fe71158c..bcf187725ed 100644 --- a/src/google/adk/tools/data_agent/data_agent_tool.py +++ b/src/google/adk/tools/data_agent/data_agent_tool.py @@ -13,7 +13,6 @@ # limitations under the License. from __future__ import annotations -import json from typing import Any from google.auth.credentials import Credentials @@ -23,30 +22,9 @@ from ..tool_context import ToolContext from .config import DataAgentToolConfig -BASE_URL = "https://geminidataanalytics.googleapis.com/v1beta" _GDA_CLIENT_ID = "GOOGLE_ADK" -def _get_http_headers( - credentials: Credentials, -) -> dict[str, str]: - """Prepares headers for HTTP requests.""" - if not credentials.token: - error_details = ( - "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." - ) - raise ValueError(error_details) - return { - "Authorization": f"Bearer {credentials.token}", - "Content-Type": "application/json", - "X-Goog-API-Client": _GDA_CLIENT_ID, - } - - def list_accessible_data_agents( project_id: str, credentials: Credentials, @@ -116,12 +94,18 @@ def list_accessible_data_agents( } """ try: - headers = _get_http_headers(credentials) - list_url = f"{BASE_URL}/projects/{project_id}/locations/global/dataAgents:listAccessible" - resp = requests.get( - list_url, - headers=headers, - ) + session, endpoint = _gda_stream_util.get_gda_session(credentials) + base_url = f"{endpoint}/v1beta" + headers = { + "Content-Type": "application/json", + "X-Goog-API-Client": _GDA_CLIENT_ID, + } + list_url = f"{base_url}/projects/{project_id}/locations/global/dataAgents:listAccessible" + with session: + resp = session.get( + list_url, + headers=headers, + ) resp.raise_for_status() return { "status": "SUCCESS", @@ -134,6 +118,45 @@ def list_accessible_data_agents( } +def _get_data_agent_info( + data_agent_name: str, + credentials: Credentials, + session: requests.Session | None = None, +) -> dict[str, Any]: + try: + endpoint = _gda_stream_util.get_gda_endpoint() + base_url = f"{endpoint}/v1beta" + headers = { + "Content-Type": "application/json", + "X-Goog-API-Client": _GDA_CLIENT_ID, + } + get_url = f"{base_url}/{data_agent_name}" + + if session: + resp = session.get( + get_url, + headers=headers, + ) + else: + local_session, _ = _gda_stream_util.get_gda_session(credentials) + with local_session: + resp = local_session.get( + get_url, + headers=headers, + ) + + resp.raise_for_status() + return { + "status": "SUCCESS", + "response": resp.json(), + } + except Exception as ex: # pylint: disable=broad-except + return { + "status": "ERROR", + "error_details": str(ex), + } + + def get_data_agent_info( data_agent_name: str, credentials: Credentials, @@ -182,23 +205,7 @@ def get_data_agent_info( } } """ - try: - headers = _get_http_headers(credentials) - get_url = f"{BASE_URL}/{data_agent_name}" - resp = requests.get( - get_url, - headers=headers, - ) - resp.raise_for_status() - return { - "status": "SUCCESS", - "response": resp.json(), - } - except Exception as ex: # pylint: disable=broad-except - return { - "status": "ERROR", - "error_details": str(ex), - } + return _get_data_agent_info(data_agent_name, credentials) def ask_data_agent( @@ -280,26 +287,35 @@ def ask_data_agent( } """ try: - headers = _get_http_headers(credentials) + session, endpoint = _gda_stream_util.get_gda_session(credentials) + with session: + base_url = f"{endpoint}/v1beta" + headers = { + "Content-Type": "application/json", + "X-Goog-API-Client": _GDA_CLIENT_ID, + } - agent_info = get_data_agent_info(data_agent_name, credentials) - if agent_info.get("status") == "ERROR": - return agent_info - parent = data_agent_name.rsplit("/", 2)[0] - chat_url = f"{BASE_URL}/{parent}:chat" - chat_payload = { - "messages": [{"userMessage": {"text": query}}], - "dataAgentContext": { - "dataAgent": data_agent_name, - }, - "clientIdEnum": _GDA_CLIENT_ID, - } - resp = _gda_stream_util.get_stream( - chat_url, - chat_payload, - headers, - settings.max_query_result_rows, - ) + agent_info = _get_data_agent_info( + data_agent_name, credentials, session=session + ) + if agent_info.get("status") == "ERROR": + return agent_info + parent = data_agent_name.rsplit("/", 2)[0] + chat_url = f"{base_url}/{parent}:chat" + chat_payload = { + "messages": [{"userMessage": {"text": query}}], + "dataAgentContext": { + "dataAgent": data_agent_name, + }, + "clientIdEnum": _GDA_CLIENT_ID, + } + resp = _gda_stream_util.get_stream( + session, + chat_url, + chat_payload, + headers, + settings.max_query_result_rows, + ) return {"status": "SUCCESS", "response": resp} except Exception as ex: # pylint: disable=broad-except diff --git a/tests/unittests/integrations/bigquery/test_bigquery_data_insights_tool.py b/tests/unittests/integrations/bigquery/test_bigquery_data_insights_tool.py index 4c022ffe215..3614d7ae362 100644 --- a/tests/unittests/integrations/bigquery/test_bigquery_data_insights_tool.py +++ b/tests/unittests/integrations/bigquery/test_bigquery_data_insights_tool.py @@ -26,8 +26,10 @@ pytest.param("test_data/ask_data_insights_penguins_highest_mass.yaml"), ], ) -@mock.patch.object(data_insights_tool.requests.Session, "post") -def test_ask_data_insights_pipeline_from_file(mock_post, case_file_path): +@mock.patch.object( + data_insights_tool._gda_stream_util, "get_gda_session", autospec=True +) +def test_ask_data_insights_pipeline_from_file(mock_get_session, case_file_path): """Runs a full integration test for the ask_data_insights pipeline using data from a specific file.""" # 1. Construct the full, absolute path to the data file full_path = pathlib.Path(__file__).parent / case_file_path @@ -45,15 +47,19 @@ def test_ask_data_insights_pipeline_from_file(mock_post, case_file_path): expected_final_list = case_data["expected_output"] # 4. Configure the mock for requests.post + mock_session = mock.MagicMock() mock_response = mock.Mock() mock_response.iter_lines.return_value = fake_stream_lines # Add raise_for_status mock which is called in the updated code mock_response.raise_for_status.return_value = None - mock_post.return_value.__enter__.return_value = mock_response + mock_session.post.return_value.__enter__.return_value = mock_response + mock_get_session.return_value = ( + mock_session, + "https://geminidataanalytics.mtls.googleapis.com", + ) # 5. Call the function under test mock_creds = mock.Mock() - mock_creds.token = "fake-token" mock_settings = mock.Mock() mock_settings.max_query_result_rows = 50 result = data_insights_tool.ask_data_insights( @@ -67,10 +73,32 @@ def test_ask_data_insights_pipeline_from_file(mock_post, case_file_path): # 6. Assert that the final list of dicts matches the expected output assert result["status"] == "SUCCESS" assert result["response"] == expected_final_list + mock_get_session.assert_called_once_with(mock_creds) + mock_session.post.assert_called_once_with( + "https://geminidataanalytics.mtls.googleapis.com/v1beta/projects/test-project/locations/global:chat", + json={ + "project": "projects/test-project", + "messages": [{"userMessage": {"text": case_data["user_question"]}}], + "inlineContext": { + "datasourceReferences": {"bq": {"tableReferences": []}}, + "systemInstruction": mock.ANY, + "options": {"chart": {"image": {"noImage": {}}}}, + }, + "clientIdEnum": "GOOGLE_ADK", + }, + headers={ + "Content-Type": "application/json", + "X-Goog-API-Client": "GOOGLE_ADK", + }, + stream=True, + ) @mock.patch.object(data_insights_tool._gda_stream_util, "get_stream") -def test_ask_data_insights_success(mock_get_stream): +@mock.patch.object( + data_insights_tool._gda_stream_util, "get_gda_session", autospec=True +) +def test_ask_data_insights_success(mock_get_session, mock_get_stream): """Tests the success path of ask_data_insights using decorators.""" # 1. Configure the behavior of the mocked functions mock_stream = [ @@ -78,10 +106,14 @@ def test_ask_data_insights_success(mock_get_stream): {"text": {"parts": ["response2"], "textType": "FINAL_RESPONSE"}}, ] mock_get_stream.return_value = mock_stream + mock_session = mock.MagicMock() + mock_get_session.return_value = ( + mock_session, + "https://geminidataanalytics.mtls.googleapis.com", + ) # 2. Create mock inputs for the function call mock_creds = mock.Mock() - mock_creds.token = "fake-token" mock_settings = mock.Mock() mock_settings.max_query_result_rows = 100 @@ -97,24 +129,44 @@ def test_ask_data_insights_success(mock_get_stream): # 4. Assert the results are as expected assert result["status"] == "SUCCESS" assert result["response"] == mock_stream - mock_get_stream.assert_called_once() - - # Verify that the correct headers and client ID were passed to _get_stream - args, _ = mock_get_stream.call_args - headers = args[2] - assert headers["X-Goog-API-Client"] == "GOOGLE_ADK" - assert headers["Authorization"] == "Bearer fake-token" + mock_get_session.assert_called_once_with(mock_creds) + mock_get_stream.assert_called_once_with( + mock_session, + "https://geminidataanalytics.mtls.googleapis.com/v1beta/projects/test-project/locations/global:chat", + { + "project": "projects/test-project", + "messages": [{"userMessage": {"text": "test query"}}], + "inlineContext": { + "datasourceReferences": {"bq": {"tableReferences": []}}, + "systemInstruction": mock.ANY, + "options": {"chart": {"image": {"noImage": {}}}}, + }, + "clientIdEnum": "GOOGLE_ADK", + }, + { + "Content-Type": "application/json", + "X-Goog-API-Client": "GOOGLE_ADK", + }, + 100, + ) @mock.patch.object(data_insights_tool._gda_stream_util, "get_stream") -def test_ask_data_insights_handles_exception(mock_get_stream): +@mock.patch.object( + data_insights_tool._gda_stream_util, "get_gda_session", autospec=True +) +def test_ask_data_insights_handles_exception(mock_get_session, mock_get_stream): """Tests the exception path of ask_data_insights using decorators.""" # 1. Configure one of the mocks to raise an error mock_get_stream.side_effect = Exception("API call failed!") + mock_session = mock.MagicMock() + mock_get_session.return_value = ( + mock_session, + "https://geminidataanalytics.mtls.googleapis.com", + ) # 2. Create mock inputs mock_creds = mock.Mock() - mock_creds.token = "fake-token" mock_settings = mock.Mock() # 3. Call the function @@ -129,4 +181,5 @@ def test_ask_data_insights_handles_exception(mock_get_stream): # 4. Assert that the error was caught and formatted correctly assert result["status"] == "ERROR" assert "API call failed!" in result["error_details"] + mock_get_session.assert_called_once_with(mock_creds) mock_get_stream.assert_called_once() diff --git a/tests/unittests/tools/data_agent/test_data_agent_tool.py b/tests/unittests/tools/data_agent/test_data_agent_tool.py index 47be0ab7654..43556f74ada 100644 --- a/tests/unittests/tools/data_agent/test_data_agent_tool.py +++ b/tests/unittests/tools/data_agent/test_data_agent_tool.py @@ -12,99 +12,137 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pathlib from unittest import mock from google.adk.tools.data_agent import data_agent_tool from google.adk.tools.tool_context import ToolContext -import pytest -import requests -import yaml -@mock.patch.object(data_agent_tool, "requests", autospec=True) -def test_list_accessible_data_agents_success(mock_requests): +@mock.patch.object( + data_agent_tool._gda_stream_util, "get_gda_session", autospec=True +) +def test_list_accessible_data_agents_success(mock_get_session): """Tests list_accessible_data_agents success path.""" mock_creds = mock.Mock() - mock_creds.token = "fake-token" + mock_session = mock.MagicMock() mock_response = mock.Mock() mock_response.json.return_value = {"dataAgents": ["agent1", "agent2"]} mock_response.raise_for_status.return_value = None - mock_requests.get.return_value = mock_response + mock_session.get.return_value = mock_response + mock_get_session.return_value = ( + mock_session, + "https://geminidataanalytics.googleapis.com", + ) result = data_agent_tool.list_accessible_data_agents( "test-project", mock_creds ) assert result["status"] == "SUCCESS" assert result["response"] == ["agent1", "agent2"] - mock_requests.get.assert_called_once_with( + mock_get_session.assert_called_once_with(mock_creds) + mock_session.get.assert_called_once_with( "https://geminidataanalytics.googleapis.com/v1beta/projects/test-project/locations/global/dataAgents:listAccessible", headers={ - "Authorization": "Bearer fake-token", "Content-Type": "application/json", "X-Goog-API-Client": "GOOGLE_ADK", }, ) -@mock.patch.object(data_agent_tool, "requests", autospec=True) -def test_list_accessible_data_agents_exception(mock_requests): +@mock.patch.object( + data_agent_tool._gda_stream_util, "get_gda_session", autospec=True +) +def test_list_accessible_data_agents_exception(mock_get_session): """Tests list_accessible_data_agents exception path.""" mock_creds = mock.Mock() - mock_creds.token = "fake-token" - mock_requests.get.side_effect = Exception("List failed!") + mock_session = mock.MagicMock() + mock_session.get.side_effect = Exception("List failed!") + mock_get_session.return_value = ( + mock_session, + "https://geminidataanalytics.googleapis.com", + ) result = data_agent_tool.list_accessible_data_agents( "test-project", mock_creds ) assert result["status"] == "ERROR" assert "List failed!" in result["error_details"] - mock_requests.get.assert_called_once() + mock_get_session.assert_called_once_with(mock_creds) + mock_session.get.assert_called_once() -@mock.patch.object(data_agent_tool, "requests", autospec=True) -def test_get_data_agent_info_success(mock_requests): +@mock.patch.object( + data_agent_tool._gda_stream_util, "get_gda_endpoint", autospec=True +) +@mock.patch.object( + data_agent_tool._gda_stream_util, "get_gda_session", autospec=True +) +def test_get_data_agent_info_success(mock_get_session, mock_get_endpoint): """Tests get_data_agent_info success path.""" mock_creds = mock.Mock() - mock_creds.token = "fake-token" + mock_session = mock.MagicMock() mock_response = mock.Mock() mock_response.json.return_value = "agent_info" mock_response.raise_for_status.return_value = None - mock_requests.get.return_value = mock_response + mock_session.get.return_value = mock_response + mock_get_endpoint.return_value = "https://geminidataanalytics.googleapis.com" + mock_get_session.return_value = ( + mock_session, + "https://geminidataanalytics.googleapis.com", + ) result = data_agent_tool.get_data_agent_info("agent_name", mock_creds) assert result["status"] == "SUCCESS" assert result["response"] == "agent_info" - mock_requests.get.assert_called_once_with( + mock_get_session.assert_called_once_with(mock_creds) + mock_get_endpoint.assert_called_once() + mock_session.get.assert_called_once_with( "https://geminidataanalytics.googleapis.com/v1beta/agent_name", headers={ - "Authorization": "Bearer fake-token", "Content-Type": "application/json", "X-Goog-API-Client": "GOOGLE_ADK", }, ) -@mock.patch.object(data_agent_tool, "requests", autospec=True) -def test_get_data_agent_info_exception(mock_requests): +@mock.patch.object( + data_agent_tool._gda_stream_util, "get_gda_endpoint", autospec=True +) +@mock.patch.object( + data_agent_tool._gda_stream_util, "get_gda_session", autospec=True +) +def test_get_data_agent_info_exception(mock_get_session, mock_get_endpoint): """Tests get_data_agent_info exception path.""" mock_creds = mock.Mock() - mock_creds.token = "fake-token" - mock_requests.get.side_effect = Exception("Get failed!") + mock_session = mock.MagicMock() + mock_session.get.side_effect = Exception("Get failed!") + mock_get_endpoint.return_value = "https://geminidataanalytics.googleapis.com" + mock_get_session.return_value = ( + mock_session, + "https://geminidataanalytics.googleapis.com", + ) result = data_agent_tool.get_data_agent_info("agent_name", mock_creds) assert result["status"] == "ERROR" assert "Get failed!" in result["error_details"] - mock_requests.get.assert_called_once() + mock_get_session.assert_called_once_with(mock_creds) + mock_get_endpoint.assert_called_once() + mock_session.get.assert_called_once() @mock.patch.object( data_agent_tool._gda_stream_util, "get_stream", autospec=True ) -@mock.patch.object(data_agent_tool, "requests", autospec=True) -@mock.patch.object(data_agent_tool, "get_data_agent_info", autospec=True) +@mock.patch.object( + data_agent_tool._gda_stream_util, "get_gda_session", autospec=True +) +@mock.patch.object(data_agent_tool, "_get_data_agent_info", autospec=True) def test_ask_data_agent_success( - mock_get_agent_info, mock_requests, mock_get_stream + mock_get_agent_info, mock_get_session, mock_get_stream ): """Tests ask_data_agent success path.""" mock_creds = mock.Mock() - mock_creds.token = "fake-token" + mock_session = mock.MagicMock() + mock_get_session.return_value = ( + mock_session, + "https://geminidataanalytics.googleapis.com", + ) mock_get_agent_info.return_value = {"status": "SUCCESS", "response": {}} mock_get_stream.return_value = [ {"text": {"parts": ["response1"], "textType": "THOUGHT"}}, @@ -127,8 +165,12 @@ def test_ask_data_agent_success( {"text": {"parts": ["response1"], "textType": "THOUGHT"}}, {"text": {"parts": ["response2"], "textType": "FINAL_RESPONSE"}}, ] - mock_get_agent_info.assert_called_once() + mock_get_agent_info.assert_called_once_with( + "projects/p/locations/l/dataAgents/a", mock_creds, session=mock_session + ) + mock_get_session.assert_called_once_with(mock_creds) mock_get_stream.assert_called_once_with( + mock_session, "https://geminidataanalytics.googleapis.com/v1beta/projects/p/locations/l:chat", { "messages": [{"userMessage": {"text": "query"}}], @@ -138,7 +180,6 @@ def test_ask_data_agent_success( "clientIdEnum": "GOOGLE_ADK", }, { - "Authorization": "Bearer fake-token", "Content-Type": "application/json", "X-Goog-API-Client": "GOOGLE_ADK", }, @@ -149,14 +190,20 @@ def test_ask_data_agent_success( @mock.patch.object( data_agent_tool._gda_stream_util, "get_stream", autospec=True ) -@mock.patch.object(data_agent_tool, "requests", autospec=True) -@mock.patch.object(data_agent_tool, "get_data_agent_info", autospec=True) +@mock.patch.object( + data_agent_tool._gda_stream_util, "get_gda_session", autospec=True +) +@mock.patch.object(data_agent_tool, "_get_data_agent_info", autospec=True) def test_ask_data_agent_exception( - mock_get_agent_info, mock_requests, mock_get_stream + mock_get_agent_info, mock_get_session, mock_get_stream ): """Tests ask_data_agent exception path.""" mock_creds = mock.Mock() - mock_creds.token = "fake-token" + mock_session = mock.MagicMock() + mock_get_session.return_value = ( + mock_session, + "https://geminidataanalytics.googleapis.com", + ) mock_get_agent_info.return_value = {"status": "SUCCESS", "response": {}} mock_get_stream.side_effect = Exception("Chat failed!") mock_invocation_context = mock.Mock() @@ -173,4 +220,5 @@ def test_ask_data_agent_exception( ) assert result["status"] == "ERROR" assert "Chat failed!" in result["error_details"] + mock_get_session.assert_called_once_with(mock_creds) mock_get_stream.assert_called_once() diff --git a/tests/unittests/tools/test__gda_stream_util.py b/tests/unittests/tools/test__gda_stream_util.py index 6945e678334..3b5fc2ebc52 100644 --- a/tests/unittests/tools/test__gda_stream_util.py +++ b/tests/unittests/tools/test__gda_stream_util.py @@ -15,7 +15,6 @@ from unittest import mock from google.adk.tools import _gda_stream_util -import requests class MockResponse: @@ -114,8 +113,7 @@ def test_format_data_retrieved_missing_schema(self): }, ) - @mock.patch("requests.Session.post") - def test_get_stream(self, mock_post): + def test_get_stream(self): stream_lines = [ b"[{", b'"systemMessage": {"text": "msg1"}', @@ -139,8 +137,9 @@ def test_get_stream(self, mock_post): b'"systemMessage": {"text": "msg4"}', b"}]", ] - mock_post.return_value = MockResponse(stream_lines) - messages = _gda_stream_util.get_stream("url", {}, {}, 10) + mock_session = mock.MagicMock() + mock_session.post.return_value = MockResponse(stream_lines) + messages = _gda_stream_util.get_stream(mock_session, "url", {}, {}, 10) self.assertEqual(len(messages), 4) self.assertEqual(messages[0], {"text": "msg1"}) self.assertEqual( @@ -158,6 +157,93 @@ def test_get_stream(self, mock_post): ) self.assertEqual(messages[3], {"text": "msg4"}) + @mock.patch.object( + _gda_stream_util._mtls_utils, "get_api_endpoint", autospec=True + ) + @mock.patch.object( + _gda_stream_util.mtls, "has_default_client_cert_source", autospec=True + ) + @mock.patch.object( + _gda_stream_util.auth_requests, "AuthorizedSession", autospec=True + ) + def test_get_gda_session_use_mtls_and_cert( + self, mock_authorized_session, mock_use_client_cert, mock_get_api_endpoint + ): + mock_session = mock.MagicMock() + mock_authorized_session.return_value = mock_session + mock_use_client_cert.return_value = True + mock_get_api_endpoint.return_value = ( + "https://geminidataanalytics.mtls.googleapis.com" + ) + + creds = mock.MagicMock() + session, endpoint = _gda_stream_util.get_gda_session(creds) + + self.assertEqual(session, mock_session) + self.assertEqual( + endpoint, "https://geminidataanalytics.mtls.googleapis.com" + ) + mock_session.configure_mtls_channel.assert_called_once() + mock_get_api_endpoint.assert_called_once_with( + location="", + default_template="https://geminidataanalytics.googleapis.com", + mtls_template="https://geminidataanalytics.mtls.googleapis.com", + ) + + @mock.patch.object( + _gda_stream_util._mtls_utils, "get_api_endpoint", autospec=True + ) + @mock.patch.object( + _gda_stream_util.mtls, "has_default_client_cert_source", autospec=True + ) + @mock.patch.object( + _gda_stream_util.auth_requests, "AuthorizedSession", autospec=True + ) + def test_get_gda_session_use_mtls_no_cert( + self, mock_authorized_session, mock_use_client_cert, mock_get_api_endpoint + ): + mock_session = mock.MagicMock() + mock_authorized_session.return_value = mock_session + mock_use_client_cert.return_value = False + mock_get_api_endpoint.return_value = ( + "https://geminidataanalytics.mtls.googleapis.com" + ) + + creds = mock.MagicMock() + with self.assertRaises(ValueError) as context: + _gda_stream_util.get_gda_session(creds) + + self.assertIn( + "mTLS endpoint is selected, but client certificate is not provisioned", + str(context.exception), + ) + + @mock.patch.object( + _gda_stream_util._mtls_utils, "get_api_endpoint", autospec=True + ) + @mock.patch.object( + _gda_stream_util.mtls, "has_default_client_cert_source", autospec=True + ) + @mock.patch.object( + _gda_stream_util.auth_requests, "AuthorizedSession", autospec=True + ) + def test_get_gda_session_regular_endpoint( + self, mock_authorized_session, mock_use_client_cert, mock_get_api_endpoint + ): + mock_session = mock.MagicMock() + mock_authorized_session.return_value = mock_session + mock_use_client_cert.return_value = True + mock_get_api_endpoint.return_value = ( + "https://geminidataanalytics.googleapis.com" + ) + + creds = mock.MagicMock() + session, endpoint = _gda_stream_util.get_gda_session(creds) + + self.assertEqual(session, mock_session) + self.assertEqual(endpoint, "https://geminidataanalytics.googleapis.com") + mock_session.configure_mtls_channel.assert_not_called() + if __name__ == "__main__": unittest.main()