Compare commits

...

7 Commits

Author SHA1 Message Date
Shangjie Chen 8677d5c8dc chore: bump version number to 1.4.2
PiperOrigin-RevId: 773867075
2025-06-20 15:41:50 -07:00
Shangjie Chen 4d72d31b13 fix: Add type checking to handle different response type of genai API client
Fixes https://github.com/google/adk-python/issues/1514

PiperOrigin-RevId: 773838035
2025-06-20 14:12:55 -07:00
Xiang (Sean) Zhou 742478fdb7 chore: Add event converters to convert adk event to a2a event (WIP)
PiperOrigin-RevId: 773795427
2025-06-20 12:14:12 -07:00
Xiang (Sean) Zhou ffcba70686 chore: skip mcp and a2a tests for python 3.9
PiperOrigin-RevId: 773785385
2025-06-20 11:46:31 -07:00
Google Team Member 2f716ada7f fix: Allow more credentials types for BigQuery tools
This change accepts the `google.auth.credentials.Credentials` type for `BigQueryCredentialsConfig`, so any subclass of that, including `google.oauth2.credentials.Credentials` would work to integrate with BigQuery service. This opens up a whole range of possibilities, such as using service account credentials to deploy an agent using these tools.

PiperOrigin-RevId: 773190440
2025-06-18 22:02:09 -07:00
Shangjie Chen 17beb32880 chore: Bump version number and update changelog for 1.4.1 release
PiperOrigin-RevId: 773148349
2025-06-18 19:08:43 -07:00
Xiang (Sean) Zhou 7f8dc8927a chore: fix the mcp_sse_agent
PiperOrigin-RevId: 773140021
2025-06-18 18:32:06 -07:00
21 changed files with 1355 additions and 97 deletions
+44
View File
@@ -1,5 +1,49 @@
# Changelog
## [1.4.2](https://github.com/google/adk-python/compare/v1.4.1...v1.4.2) (2025-06-20)
### Bug Fixes
* Add type checking to handle different response type of genai API client ([4d72d31](https://github.com/google/adk-python/commit/4d72d31b13f352245baa72b78502206dcbe25406))
* This fixes the broken VertexAiSessionService
* Allow more credentials types for BigQuery tools ([2f716ad](https://github.com/google/adk-python/commit/2f716ada7fbcf8e03ff5ae16ce26a80ca6fd7bf6))
## [1.4.1](https://github.com/google/adk-python/compare/v1.3.0...v1.4.1) (2025-06-18)
### Features
* Add Authenticated Tool (Experimental) ([dcea776](https://github.com/google/adk-python/commit/dcea7767c67c7edfb694304df32dca10b74c9a71))
* Add enable_affective_dialog and proactivity to run_config and llm_request ([fe1d5aa](https://github.com/google/adk-python/commit/fe1d5aa439cc56b89d248a52556c0a9b4cbd15e4))
* Add import session API in the fast API ([233fd20](https://github.com/google/adk-python/commit/233fd2024346abd7f89a16c444de0cf26da5c1a1))
* Add integration tests for litellm with and without turning on add_function_to_prompt ([8e28587](https://github.com/google/adk-python/commit/8e285874da7f5188ea228eb4d7262dbb33b1ae6f))
* Allow data_store_specs pass into ADK VAIS built-in tool ([675faef](https://github.com/google/adk-python/commit/675faefc670b5cd41991939fe0fc604df331111a))
* Enable MCP Tool Auth (Experimental) ([157d9be](https://github.com/google/adk-python/commit/157d9be88d92f22320604832e5a334a6eb81e4af))
* Implement GcsEvalSetResultsManager to handle storage of eval sets on GCS, and refactor eval set results manager ([0a5cf45](https://github.com/google/adk-python/commit/0a5cf45a75aca7b0322136b65ca5504a0c3c7362))
* Re-factor some eval sets manager logic, and implement GcsEvalSetsManager to handle storage of eval sets on GCS ([1551bd4](https://github.com/google/adk-python/commit/1551bd4f4d7042fffb497d9308b05f92d45d818f))
* Support real time input config ([d22920b](https://github.com/google/adk-python/commit/d22920bd7f827461afd649601326b0c58aea6716))
* Support refresh access token automatically for rest_api_tool ([1779801](https://github.com/google/adk-python/commit/177980106b2f7be9a8c0a02f395ff0f85faa0c5a))
### Bug Fixes
* Fix Agent generate config err ([#1305](https://github.com/google/adk-python/issues/1305)) ([badbcbd](https://github.com/google/adk-python/commit/badbcbd7a464e6b323cf3164d2bcd4e27cbc057f))
* Fix Agent generate config error ([#1450](https://github.com/google/adk-python/issues/1450)) ([694b712](https://github.com/google/adk-python/commit/694b71256c631d44bb4c4488279ea91d82f43e26))
* Fix liteLLM test failures ([fef8778](https://github.com/google/adk-python/commit/fef87784297b806914de307f48c51d83f977298f))
* Fix tracing for live ([58e07ca](https://github.com/google/adk-python/commit/58e07cae83048d5213d822be5197a96be9ce2950))
* Merge custom http options with adk specific http options in model api request ([4ccda99](https://github.com/google/adk-python/commit/4ccda99e8ec7aa715399b4b83c3f101c299a95e8))
* Remove unnecessary double quote on Claude docstring ([bbceb4f](https://github.com/google/adk-python/commit/bbceb4f2e89f720533b99cf356c532024a120dc4))
* Set explicit project in the BigQuery client ([6d174eb](https://github.com/google/adk-python/commit/6d174eba305a51fcf2122c0fd481378752d690ef))
* Support streaming in litellm + adk and add corresponding integration tests ([aafa80b](https://github.com/google/adk-python/commit/aafa80bd85a49fb1c1a255ac797587cffd3fa567))
* Support project-based gemini model path to use google_search_tool ([b2fc774](https://github.com/google/adk-python/commit/b2fc7740b363a4e33ec99c7377f396f5cee40b5a))
* Update conversion between Celsius and Fahrenheit ([1ae176a](https://github.com/google/adk-python/commit/1ae176ad2fa2b691714ac979aec21f1cf7d35e45))
### Chores
* Set `agent_engine_id` in the VertexAiSessionService constructor, also use the `agent_engine_id` field instead of overriding `app_name` in FastAPI endpoint ([fc65873](https://github.com/google/adk-python/commit/fc65873d7c31be607f6cd6690f142a031631582a))
## [1.3.0](https://github.com/google/adk-python/compare/v1.2.1...v1.3.0) (2025-06-11)
+19 -4
View File
@@ -40,13 +40,28 @@ would set:
### With Application Default Credentials
This mode is useful for quick development when the agent builder is the only
user interacting with the agent. The tools are initialized with the default
credentials present on the machine running the agent.
user interacting with the agent. The tools are run with these credentials.
1. Create application default credentials on the machine where the agent would
be running by following https://cloud.google.com/docs/authentication/provide-credentials-adc.
1. Set `RUN_WITH_ADC=True` in `agent.py` and run the agent
1. Set `CREDENTIALS_TYPE=None` in `agent.py`
1. Run the agent
### With Service Account Keys
This mode is useful for quick development when the agent builder wants to run
the agent with service account credentials. The tools are run with these
credentials.
1. Create service account key by following https://cloud.google.com/iam/docs/service-account-creds#user-managed-keys.
1. Set `CREDENTIALS_TYPE=AuthCredentialTypes.SERVICE_ACCOUNT` in `agent.py`
1. Download the key file and replace `"service_account_key.json"` with the path
1. Run the agent
### With Interactive OAuth
@@ -72,7 +87,7 @@ type.
Note: don't create a separate .env, instead put it to the same .env file that
stores your Vertex AI or Dev ML credentials
1. Set `RUN_WITH_ADC=False` in `agent.py` and run the agent
1. Set `CREDENTIALS_TYPE=AuthCredentialTypes.OAUTH2` in `agent.py` and run the agent
## Sample prompts
+19 -8
View File
@@ -15,24 +15,21 @@
import os
from google.adk.agents import llm_agent
from google.adk.auth import AuthCredentialTypes
from google.adk.tools.bigquery import BigQueryCredentialsConfig
from google.adk.tools.bigquery import BigQueryToolset
from google.adk.tools.bigquery.config import BigQueryToolConfig
from google.adk.tools.bigquery.config import WriteMode
import google.auth
RUN_WITH_ADC = False
# Define an appropriate credential type
CREDENTIALS_TYPE = AuthCredentialTypes.OAUTH2
# Define BigQuery tool config
tool_config = BigQueryToolConfig(write_mode=WriteMode.ALLOWED)
if RUN_WITH_ADC:
# Initialize the tools to use the application default credentials.
application_default_credentials, _ = google.auth.default()
credentials_config = BigQueryCredentialsConfig(
credentials=application_default_credentials
)
else:
if CREDENTIALS_TYPE == AuthCredentialTypes.OAUTH2:
# Initiaze the tools to do interactive OAuth
# The environment variables OAUTH_CLIENT_ID and OAUTH_CLIENT_SECRET
# must be set
@@ -40,6 +37,20 @@ else:
client_id=os.getenv("OAUTH_CLIENT_ID"),
client_secret=os.getenv("OAUTH_CLIENT_SECRET"),
)
elif CREDENTIALS_TYPE == AuthCredentialTypes.SERVICE_ACCOUNT:
# Initialize the tools to use the credentials in the service account key.
# If this flow is enabled, make sure to replace the file path with your own
# service account key file
# https://cloud.google.com/iam/docs/service-account-creds#user-managed-keys
creds, _ = google.auth.load_credentials_from_file("service_account_key.json")
credentials_config = BigQueryCredentialsConfig(credentials=creds)
else:
# Initialize the tools to use the application default credentials.
# https://cloud.google.com/docs/authentication/provide-credentials-adc
application_default_credentials, _ = google.auth.default()
credentials_config = BigQueryCredentialsConfig(
credentials=application_default_credentials
)
bigquery_toolset = BigQueryToolset(
credentials_config=credentials_config, bigquery_tool_config=tool_config
+2 -2
View File
@@ -16,8 +16,8 @@
import os
from google.adk.agents.llm_agent import LlmAgent
from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams
from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
from google.adk.tools.mcp_tool.mcp_toolset import SseServerParams
_allowed_path = os.path.dirname(os.path.abspath(__file__))
@@ -31,7 +31,7 @@ Allowed directory: {_allowed_path}
""",
tools=[
MCPToolset(
connection_params=SseServerParams(
connection_params=SseConnectionParams(
url='http://localhost:3000/sse',
headers={'Accept': 'text/event-stream'},
),
@@ -0,0 +1,382 @@
# 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
import datetime
import logging
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
import uuid
from a2a.server.events import Event as A2AEvent
from a2a.types import Artifact
from a2a.types import DataPart
from a2a.types import Message
from a2a.types import Role
from a2a.types import TaskArtifactUpdateEvent
from a2a.types import TaskState
from a2a.types import TaskStatus
from a2a.types import TaskStatusUpdateEvent
from a2a.types import TextPart
from ...agents.invocation_context import InvocationContext
from ...events.event import Event
from ...utils.feature_decorator import working_in_progress
from .part_converter import A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL
from .part_converter import A2A_DATA_PART_METADATA_TYPE_KEY
from .part_converter import convert_genai_part_to_a2a_part
from .utils import _get_adk_metadata_key
# Constants
ARTIFACT_ID_SEPARATOR = "-"
DEFAULT_ERROR_MESSAGE = "An error occurred during processing"
# Logger
logger = logging.getLogger("google_adk." + __name__)
def _serialize_metadata_value(value: Any) -> str:
"""Safely serializes metadata values to string format.
Args:
value: The value to serialize.
Returns:
String representation of the value.
"""
if hasattr(value, "model_dump"):
try:
return value.model_dump(exclude_none=True, by_alias=True)
except Exception as e:
logger.warning("Failed to serialize metadata value: %s", e)
return str(value)
return str(value)
def _get_context_metadata(
event: Event, invocation_context: InvocationContext
) -> Dict[str, str]:
"""Gets the context metadata for the event.
Args:
event: The ADK event to extract metadata from.
invocation_context: The invocation context containing session information.
Returns:
A dictionary containing the context metadata.
Raises:
ValueError: If required fields are missing from event or context.
"""
if not event:
raise ValueError("Event cannot be None")
if not invocation_context:
raise ValueError("Invocation context cannot be None")
try:
metadata = {
_get_adk_metadata_key("app_name"): invocation_context.app_name,
_get_adk_metadata_key("user_id"): invocation_context.user_id,
_get_adk_metadata_key("session_id"): invocation_context.session.id,
_get_adk_metadata_key("invocation_id"): event.invocation_id,
_get_adk_metadata_key("author"): event.author,
}
# Add optional metadata fields if present
optional_fields = [
("branch", event.branch),
("grounding_metadata", event.grounding_metadata),
("custom_metadata", event.custom_metadata),
("usage_metadata", event.usage_metadata),
("error_code", event.error_code),
]
for field_name, field_value in optional_fields:
if field_value is not None:
metadata[_get_adk_metadata_key(field_name)] = _serialize_metadata_value(
field_value
)
return metadata
except Exception as e:
logger.error("Failed to create context metadata: %s", e)
raise
def _create_artifact_id(
app_name: str, user_id: str, session_id: str, filename: str, version: int
) -> str:
"""Creates a unique artifact ID.
Args:
app_name: The application name.
user_id: The user ID.
session_id: The session ID.
filename: The artifact filename.
version: The artifact version.
Returns:
A unique artifact ID string.
"""
components = [app_name, user_id, session_id, filename, str(version)]
return ARTIFACT_ID_SEPARATOR.join(components)
def _convert_artifact_to_a2a_events(
event: Event,
invocation_context: InvocationContext,
filename: str,
version: int,
) -> TaskArtifactUpdateEvent:
"""Converts a new artifact version to an A2A TaskArtifactUpdateEvent.
Args:
event: The ADK event containing the artifact information.
invocation_context: The invocation context.
filename: The name of the artifact file.
version: The version number of the artifact.
Returns:
A TaskArtifactUpdateEvent representing the artifact update.
Raises:
ValueError: If required parameters are invalid.
RuntimeError: If artifact loading fails.
"""
if not filename:
raise ValueError("Filename cannot be empty")
if version < 0:
raise ValueError("Version must be non-negative")
try:
artifact_part = invocation_context.artifact_service.load_artifact(
app_name=invocation_context.app_name,
user_id=invocation_context.user_id,
session_id=invocation_context.session.id,
filename=filename,
version=version,
)
converted_part = convert_genai_part_to_a2a_part(part=artifact_part)
if not converted_part:
raise RuntimeError(f"Failed to convert artifact part for {filename}")
artifact_id = _create_artifact_id(
invocation_context.app_name,
invocation_context.user_id,
invocation_context.session.id,
filename,
version,
)
return TaskArtifactUpdateEvent(
taskId=str(uuid.uuid4()),
append=False,
contextId=invocation_context.session.id,
lastChunk=True,
artifact=Artifact(
artifactId=artifact_id,
name=filename,
metadata={
"filename": filename,
"version": version,
},
parts=[converted_part],
),
)
except Exception as e:
logger.error(
"Failed to convert artifact for %s, version %s: %s",
filename,
version,
e,
)
raise RuntimeError(f"Artifact conversion failed: {e}") from e
def _process_long_running_tool(a2a_part, event: Event) -> None:
"""Processes long-running tool metadata for an A2A part.
Args:
a2a_part: The A2A part to potentially mark as long-running.
event: The ADK event containing long-running tool information.
"""
if (
isinstance(a2a_part.root, DataPart)
and event.long_running_tool_ids
and a2a_part.root.metadata.get(
_get_adk_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY)
)
== A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL
and a2a_part.root.metadata.get("id") in event.long_running_tool_ids
):
a2a_part.root.metadata[_get_adk_metadata_key("is_long_running")] = True
@working_in_progress
def convert_event_to_a2a_status_message(
event: Event, invocation_context: InvocationContext
) -> Optional[Message]:
"""Converts an ADK event to an A2A message.
Args:
event: The ADK event to convert.
invocation_context: The invocation context.
Returns:
An A2A Message if the event has content, None otherwise.
Raises:
ValueError: If required parameters are invalid.
"""
if not event:
raise ValueError("Event cannot be None")
if not invocation_context:
raise ValueError("Invocation context cannot be None")
if not event.content or not event.content.parts:
return None
try:
a2a_parts = []
for part in event.content.parts:
a2a_part = convert_genai_part_to_a2a_part(part)
if a2a_part:
a2a_parts.append(a2a_part)
_process_long_running_tool(a2a_part, event)
if a2a_parts:
return Message(
messageId=str(uuid.uuid4()), role=Role.agent, parts=a2a_parts
)
except Exception as e:
logger.error("Failed to convert event to status message: %s", e)
raise
return None
def _create_error_status_event(
event: Event, invocation_context: InvocationContext
) -> TaskStatusUpdateEvent:
"""Creates a TaskStatusUpdateEvent for error scenarios.
Args:
event: The ADK event containing error information.
invocation_context: The invocation context.
Returns:
A TaskStatusUpdateEvent with FAILED state.
"""
error_message = getattr(event, "error_message", None) or DEFAULT_ERROR_MESSAGE
return TaskStatusUpdateEvent(
taskId=str(uuid.uuid4()),
contextId=invocation_context.session.id,
final=False,
metadata=_get_context_metadata(event, invocation_context),
status=TaskStatus(
state=TaskState.failed,
message=Message(
messageId=str(uuid.uuid4()),
role=Role.agent,
parts=[TextPart(text=error_message)],
),
timestamp=datetime.datetime.now().isoformat(),
),
)
def _create_running_status_event(
message: Message, invocation_context: InvocationContext, event: Event
) -> TaskStatusUpdateEvent:
"""Creates a TaskStatusUpdateEvent for running scenarios.
Args:
message: The A2A message to include.
invocation_context: The invocation context.
event: The ADK event.
Returns:
A TaskStatusUpdateEvent with RUNNING state.
"""
return TaskStatusUpdateEvent(
taskId=str(uuid.uuid4()),
contextId=invocation_context.session.id,
final=False,
status=TaskStatus(
state=TaskState.working,
message=message,
timestamp=datetime.datetime.now().isoformat(),
),
metadata=_get_context_metadata(event, invocation_context),
)
@working_in_progress
def convert_event_to_a2a_events(
event: Event, invocation_context: InvocationContext
) -> List[A2AEvent]:
"""Converts a GenAI event to a list of A2A events.
Args:
event: The ADK event to convert.
invocation_context: The invocation context.
Returns:
A list of A2A events representing the converted ADK event.
Raises:
ValueError: If required parameters are invalid.
"""
if not event:
raise ValueError("Event cannot be None")
if not invocation_context:
raise ValueError("Invocation context cannot be None")
a2a_events = []
try:
# Handle artifact deltas
if event.actions.artifact_delta:
for filename, version in event.actions.artifact_delta.items():
artifact_event = _convert_artifact_to_a2a_events(
event, invocation_context, filename, version
)
a2a_events.append(artifact_event)
# Handle error scenarios
if event.error_code:
error_event = _create_error_status_event(event, invocation_context)
a2a_events.append(error_event)
# Handle regular message content
message = convert_event_to_a2a_status_message(event, invocation_context)
if message:
running_event = _create_running_status_event(
message, invocation_context, event
)
a2a_events.append(running_event)
except Exception as e:
logger.error("Failed to convert event to A2A events: %s", e)
raise
return a2a_events
@@ -23,6 +23,8 @@ import logging
import sys
from typing import Optional
from .utils import _get_adk_metadata_key
try:
from a2a import types as a2a_types
except ImportError as e:
@@ -84,7 +86,7 @@ def convert_a2a_part_to_genai_part(
# logic accordinlgy
if part.metadata and A2A_DATA_PART_METADATA_TYPE_KEY in part.metadata:
if (
part.metadata[A2A_DATA_PART_METADATA_TYPE_KEY]
part.metadata[_get_adk_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY)]
== A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL
):
return genai_types.Part(
@@ -93,7 +95,7 @@ def convert_a2a_part_to_genai_part(
)
)
if (
part.metadata[A2A_DATA_PART_METADATA_TYPE_KEY]
part.metadata[_get_adk_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY)]
== A2A_DATA_PART_METADATA_TYPE_FUNCTION_RESPONSE
):
return genai_types.Part(
+34
View File
@@ -0,0 +1,34 @@
# 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
ADK_METADATA_KEY_PREFIX = "adk_"
def _get_adk_metadata_key(key: str) -> str:
"""Gets the A2A event metadata key for the given key.
Args:
key: The metadata key to prefix.
Returns:
The prefixed metadata key.
Raises:
ValueError: If key is empty or None.
"""
if not key:
raise ValueError("Metadata key cannot be empty or None")
return f"{ADK_METADATA_KEY_PREFIX}{key}"
@@ -14,6 +14,7 @@
from __future__ import annotations
import asyncio
import json
import logging
import re
from typing import Any
@@ -87,6 +88,7 @@ class VertexAiSessionService(BaseSessionService):
path=f'reasoningEngines/{reasoning_engine_id}/sessions',
request_dict=session_json_dict,
)
api_response = _convert_api_response(api_response)
logger.info(f'Create Session response {api_response}')
session_id = api_response['name'].split('/')[-3]
@@ -100,6 +102,7 @@ class VertexAiSessionService(BaseSessionService):
path=f'operations/{operation_id}',
request_dict={},
)
lro_response = _convert_api_response(lro_response)
if lro_response.get('done', None):
break
@@ -118,6 +121,7 @@ class VertexAiSessionService(BaseSessionService):
path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
request_dict={},
)
get_session_api_response = _convert_api_response(get_session_api_response)
update_timestamp = isoparse(
get_session_api_response['updateTime']
@@ -149,6 +153,7 @@ class VertexAiSessionService(BaseSessionService):
path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
request_dict={},
)
get_session_api_response = _convert_api_response(get_session_api_response)
session_id = get_session_api_response['name'].split('/')[-1]
update_timestamp = isoparse(
@@ -167,9 +172,12 @@ class VertexAiSessionService(BaseSessionService):
path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}/events',
request_dict={},
)
list_events_api_response = _convert_api_response(list_events_api_response)
# Handles empty response case
if list_events_api_response.get('httpHeaders', None):
if not list_events_api_response or list_events_api_response.get(
'httpHeaders', None
):
return session
session.events += [
@@ -226,9 +234,10 @@ class VertexAiSessionService(BaseSessionService):
path=path,
request_dict={},
)
api_response = _convert_api_response(api_response)
# Handles empty response case
if api_response.get('httpHeaders', None):
if not api_response or api_response.get('httpHeaders', None):
return ListSessionsResponse()
sessions = []
@@ -303,6 +312,13 @@ class VertexAiSessionService(BaseSessionService):
return client._api_client
def _convert_api_response(api_response):
"""Converts the API response to a JSON object based on the type."""
if hasattr(api_response, 'body'):
return json.loads(api_response.body)
return api_response
def _convert_event_to_json(event: Event) -> Dict[str, Any]:
metadata_json = {
'partial': event.partial,
@@ -21,9 +21,10 @@ from typing import Optional
from fastapi.openapi.models import OAuth2
from fastapi.openapi.models import OAuthFlowAuthorizationCode
from fastapi.openapi.models import OAuthFlows
import google.auth.credentials
from google.auth.exceptions import RefreshError
from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
import google.oauth2.credentials
from pydantic import BaseModel
from pydantic import model_validator
@@ -40,26 +41,35 @@ BIGQUERY_DEFAULT_SCOPE = ["https://www.googleapis.com/auth/bigquery"]
@experimental
class BigQueryCredentialsConfig(BaseModel):
"""Configuration for Google API tools. (Experimental)"""
"""Configuration for Google API tools (Experimental).
Please do not use this in production, as it may be deprecated later.
"""
# Configure the model to allow arbitrary types like Credentials
model_config = {"arbitrary_types_allowed": True}
credentials: Optional[Credentials] = None
"""the existing oauth credentials to use. If set,this credential will be used
credentials: Optional[google.auth.credentials.Credentials] = None
"""The existing auth credentials to use. If set, this credential will be used
for every end user, end users don't need to be involved in the oauthflow. This
field is mutually exclusive with client_id, client_secret and scopes.
Don't set this field unless you are sure this credential has the permission to
access every end user's data.
Example usage: when the agent is deployed in Google Cloud environment and
Example usage 1: When the agent is deployed in Google Cloud environment and
the service account (used as application default credentials) has access to
all the required BigQuery resource. Setting this credential to allow user to
access the BigQuery resource without end users going through oauth flow.
To get application default credential: `google.auth.default(...)`. See more
To get application default credential, use: `google.auth.default(...)`. See more
details in https://cloud.google.com/docs/authentication/application-default-credentials.
Example usage 2: When the agent wants to access the user's BigQuery resources
using the service account key credentials.
To load service account key credentials, use: `google.auth.load_credentials_from_file(...)`.
See more details in https://cloud.google.com/iam/docs/service-account-creds#user-managed-keys.
When the deployed environment cannot provide a pre-existing credential,
consider setting below client_id, client_secret and scope for end users to go
through oauth flow, so that agent can access the user data.
@@ -86,7 +96,9 @@ class BigQueryCredentialsConfig(BaseModel):
" client_id/client_secret/scopes."
)
if self.credentials:
if self.credentials and isinstance(
self.credentials, google.oauth2.credentials.Credentials
):
self.client_id = self.credentials.client_id
self.client_secret = self.credentials.client_secret
self.scopes = self.credentials.scopes
@@ -115,7 +127,7 @@ class BigQueryCredentialsManager:
async def get_valid_credentials(
self, tool_context: ToolContext
) -> Optional[Credentials]:
) -> Optional[google.auth.credentials.Credentials]:
"""Get valid credentials, handling refresh and OAuth flow as needed.
Args:
@@ -127,7 +139,7 @@ class BigQueryCredentialsManager:
# First, try to get credentials from the tool context
creds_json = tool_context.state.get(BIGQUERY_TOKEN_CACHE_KEY, None)
creds = (
Credentials.from_authorized_user_info(
google.oauth2.credentials.Credentials.from_authorized_user_info(
json.loads(creds_json), self.credentials_config.scopes
)
if creds_json
@@ -138,6 +150,11 @@ class BigQueryCredentialsManager:
if not creds:
creds = self.credentials_config.credentials
# If non-oauth credentials are provided then use them as is. This helps
# in flows such as service account keys
if creds and not isinstance(creds, google.oauth2.credentials.Credentials):
return creds
# Check if we have valid credentials
if creds and creds.valid:
return creds
@@ -159,7 +176,7 @@ class BigQueryCredentialsManager:
async def _perform_oauth_flow(
self, tool_context: ToolContext
) -> Optional[Credentials]:
) -> Optional[google.oauth2.credentials.Credentials]:
"""Perform OAuth flow to get new credentials.
Args:
@@ -199,7 +216,7 @@ class BigQueryCredentialsManager:
if auth_response:
# OAuth flow completed, create credentials
creds = Credentials(
creds = google.oauth2.credentials.Credentials(
token=auth_response.oauth2.access_token,
refresh_token=auth_response.oauth2.refresh_token,
token_uri=auth_scheme.flows.authorizationCode.tokenUrl,
@@ -19,7 +19,7 @@ from typing import Any
from typing import Callable
from typing import Optional
from google.oauth2.credentials import Credentials
from google.auth.credentials import Credentials
from typing_extensions import override
from ...utils.feature_decorator import experimental
+1 -1
View File
@@ -15,8 +15,8 @@
from __future__ import annotations
import google.api_core.client_info
from google.auth.credentials import Credentials
from google.cloud import bigquery
from google.oauth2.credentials import Credentials
from ... import version
@@ -12,8 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from google.auth.credentials import Credentials
from google.cloud import bigquery
from google.oauth2.credentials import Credentials
from . import client
+1 -1
View File
@@ -16,8 +16,8 @@ import functools
import types
from typing import Callable
from google.auth.credentials import Credentials
from google.cloud import bigquery
from google.oauth2.credentials import Credentials
from . import client
from .config import BigQueryToolConfig
+1 -1
View File
@@ -13,4 +13,4 @@
# limitations under the License.
# version: major.minor.patch
__version__ = "1.3.0"
__version__ = "1.4.2"
File diff suppressed because it is too large Load Diff
@@ -13,18 +13,43 @@
# limitations under the License.
import json
import sys
from unittest.mock import Mock
from unittest.mock import patch
from a2a import types as a2a_types
from google.adk.a2a.converters.part_converter import A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL
from google.adk.a2a.converters.part_converter import A2A_DATA_PART_METADATA_TYPE_FUNCTION_RESPONSE
from google.adk.a2a.converters.part_converter import A2A_DATA_PART_METADATA_TYPE_KEY
from google.adk.a2a.converters.part_converter import convert_a2a_part_to_genai_part
from google.adk.a2a.converters.part_converter import convert_genai_part_to_a2a_part
from google.genai import types as genai_types
import pytest
# Skip all tests in this module if Python version is less than 3.10
pytestmark = pytest.mark.skipif(
sys.version_info < (3, 10), reason="A2A tool requires Python 3.10+"
)
# Import dependencies with version checking
try:
from a2a import types as a2a_types
from google.adk.a2a.converters.part_converter import A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL
from google.adk.a2a.converters.part_converter import A2A_DATA_PART_METADATA_TYPE_FUNCTION_RESPONSE
from google.adk.a2a.converters.part_converter import A2A_DATA_PART_METADATA_TYPE_KEY
from google.adk.a2a.converters.part_converter import convert_a2a_part_to_genai_part
from google.adk.a2a.converters.part_converter import convert_genai_part_to_a2a_part
from google.genai import types as genai_types
except ImportError as e:
if sys.version_info < (3, 10):
# Create dummy classes to prevent NameError during test collection
# Tests will be skipped anyway due to pytestmark
class DummyTypes:
pass
a2a_types = DummyTypes()
genai_types = DummyTypes()
A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL = "function_call"
A2A_DATA_PART_METADATA_TYPE_FUNCTION_RESPONSE = "function_response"
A2A_DATA_PART_METADATA_TYPE_KEY = "type"
convert_a2a_part_to_genai_part = lambda x: None
convert_genai_part_to_a2a_part = lambda x: None
else:
raise e
class TestConvertA2aPartToGenaiPart:
"""Test cases for convert_a2a_part_to_genai_part function."""
@@ -100,7 +125,8 @@ class TestConvertA2aPartToGenaiPart:
metadata={
A2A_DATA_PART_METADATA_TYPE_KEY: (
A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL
)
),
"adk_type": A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL,
},
)
)
@@ -128,7 +154,8 @@ class TestConvertA2aPartToGenaiPart:
metadata={
A2A_DATA_PART_METADATA_TYPE_KEY: (
A2A_DATA_PART_METADATA_TYPE_FUNCTION_RESPONSE
)
),
"adk_type": A2A_DATA_PART_METADATA_TYPE_FUNCTION_RESPONSE,
},
)
)
@@ -12,11 +12,12 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from unittest.mock import Mock
from unittest import mock
from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsConfig
# Mock the Google OAuth and API dependencies
from google.oauth2.credentials import Credentials
import google.auth.credentials
import google.oauth2.credentials
import pytest
@@ -27,22 +28,46 @@ class TestBigQueryCredentials:
either existing credentials or client ID/secret pairs are provided.
"""
def test_valid_credentials_object(self):
"""Test that providing valid Credentials object works correctly.
def test_valid_credentials_object_auth_credentials(self):
"""Test that providing valid Credentials object works correctly with
google.auth.credentials.Credentials.
When a user already has valid OAuth credentials, they should be able
to pass them directly without needing to provide client ID/secret.
"""
# Create a mock credentials object with the expected attributes
mock_creds = Mock(spec=Credentials)
mock_creds.client_id = "test_client_id"
mock_creds.client_secret = "test_client_secret"
mock_creds.scopes = ["https://www.googleapis.com/auth/calendar"]
# Create a mock auth credentials object
# auth_creds = google.auth.credentials.Credentials()
auth_creds = mock.create_autospec(
google.auth.credentials.Credentials, instance=True
)
config = BigQueryCredentialsConfig(credentials=mock_creds)
config = BigQueryCredentialsConfig(credentials=auth_creds)
# Verify that the credentials are properly stored and attributes are extracted
assert config.credentials == mock_creds
assert config.credentials == auth_creds
assert config.client_id is None
assert config.client_secret is None
assert config.scopes == ["https://www.googleapis.com/auth/bigquery"]
def test_valid_credentials_object_oauth2_credentials(self):
"""Test that providing valid Credentials object works correctly with
google.oauth2.credentials.Credentials.
When a user already has valid OAuth credentials, they should be able
to pass them directly without needing to provide client ID/secret.
"""
# Create a mock oauth2 credentials object
oauth2_creds = google.oauth2.credentials.Credentials(
"test_token",
client_id="test_client_id",
client_secret="test_client_secret",
scopes=["https://www.googleapis.com/auth/calendar"],
)
config = BigQueryCredentialsConfig(credentials=oauth2_creds)
# Verify that the credentials are properly stored and attributes are extracted
assert config.credentials == oauth2_creds
assert config.client_id == "test_client_id"
assert config.client_secret == "test_client_secret"
assert config.scopes == ["https://www.googleapis.com/auth/calendar"]
@@ -22,9 +22,10 @@ from google.adk.tools import ToolContext
from google.adk.tools.bigquery.bigquery_credentials import BIGQUERY_TOKEN_CACHE_KEY
from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsConfig
from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsManager
from google.auth.credentials import Credentials as AuthCredentials
from google.auth.exceptions import RefreshError
# Mock the Google OAuth and API dependencies
from google.oauth2.credentials import Credentials
from google.oauth2.credentials import Credentials as OAuthCredentials
import pytest
@@ -64,9 +65,16 @@ class TestBigQueryCredentialsManager:
"""Create a credentials manager instance for testing."""
return BigQueryCredentialsManager(credentials_config)
@pytest.mark.parametrize(
("credentials_class",),
[
pytest.param(OAuthCredentials, id="oauth"),
pytest.param(AuthCredentials, id="auth"),
],
)
@pytest.mark.asyncio
async def test_get_valid_credentials_with_valid_existing_creds(
self, manager, mock_tool_context
self, manager, mock_tool_context, credentials_class
):
"""Test that valid existing credentials are returned immediately.
@@ -74,7 +82,7 @@ class TestBigQueryCredentialsManager:
should be needed. This is the optimal happy path scenario.
"""
# Create mock credentials that are already valid
mock_creds = Mock(spec=Credentials)
mock_creds = Mock(spec=credentials_class)
mock_creds.valid = True
manager.credentials_config.credentials = mock_creds
@@ -85,6 +93,34 @@ class TestBigQueryCredentialsManager:
mock_tool_context.get_auth_response.assert_not_called()
mock_tool_context.request_credential.assert_not_called()
@pytest.mark.parametrize(
("valid",),
[
pytest.param(False, id="invalid"),
pytest.param(True, id="valid"),
],
)
@pytest.mark.asyncio
async def test_get_valid_credentials_with_existing_non_oauth_creds(
self, manager, mock_tool_context, valid
):
"""Test that existing non-oauth credentials are returned immediately.
When credentials are of non-oauth type, no refresh or OAuth flow
is triggered irrespective of whether it is valid or not.
"""
# Create mock credentials that are already valid
mock_creds = Mock(spec=AuthCredentials)
mock_creds.valid = valid
manager.credentials_config.credentials = mock_creds
result = await manager.get_valid_credentials(mock_tool_context)
assert result == mock_creds
# Verify no OAuth flow was triggered
mock_tool_context.get_auth_response.assert_not_called()
mock_tool_context.request_credential.assert_not_called()
@pytest.mark.asyncio
async def test_get_credentials_from_cache_when_none_in_manager(
self, manager, mock_tool_context
@@ -113,7 +149,7 @@ class TestBigQueryCredentialsManager:
with patch(
"google.oauth2.credentials.Credentials.from_authorized_user_info"
) as mock_from_json:
mock_creds = Mock(spec=Credentials)
mock_creds = Mock(spec=OAuthCredentials)
mock_creds.valid = True
mock_from_json.return_value = mock_creds
@@ -179,7 +215,7 @@ class TestBigQueryCredentialsManager:
mock_tool_context.state[BIGQUERY_TOKEN_CACHE_KEY] = mock_cached_creds_json
# Create expired cached credentials with refresh token
mock_cached_creds = Mock(spec=Credentials)
mock_cached_creds = Mock(spec=OAuthCredentials)
mock_cached_creds.valid = False
mock_cached_creds.expired = True
mock_cached_creds.refresh_token = "valid_refresh_token"
@@ -227,7 +263,7 @@ class TestBigQueryCredentialsManager:
users from having to re-authenticate for every expired token.
"""
# Create expired credentials with refresh token
mock_creds = Mock(spec=Credentials)
mock_creds = Mock(spec=OAuthCredentials)
mock_creds.valid = False
mock_creds.expired = True
mock_creds.refresh_token = "refresh_token"
@@ -257,7 +293,7 @@ class TestBigQueryCredentialsManager:
gracefully fall back to requesting a new OAuth flow.
"""
# Create expired credentials that fail to refresh
mock_creds = Mock(spec=Credentials)
mock_creds = Mock(spec=OAuthCredentials)
mock_creds.valid = False
mock_creds.expired = True
mock_creds.refresh_token = "expired_refresh_token"
@@ -287,7 +323,7 @@ class TestBigQueryCredentialsManager:
mock_tool_context.get_auth_response.return_value = mock_auth_response
# Create a mock credentials instance that will represent our created credentials
mock_creds = Mock(spec=Credentials)
mock_creds = Mock(spec=OAuthCredentials)
# Make the JSON match what a real Credentials object would produce
mock_creds_json = (
'{"token": "new_access_token", "refresh_token": "new_refresh_token",'
@@ -300,7 +336,7 @@ class TestBigQueryCredentialsManager:
# Use the full module path as it appears in the project structure
with patch(
"google.adk.tools.bigquery.bigquery_credentials.Credentials",
"google.adk.tools.bigquery.bigquery_credentials.google.oauth2.credentials.Credentials",
return_value=mock_creds,
) as mock_credentials_class:
result = await manager.get_valid_credentials(mock_tool_context)
@@ -361,7 +397,7 @@ class TestBigQueryCredentialsManager:
mock_tool_context.get_auth_response.return_value = mock_auth_response
# Create the mock credentials instance that will be returned by the constructor
mock_creds = Mock(spec=Credentials)
mock_creds = Mock(spec=OAuthCredentials)
# Make sure our mock JSON matches the structure that real Credentials objects produce
mock_creds_json = (
'{"token": "cached_access_token", "refresh_token":'
@@ -376,7 +412,7 @@ class TestBigQueryCredentialsManager:
# Use the correct module path - without the 'src.' prefix
with patch(
"google.adk.tools.bigquery.bigquery_credentials.Credentials",
"google.adk.tools.bigquery.bigquery_credentials.google.oauth2.credentials.Credentials",
return_value=mock_creds,
) as mock_credentials_class:
# Complete OAuth flow with first manager
@@ -396,9 +432,9 @@ class TestBigQueryCredentialsManager:
# Mock the from_authorized_user_info method for the second manager
with patch(
"google.adk.tools.bigquery.bigquery_credentials.Credentials.from_authorized_user_info"
"google.adk.tools.bigquery.bigquery_credentials.google.oauth2.credentials.Credentials.from_authorized_user_info"
) as mock_from_json:
mock_cached_creds = Mock(spec=Credentials)
mock_cached_creds = Mock(spec=OAuthCredentials)
mock_cached_creds.valid = True
mock_from_json.return_value = mock_cached_creds
@@ -20,13 +20,35 @@ from unittest.mock import AsyncMock
from unittest.mock import Mock
from unittest.mock import patch
from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager
from google.adk.tools.mcp_tool.mcp_session_manager import retry_on_closed_resource
from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams
from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
import pytest
# Skip all tests in this module if Python version is less than 3.10
pytestmark = pytest.mark.skipif(
sys.version_info < (3, 10), reason="MCP tool requires Python 3.10+"
)
# Import dependencies with version checking
try:
from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager
from google.adk.tools.mcp_tool.mcp_session_manager import retry_on_closed_resource
from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams
from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
except ImportError as e:
if sys.version_info < (3, 10):
# Create dummy classes to prevent NameError during test collection
# Tests will be skipped anyway due to pytestmark
class DummyClass:
pass
MCPSessionManager = DummyClass
retry_on_closed_resource = lambda x: x
SseConnectionParams = DummyClass
StdioConnectionParams = DummyClass
StreamableHTTPConnectionParams = DummyClass
else:
raise e
# Import real MCP classes
try:
from mcp import StdioServerParameters
@@ -12,7 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import sys
from typing import Any
from typing import Dict
from unittest.mock import AsyncMock
from unittest.mock import Mock
from unittest.mock import patch
@@ -23,14 +25,33 @@ from google.adk.auth.auth_credential import HttpAuth
from google.adk.auth.auth_credential import HttpCredentials
from google.adk.auth.auth_credential import OAuth2Auth
from google.adk.auth.auth_credential import ServiceAccount
from google.adk.auth.auth_schemes import AuthScheme
from google.adk.auth.auth_schemes import AuthSchemeType
from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager
from google.adk.tools.mcp_tool.mcp_tool import MCPTool
from google.adk.tools.tool_context import ToolContext
from google.genai.types import FunctionDeclaration
import pytest
# Skip all tests in this module if Python version is less than 3.10
pytestmark = pytest.mark.skipif(
sys.version_info < (3, 10), reason="MCP tool requires Python 3.10+"
)
# Import dependencies with version checking
try:
from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager
from google.adk.tools.mcp_tool.mcp_tool import MCPTool
from google.adk.tools.tool_context import ToolContext
from google.genai.types import FunctionDeclaration
except ImportError as e:
if sys.version_info < (3, 10):
# Create dummy classes to prevent NameError during test collection
# Tests will be skipped anyway due to pytestmark
class DummyClass:
pass
MCPSessionManager = DummyClass
MCPTool = DummyClass
ToolContext = DummyClass
FunctionDeclaration = DummyClass
else:
raise e
# Mock MCP Tool from mcp.types
class MockMCPTool:

Some files were not shown because too many files have changed in this diff Show More