You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
Compare commits
7 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 8677d5c8dc | |||
| 4d72d31b13 | |||
| 742478fdb7 | |||
| ffcba70686 | |||
| 2f716ada7f | |||
| 17beb32880 | |||
| 7f8dc8927a |
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
Reference in New Issue
Block a user