diff --git a/contributing/samples/adk_pr_triaging_agent/agent.py b/contributing/samples/adk_pr_triaging_agent/agent.py index e7718e50..c5fe0e0c 100644 --- a/contributing/samples/adk_pr_triaging_agent/agent.py +++ b/contributing/samples/adk_pr_triaging_agent/agent.py @@ -69,8 +69,10 @@ def get_pull_request_details(pr_number: int) -> str: repository(owner: $owner, name: $repo) { pullRequest(number: $prNumber) { id + number title body + state author { login } diff --git a/contributing/samples/adk_pr_triaging_agent/main.py b/contributing/samples/adk_pr_triaging_agent/main.py index da67fa16..ad5893d8 100644 --- a/contributing/samples/adk_pr_triaging_agent/main.py +++ b/contributing/samples/adk_pr_triaging_agent/main.py @@ -13,6 +13,7 @@ # limitations under the License. import asyncio +import logging import time from adk_pr_triaging_agent import agent @@ -21,11 +22,14 @@ from adk_pr_triaging_agent.settings import PULL_REQUEST_NUMBER from adk_pr_triaging_agent.settings import REPO from adk_pr_triaging_agent.utils import call_agent_async from adk_pr_triaging_agent.utils import parse_number_string +from google.adk.cli.utils import logs from google.adk.runners import InMemoryRunner APP_NAME = "adk_pr_triaging_app" USER_ID = "adk_pr_triaging_user" +logs.setup_adk_logger(level=logging.DEBUG) + async def main(): runner = InMemoryRunner( diff --git a/contributing/samples/built_in_multi_tools/agent.py b/contributing/samples/built_in_multi_tools/agent.py index 5194e0ef..3eb9ce8b 100644 --- a/contributing/samples/built_in_multi_tools/agent.py +++ b/contributing/samples/built_in_multi_tools/agent.py @@ -17,7 +17,7 @@ import random from dotenv import load_dotenv from google.adk import Agent -from google.adk.tools.google_search_tool import google_search +from google.adk.tools.google_search_tool import GoogleSearchTool from google.adk.tools.tool_context import ToolContext from google.adk.tools.vertex_ai_search_tool import VertexAiSearchTool @@ -57,7 +57,9 @@ root_agent = Agent( """, tools=[ roll_die, - VertexAiSearchTool(data_store_id=VERTEXAI_DATASTORE_ID), - google_search, + VertexAiSearchTool( + data_store_id=VERTEXAI_DATASTORE_ID, bypass_multi_tools_limit=True + ), + GoogleSearchTool(bypass_multi_tools_limit=True), ], ) diff --git a/contributing/samples/computer_use/README.md b/contributing/samples/computer_use/README.md new file mode 100644 index 00000000..ff7ae6c0 --- /dev/null +++ b/contributing/samples/computer_use/README.md @@ -0,0 +1,96 @@ +# Computer Use Agent + +This directory contains a computer use agent that can operate a browser to complete user tasks. The agent uses Playwright to control a Chromium browser and can interact with web pages by taking screenshots, clicking, typing, and navigating. + +This agent is to demo the usage of ComputerUseToolset. + + +## Overview + +The computer use agent consists of: +- `agent.py`: Main agent configuration using Google's gemini-2.5-computer-use-preview-10-2025 model +- `playwright.py`: Playwright-based computer implementation for browser automation +- `requirements.txt`: Python dependencies + +## Setup + +### 1. Install Python Dependencies + +Install the required Python packages from the requirements file: + +```bash +uv pip install -r internal/samples/computer_use/requirements.txt +``` + +### 2. Install Playwright Dependencies + +Install Playwright's system dependencies for Chromium: + +```bash +playwright install-deps chromium +``` + +### 3. Install Chromium Browser + +Install the Chromium browser for Playwright: + +```bash +playwright install chromium +``` + +## Usage + +### Running the Agent + +To start the computer use agent, run the following command from the project root: + +```bash +adk web internal/samples +``` + +This will start the ADK web interface where you can interact with the computer_use agent. + +### Example Queries + +Once the agent is running, you can send queries like: + +``` +find me a flight from SF to Hawaii on next Monday, coming back on next Friday. start by navigating directly to flights.google.com +``` + +The agent will: +1. Open a browser window +2. Navigate to the specified website +3. Interact with the page elements to complete your task +4. Provide updates on its progress + +### Other Example Tasks + +- Book hotel reservations +- Search for products online +- Fill out forms +- Navigate complex websites +- Research information across multiple pages + +## Technical Details + +- **Model**: Uses Google's `gemini-2.5-computer-use-preview-10-2025` model for computer use capabilities +- **Browser**: Automated Chromium browser via Playwright +- **Screen Size**: Configured for 600x800 resolution +- **Tools**: Uses ComputerUseToolset for screen capture, clicking, typing, and scrolling + +## Troubleshooting + +If you encounter issues: + +1. **Playwright not found**: Make sure you've run both `playwright install-deps chromium` and `playwright install chromium` +2. **Dependencies missing**: Verify all packages from `requirements.txt` are installed +3. **Browser crashes**: Check that your system supports Chromium and has sufficient resources +4. **Permission errors**: Ensure your user has permission to run browser automation tools + +## Notes + +- The agent operates in a controlled browser environment +- Screenshots are taken to help the agent understand the current state +- The agent will provide updates on its actions as it works +- Be patient as complex tasks may take some time to complete diff --git a/contributing/samples/computer_use/agent.py b/contributing/samples/computer_use/agent.py new file mode 100755 index 00000000..5fed9aa2 --- /dev/null +++ b/contributing/samples/computer_use/agent.py @@ -0,0 +1,35 @@ +# 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 google.adk import Agent +from google.adk.models.google_llm import Gemini +from google.adk.tools.computer_use.computer_use_toolset import ComputerUseToolset +from typing_extensions import override + +from .playwright import PlaywrightComputer + +root_agent = Agent( + model='gemini-2.5-computer-use-preview-10-2025', + name='hello_world_agent', + description=( + 'computer use agent that can operate a browser on a computer to finish' + ' user tasks' + ), + instruction=""" + you are a computer use agent + """, + tools=[ + ComputerUseToolset(computer=PlaywrightComputer(screen_size=(1280, 936))) + ], +) diff --git a/contributing/samples/computer_use/playwright.py b/contributing/samples/computer_use/playwright.py new file mode 100644 index 00000000..64ad54fd --- /dev/null +++ b/contributing/samples/computer_use/playwright.py @@ -0,0 +1,317 @@ +# 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. +import asyncio +import time +from typing import Literal + +from google.adk.tools.computer_use.base_computer import BaseComputer +from google.adk.tools.computer_use.base_computer import ComputerEnvironment +from google.adk.tools.computer_use.base_computer import ComputerState +from playwright.async_api import async_playwright +import termcolor +from typing_extensions import override + +# Define a mapping from the user-friendly key names to Playwright's expected key names. +# Playwright is generally good with case-insensitivity for these, but it's best to be canonical. +# See: https://playwright.dev/docs/api/class-keyboard#keyboard-press +# Keys like 'a', 'b', '1', '$' are passed directly. +PLAYWRIGHT_KEY_MAP = { + "backspace": "Backspace", + "tab": "Tab", + "return": "Enter", # Playwright uses 'Enter' + "enter": "Enter", + "shift": "Shift", + "control": "Control", # Or 'ControlOrMeta' for cross-platform Ctrl/Cmd + "alt": "Alt", + "escape": "Escape", + "space": "Space", # Can also just be " " + "pageup": "PageUp", + "pagedown": "PageDown", + "end": "End", + "home": "Home", + "left": "ArrowLeft", + "up": "ArrowUp", + "right": "ArrowRight", + "down": "ArrowDown", + "insert": "Insert", + "delete": "Delete", + "semicolon": ";", # For actual character ';' + "equals": "=", # For actual character '=' + "multiply": "Multiply", # NumpadMultiply + "add": "Add", # NumpadAdd + "separator": "Separator", # Numpad specific + "subtract": "Subtract", # NumpadSubtract, or just '-' for character + "decimal": "Decimal", # NumpadDecimal, or just '.' for character + "divide": "Divide", # NumpadDivide, or just '/' for character + "f1": "F1", + "f2": "F2", + "f3": "F3", + "f4": "F4", + "f5": "F5", + "f6": "F6", + "f7": "F7", + "f8": "F8", + "f9": "F9", + "f10": "F10", + "f11": "F11", + "f12": "F12", + "command": "Meta", # 'Meta' is Command on macOS, Windows key on Windows +} + + +class PlaywrightComputer(BaseComputer): + """Conputer that controls Chromium via Playwright.""" + + def __init__( + self, + screen_size: tuple[int, int], + initial_url: str = "https://www.google.com", + search_engine_url: str = "https://www.google.com", + highlight_mouse: bool = False, + ): + self._initial_url = initial_url + self._screen_size = screen_size + self._search_engine_url = search_engine_url + self._highlight_mouse = highlight_mouse + + @override + async def initialize(self): + print("Creating session...") + self._playwright = await async_playwright().start() + self._browser = await self._playwright.chromium.launch( + args=["--disable-blink-features=AutomationControlled"], + headless=False, + ) + self._context = await self._browser.new_context( + viewport={ + "width": self._screen_size[0], + "height": self._screen_size[1], + } + ) + self._page = await self._context.new_page() + await self._page.goto(self._initial_url) + + termcolor.cprint( + f"Started local playwright.", + color="green", + attrs=["bold"], + ) + + @override + async def environment(self): + return ComputerEnvironment.ENVIRONMENT_BROWSER + + @override + async def close(self, exc_type, exc_val, exc_tb): + if self._context: + self._context.close() + try: + self._browser.close() + except Exception as e: + # Browser was already shut down because of SIGINT or such. + if ( + "Browser.close: Connection closed while reading from the driver" + in str(e) + ): + pass + else: + raise + + self._playwright.stop() + + async def open_web_browser(self) -> ComputerState: + return await self.current_state() + + async def click_at(self, x: int, y: int): + await self.highlight_mouse(x, y) + await self._page.mouse.click(x, y) + await self._page.wait_for_load_state() + return await self.current_state() + + async def hover_at(self, x: int, y: int): + await self.highlight_mouse(x, y) + await self._page.mouse.move(x, y) + await self._page.wait_for_load_state() + return await self.current_state() + + async def type_text_at( + self, + x: int, + y: int, + text: str, + press_enter: bool = True, + clear_before_typing: bool = True, + ) -> ComputerState: + await self.highlight_mouse(x, y) + await self._page.mouse.click(x, y) + await self._page.wait_for_load_state() + + if clear_before_typing: + await self.key_combination(["Control", "A"]) + await self.key_combination(["Delete"]) + + await self._page.keyboard.type(text) + await self._page.wait_for_load_state() + + if press_enter: + await self.key_combination(["Enter"]) + await self._page.wait_for_load_state() + return await self.current_state() + + async def _horizontal_document_scroll( + self, direction: Literal["left", "right"] + ) -> ComputerState: + # Scroll by 50% of the viewport size. + horizontal_scroll_amount = await self.screen_size()[0] // 2 + if direction == "left": + sign = "-" + else: + sign = "" + scroll_argument = f"{sign}{horizontal_scroll_amount}" + # Scroll using JS. + await self._page.evaluate(f"window.scrollBy({scroll_argument}, 0); ") + await self._page.wait_for_load_state() + return await self.current_state() + + async def scroll_document( + self, direction: Literal["up", "down", "left", "right"] + ) -> ComputerState: + if direction == "down": + return await self.key_combination(["PageDown"]) + elif direction == "up": + return await self.key_combination(["PageUp"]) + elif direction in ("left", "right"): + return await self._horizontal_document_scroll(direction) + else: + raise ValueError("Unsupported direction: ", direction) + + async def scroll_at( + self, + x: int, + y: int, + direction: Literal["up", "down", "left", "right"], + magnitude: int, + ) -> ComputerState: + await self.highlight_mouse(x, y) + + await self._page.mouse.move(x, y) + await self._page.wait_for_load_state() + + dx = 0 + dy = 0 + if direction == "up": + dy = -magnitude + elif direction == "down": + dy = magnitude + elif direction == "left": + dx = -magnitude + elif direction == "right": + dx = magnitude + else: + raise ValueError("Unsupported direction: ", direction) + + await self._page.mouse.wheel(dx, dy) + await self._page.wait_for_load_state() + return await self.current_state() + + async def wait(self, seconds: int) -> ComputerState: + await asyncio.sleep(seconds) + return await self.current_state() + + async def go_back(self) -> ComputerState: + await self._page.go_back() + await self._page.wait_for_load_state() + return await self.current_state() + + async def go_forward(self) -> ComputerState: + await self._page.go_forward() + await self._page.wait_for_load_state() + return await self.current_state() + + async def search(self) -> ComputerState: + return await self.navigate(self._search_engine_url) + + async def navigate(self, url: str) -> ComputerState: + await self._page.goto(url) + await self._page.wait_for_load_state() + return await self.current_state() + + async def key_combination(self, keys: list[str]) -> ComputerState: + # Normalize all keys to the Playwright compatible version. + keys = [PLAYWRIGHT_KEY_MAP.get(k.lower(), k) for k in keys] + + for key in keys[:-1]: + await self._page.keyboard.down(key) + + await self._page.keyboard.press(keys[-1]) + + for key in reversed(keys[:-1]): + await self._page.keyboard.up(key) + + return await self.current_state() + + async def drag_and_drop( + self, x: int, y: int, destination_x: int, destination_y: int + ) -> ComputerState: + await self.highlight_mouse(x, y) + await self._page.mouse.move(x, y) + await self._page.wait_for_load_state() + await self._page.mouse.down() + await self._page.wait_for_load_state() + + await self.highlight_mouse(destination_x, destination_y) + await self._page.mouse.move(destination_x, destination_y) + await self._page.wait_for_load_state() + await self._page.mouse.up() + return await self.current_state() + + async def current_state(self) -> ComputerState: + await self._page.wait_for_load_state() + # Even if Playwright reports the page as loaded, it may not be so. + # Add a manual sleep to make sure the page has finished rendering. + time.sleep(0.5) + screenshot_bytes = await self._page.screenshot(type="png", full_page=False) + return ComputerState(screenshot=screenshot_bytes, url=self._page.url) + + async def screen_size(self) -> tuple[int, int]: + return self._screen_size + + async def highlight_mouse(self, x: int, y: int): + if not self._highlight_mouse: + return + await self._page.evaluate(f""" + () => {{ + const element_id = "playwright-feedback-circle"; + const div = document.createElement('div'); + div.id = element_id; + div.style.pointerEvents = 'none'; + div.style.border = '4px solid red'; + div.style.borderRadius = '50%'; + div.style.width = '20px'; + div.style.height = '20px'; + div.style.position = 'fixed'; + div.style.zIndex = '9999'; + document.body.appendChild(div); + + div.hidden = false; + div.style.left = {x} - 10 + 'px'; + div.style.top = {y} - 10 + 'px'; + + setTimeout(() => {{ + div.hidden = true; + }}, 2000); + }} + """) + # Wait a bit for the user to see the cursor. + time.sleep(1) diff --git a/contributing/samples/computer_use/requirements.txt b/contributing/samples/computer_use/requirements.txt new file mode 100644 index 00000000..5b1df138 --- /dev/null +++ b/contributing/samples/computer_use/requirements.txt @@ -0,0 +1,4 @@ +termcolor==3.1.0 +playwright==1.52.0 +browserbase==1.3.0 +rich diff --git a/contributing/samples/mcp_dynamic_header_agent/README.md b/contributing/samples/mcp_dynamic_header_agent/README.md new file mode 100644 index 00000000..50f71255 --- /dev/null +++ b/contributing/samples/mcp_dynamic_header_agent/README.md @@ -0,0 +1,8 @@ +This agent connects to a local MCP server via Streamable HTTP and provides +custom per-request headers to the MCP server. + +To run this agent, start the local MCP server first by running: + +```bash +uv run header_server.py +``` \ No newline at end of file diff --git a/contributing/samples/mcp_dynamic_header_agent/__init__.py b/contributing/samples/mcp_dynamic_header_agent/__init__.py new file mode 100644 index 00000000..c48963cd --- /dev/null +++ b/contributing/samples/mcp_dynamic_header_agent/__init__.py @@ -0,0 +1,15 @@ +# 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 . import agent diff --git a/contributing/samples/mcp_dynamic_header_agent/agent.py b/contributing/samples/mcp_dynamic_header_agent/agent.py new file mode 100644 index 00000000..028d7feb --- /dev/null +++ b/contributing/samples/mcp_dynamic_header_agent/agent.py @@ -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 google.adk.agents.llm_agent import LlmAgent +from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams +from google.adk.tools.mcp_tool.mcp_toolset import McpToolset + +root_agent = LlmAgent( + model='gemini-2.0-flash', + name='tenant_agent', + instruction="""You are a helpful assistant that helps users get tenant + information. Call the get_tenant_data tool when the user asks for tenant data.""", + tools=[ + McpToolset( + connection_params=StreamableHTTPConnectionParams( + url='http://localhost:3000/mcp', + ), + tool_filter=['get_tenant_data'], + header_provider=lambda ctx: {'X-Tenant-ID': 'tenant1'}, + ) + ], +) diff --git a/contributing/samples/mcp_dynamic_header_agent/header_server.py b/contributing/samples/mcp_dynamic_header_agent/header_server.py new file mode 100644 index 00000000..386ae43b --- /dev/null +++ b/contributing/samples/mcp_dynamic_header_agent/header_server.py @@ -0,0 +1,50 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from fastapi import Request +from mcp.server.fastmcp import Context +from mcp.server.fastmcp import FastMCP + +mcp = FastMCP('Header Check Server', host='localhost', port=3000) + +TENANT_DATA = { + 'tenant1': {'name': 'Tenant 1', 'data': 'Data for tenant 1'}, + 'tenant2': {'name': 'Tenant 2', 'data': 'Data for tenant 2'}, +} + + +@mcp.tool( + description='Returns tenant specific data based on X-Tenant-ID header.' +) +def get_tenant_data(context: Context) -> dict: + """Return tenant specific data.""" + if context.request_context and context.request_context.request: + headers = context.request_context.request.headers + tenant_id = headers.get('x-tenant-id') + if tenant_id in TENANT_DATA: + return TENANT_DATA[tenant_id] + else: + return {'error': f'Tenant {tenant_id} not found'} + else: + return {'error': 'Could not get request context'} + + +if __name__ == '__main__': + try: + print('Starting Header Check MCP server on http://localhost:3000') + mcp.run(transport='streamable-http') + except KeyboardInterrupt: + print('\nServer stopped.') diff --git a/contributing/samples/plugin_reflect_tool_retry/README.md b/contributing/samples/plugin_reflect_tool_retry/README.md index 084862e4..77362316 100644 --- a/contributing/samples/plugin_reflect_tool_retry/README.md +++ b/contributing/samples/plugin_reflect_tool_retry/README.md @@ -46,8 +46,30 @@ You can run the agent with: $ adk web contributing/samples/plugin_reflect_tool_retry ``` -You can provide the following prompt to see the agent retrying tool calls: +Select "basic" and provide the following prompt to see the agent retrying tool +calls: ``` Please guess a number! Tell me what number you guess and how is it. ``` + +### Hallucinating tool calls + +The "hallucinating_func_name" agent is an example to show the plugin can retry +hallucinating tool calls. + +For example, we used the `after_model_callback` to hack a tool call with the +wrong name then the agent can retry calling with the right tool name. + +You can run the agent with: + +```bash +$ adk web contributing/samples/plugin_reflect_tool_retry +``` + +Select "hallucinating_func_name" and provide the following prompt to see the +agent retrying tool calls: + +``` +Roll a 6 sided die +``` diff --git a/contributing/samples/plugin_reflect_tool_retry/hallucinating_func_name/__init__.py b/contributing/samples/plugin_reflect_tool_retry/hallucinating_func_name/__init__.py new file mode 100644 index 00000000..c48963cd --- /dev/null +++ b/contributing/samples/plugin_reflect_tool_retry/hallucinating_func_name/__init__.py @@ -0,0 +1,15 @@ +# 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 . import agent diff --git a/contributing/samples/plugin_reflect_tool_retry/hallucinating_func_name/agent.py b/contributing/samples/plugin_reflect_tool_retry/hallucinating_func_name/agent.py new file mode 100644 index 00000000..5b8d9262 --- /dev/null +++ b/contributing/samples/plugin_reflect_tool_retry/hallucinating_func_name/agent.py @@ -0,0 +1,81 @@ +# 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. + +import random + +from google.adk.agents import LlmAgent +from google.adk.agents.callback_context import CallbackContext +from google.adk.apps.app import App +from google.adk.models.llm_response import LlmResponse +from google.adk.plugins import ReflectAndRetryToolPlugin +from google.adk.tools.tool_context import ToolContext + +APP_NAME = "hallucinating_func_name" +USER_ID = "test_user" + +hallucinated = False # Whether the tool name is hallucinated + + +def roll_die(sides: int, tool_context: ToolContext) -> int: + """Roll a die and return the rolled result. + + Args: + sides: The integer number of sides the die has. + + Returns: + An integer of the result of rolling the die. + """ + result = random.randint(1, sides) + if not "rolls" in tool_context.state: + tool_context.state["rolls"] = [] + + tool_context.state["rolls"] = tool_context.state["rolls"] + [result] + return result + + +def after_model_callback( + callback_context: CallbackContext, llm_response: LlmResponse +): + """After model callback to produce one hallucinating tool call.""" + global hallucinated + + if hallucinated: + return None + + if ( + llm_response.content + and llm_response.content.parts[0].function_call.name == "roll_die" + ): + llm_response.content.parts[0].function_call.name = "roll_die_wrong_name" + hallucinated = True + return None + + +root_agent = LlmAgent( + name="hello_world", + description="Helpful agent", + instruction="""Use guess_number_tool to guess a number.""", + model="gemini-2.5-flash", + tools=[roll_die], + after_model_callback=after_model_callback, +) + + +app = App( + name=APP_NAME, + root_agent=root_agent, + plugins=[ + ReflectAndRetryToolPlugin(max_retries=3), + ], +) diff --git a/contributing/samples/rewind_session/__init__.py b/contributing/samples/rewind_session/__init__.py new file mode 100644 index 00000000..c48963cd --- /dev/null +++ b/contributing/samples/rewind_session/__init__.py @@ -0,0 +1,15 @@ +# 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 . import agent diff --git a/contributing/samples/rewind_session/agent.py b/contributing/samples/rewind_session/agent.py new file mode 100644 index 00000000..569bde07 --- /dev/null +++ b/contributing/samples/rewind_session/agent.py @@ -0,0 +1,71 @@ +# 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 google.adk import Agent +from google.adk.tools.tool_context import ToolContext +from google.genai import types + + +async def update_state(tool_context: ToolContext, key: str, value: str) -> dict: + """Updates a state value.""" + tool_context.state[key] = value + return {"status": f"Updated state '{key}' to '{value}'"} + + +async def load_state(tool_context: ToolContext, key: str) -> dict: + """Loads a state value.""" + return {key: tool_context.state.get(key)} + + +async def save_artifact( + tool_context: ToolContext, filename: str, content: str +) -> dict: + """Saves an artifact with the given filename and content.""" + artifact_bytes = content.encode("utf-8") + artifact_part = types.Part( + inline_data=types.Blob(mime_type="text/plain", data=artifact_bytes) + ) + version = await tool_context.save_artifact(filename, artifact_part) + return {"status": "success", "filename": filename, "version": version} + + +async def load_artifact(tool_context: ToolContext, filename: str) -> dict: + """Loads an artifact with the given filename.""" + artifact = await tool_context.load_artifact(filename) + if not artifact: + return {"error": f"Artifact '{filename}' not found"} + content = artifact.inline_data.data.decode("utf-8") + return {"filename": filename, "content": content} + + +# Create the agent +root_agent = Agent( + name="state_agent", + model="gemini-2.0-flash", + instruction="""You are an agent that manages state and artifacts. + + You can: + - Update state value + - Load state value + - Save artifact + - Load artifact + + Use the appropriate tool based on what the user asks for.""", + tools=[ + update_state, + load_state, + save_artifact, + load_artifact, + ], +) diff --git a/llms-full.txt b/llms-full.txt index a053d5d7..35ffd56d 100644 --- a/llms-full.txt +++ b/llms-full.txt @@ -92,7 +92,7 @@ from google.adk.tools import google_search root_agent = Agent( name="search_assistant", - model="gemini-2.0-flash", # Or your preferred Gemini model + model="gemini-2.5-flash", # Or your preferred Gemini model instruction="You are a helpful assistant. Answer user questions using Google Search when needed.", description="An assistant that can search the web.", tools=[google_search] @@ -107,13 +107,13 @@ Define a multi-agent system with coordinator agent, greeter agent, and task exec from google.adk.agents import LlmAgent, BaseAgent # Define individual agents -greeter = LlmAgent(name="greeter", model="gemini-2.0-flash", ...) -task_executor = LlmAgent(name="task_executor", model="gemini-2.0-flash", ...) +greeter = LlmAgent(name="greeter", model="gemini-2.5-flash", ...) +task_executor = LlmAgent(name="task_executor", model="gemini-2.5-flash", ...) # Create parent agent and assign children via sub_agents coordinator = LlmAgent( name="Coordinator", - model="gemini-2.0-flash", + model="gemini-2.5-flash", description="I coordinate greetings and tasks.", sub_agents=[ # Assign sub_agents here greeter, @@ -432,7 +432,7 @@ These are standard `LlmAgent` definitions, responsible for specific tasks. Their === "Python" ```python - GEMINI_2_FLASH = "gemini-2.0-flash" # Define model constant + GEMINI_2_FLASH = "gemini-2.5-flash" # Define model constant # --- Define the individual LLM agents --- story_generator = LlmAgent( name="StoryGenerator", @@ -597,7 +597,7 @@ Finally, you instantiate your `StoryFlowAgent` and use the `Runner` as usual. APP_NAME = "story_app" USER_ID = "12345" SESSION_ID = "123344" - GEMINI_2_FLASH = "gemini-2.0-flash" + GEMINI_2_FLASH = "gemini-2.5-flash" # --- Configure Logging --- logging.basicConfig(level=logging.INFO) @@ -951,7 +951,7 @@ First, you need to establish what the agent *is* and what it's *for*. inquiries about current billing statements," not just "Billing agent"). * **`model` (Required):** Specify the underlying LLM that will power this - agent's reasoning. This is a string identifier like `"gemini-2.0-flash"`. The + agent's reasoning. This is a string identifier like `"gemini-2.5-flash"`. The choice of model impacts the agent's capabilities, cost, and performance. See the [Models](models.md) page for available options and considerations. @@ -960,7 +960,7 @@ First, you need to establish what the agent *is* and what it's *for*. ```python # Example: Defining the basic identity capital_agent = LlmAgent( - model="gemini-2.0-flash", + model="gemini-2.5-flash", name="capital_agent", description="Answers user questions about the capital city of a given country." # instruction and tools will be added next @@ -1003,7 +1003,7 @@ tells the agent: ```python # Example: Adding instructions capital_agent = LlmAgent( - model="gemini-2.0-flash", + model="gemini-2.5-flash", name="capital_agent", description="Answers user questions about the capital city of a given country.", instruction="""You are an agent that provides the capital city of a country. @@ -1053,7 +1053,7 @@ on the conversation and its instructions. # Add the tool to the agent capital_agent = LlmAgent( - model="gemini-2.0-flash", + model="gemini-2.5-flash", name="capital_agent", description="Answers user questions about the capital city of a given country.", instruction="""You are an agent that provides the capital city of a country... (previous instruction text)""", @@ -1185,7 +1185,7 @@ For more complex reasoning involving multiple steps or executing code: USER_ID = "test_user_456" SESSION_ID_TOOL_AGENT = "session_tool_agent_xyz" SESSION_ID_SCHEMA_AGENT = "session_schema_agent_xyz" - MODEL_NAME = "gemini-2.0-flash" + MODEL_NAME = "gemini-2.5-flash" # --- 2. Define Schemas --- @@ -1441,7 +1441,7 @@ export GOOGLE_GENAI_USE_VERTEXAI=FALSE # --- Example using a stable Gemini Flash model --- agent_gemini_flash = LlmAgent( # Use the latest stable Flash model identifier - model="gemini-2.0-flash", + model="gemini-2.5-flash", name="gemini_flash_agent", instruction="You are a fast and helpful Gemini assistant.", # ... other agent parameters @@ -1453,7 +1453,7 @@ export GOOGLE_GENAI_USE_VERTEXAI=FALSE # different availability or quota limitations. agent_gemini_pro = LlmAgent( # Use the latest generally available Pro model identifier - model="gemini-2.5-pro-preview-03-25", + model="gemini-2.5-pro", name="gemini_pro_agent", instruction="You are a powerful and knowledgeable Gemini assistant.", # ... other agent parameters @@ -2009,13 +2009,13 @@ The foundation for structuring multi-agent systems is the parent-child relations from google.adk.agents import LlmAgent, BaseAgent # Define individual agents - greeter = LlmAgent(name="Greeter", model="gemini-2.0-flash") + greeter = LlmAgent(name="Greeter", model="gemini-2.5-flash") task_doer = BaseAgent(name="TaskExecutor") # Custom non-LLM agent # Create parent agent and assign children via sub_agents coordinator = LlmAgent( name="Coordinator", - model="gemini-2.0-flash", + model="gemini-2.5-flash", description="I coordinate greetings and tasks.", sub_agents=[ # Assign sub_agents here greeter, @@ -2163,7 +2163,7 @@ Leverages an [`LlmAgent`](llm-agents.md)'s understanding to dynamically route ta coordinator = LlmAgent( name="Coordinator", - model="gemini-2.0-flash", + model="gemini-2.5-flash", instruction="You are an assistant. Delegate booking tasks to Booker and info requests to Info.", description="Main coordinator.", # AutoFlow is typically used implicitly here @@ -2212,7 +2212,7 @@ Allows an [`LlmAgent`](llm-agents.md) to treat another `BaseAgent` instance as a # Parent agent uses the AgentTool artist_agent = LlmAgent( name="Artist", - model="gemini-2.0-flash", + model="gemini-2.5-flash", instruction="Create a prompt and use the ImageGen tool to generate the image.", tools=[image_tool] # Include the AgentTool ) @@ -2251,7 +2251,7 @@ By combining ADK's composition primitives, you can implement various established coordinator = LlmAgent( name="HelpDeskCoordinator", - model="gemini-2.0-flash", + model="gemini-2.5-flash", instruction="Route user requests: Use Billing agent for payment issues, Support agent for technical problems.", description="Main help desk router.", # allow_transfer=True is often implicit with sub_agents in AutoFlow @@ -2357,7 +2357,7 @@ By combining ADK's composition primitives, you can implement various established # Mid-level agent combining tools research_assistant = LlmAgent( name="ResearchAssistant", - model="gemini-2.0-flash", + model="gemini-2.5-flash", description="Finds and summarizes information on a topic.", tools=[agent_tool.AgentTool(agent=web_searcher), agent_tool.AgentTool(agent=summarizer)] ) @@ -2365,7 +2365,7 @@ By combining ADK's composition primitives, you can implement various established # High-level agent delegating research report_writer = LlmAgent( name="ReportWriter", - model="gemini-2.0-flash", + model="gemini-2.5-flash", instruction="Write a report on topic X. Use the ResearchAssistant to gather information.", tools=[agent_tool.AgentTool(agent=research_assistant)] # Alternatively, could use LLM Transfer if research_assistant is a sub_agent @@ -2641,7 +2641,7 @@ In this setup, the `LoopAgent` would manage the iterative process. The `CriticA APP_NAME = "doc_writing_app_v3" # New App Name USER_ID = "dev_user_01" SESSION_ID_BASE = "loop_exit_tool_session" # New Base Session ID - GEMINI_MODEL = "gemini-2.0-flash" + GEMINI_MODEL = "gemini-2.5-flash" STATE_INITIAL_TOPIC = "initial_topic" # --- State Keys --- STATE_CURRENT_DOC = "current_document" @@ -3171,7 +3171,7 @@ Understanding artifacts involves grasping a few key components: the service that from google.adk.sessions import InMemorySessionService # Example: Configuring the Runner with an Artifact Service - my_agent = LlmAgent(name="artifact_user_agent", model="gemini-2.0-flash") + my_agent = LlmAgent(name="artifact_user_agent", model="gemini-2.5-flash") artifact_service = InMemoryArtifactService() # Choose an implementation session_service = InMemorySessionService() @@ -3290,7 +3290,7 @@ Before you can use any artifact methods via the context objects, you **must** pr from google.adk.sessions import InMemorySessionService # Your agent definition - agent = LlmAgent(name="my_agent", model="gemini-2.0-flash") + agent = LlmAgent(name="my_agent", model="gemini-2.5-flash") # Instantiate the desired artifact service artifact_service = InMemoryArtifactService() @@ -3663,7 +3663,7 @@ Callbacks are a cornerstone feature of ADK, providing a powerful mechanism to ho # --- Register it during Agent creation --- my_agent = LlmAgent( name="MyCallbackAgent", - model="gemini-2.0-flash", # Or your desired model + model="gemini-2.5-flash", # Or your desired model instruction="Be helpful.", # Other agent parameters... before_model_callback=my_before_model_logic # Pass the function here @@ -3726,7 +3726,7 @@ from typing import Optional from google.genai import types from google.adk.sessions import InMemorySessionService -GEMINI_2_FLASH="gemini-2.0-flash" +GEMINI_2_FLASH="gemini-2.5-flash" # --- Define the Callback Function --- def simple_before_model_modifier( @@ -3839,7 +3839,7 @@ await call_agent_async("write a joke on BLOCK") from google.genai import types from google.adk.sessions import InMemorySessionService - GEMINI_2_FLASH="gemini-2.0-flash" + GEMINI_2_FLASH="gemini-2.5-flash" # --- Define the Callback Function --- def simple_before_model_modifier( @@ -3990,7 +3990,7 @@ These callbacks are available on *any* agent that inherits from `BaseAgent` (inc from typing import Optional # Define the model - Use the specific model name requested - GEMINI_2_FLASH="gemini-2.0-flash" + GEMINI_2_FLASH="gemini-2.5-flash" # --- 1. Define the Callback Function --- def check_if_agent_should_run(callback_context: CallbackContext) -> Optional[types.Content]: @@ -4157,7 +4157,7 @@ These callbacks are available on *any* agent that inherits from `BaseAgent` (inc from typing import Optional # Define the model - Use the specific model name requested - GEMINI_2_FLASH="gemini-2.0-flash" + GEMINI_2_FLASH="gemini-2.5-flash" # --- 1. Define the Callback Function --- def modify_output_after_agent(callback_context: CallbackContext) -> Optional[types.Content]: @@ -4319,7 +4319,7 @@ If the callback returns `None` (or a `Maybe.empty()` object in Java), the LLM co from google.genai import types from google.adk.sessions import InMemorySessionService - GEMINI_2_FLASH="gemini-2.0-flash" + GEMINI_2_FLASH="gemini-2.5-flash" # --- Define the Callback Function --- def simple_before_model_modifier( @@ -4449,7 +4449,7 @@ If the callback returns `None` (or a `Maybe.empty()` object in Java), the LLM co from google.adk.sessions import InMemorySessionService from google.adk.models import LlmResponse - GEMINI_2_FLASH="gemini-2.0-flash" + GEMINI_2_FLASH="gemini-2.5-flash" # --- Define the Callback Function --- def simple_after_model_modifier( @@ -4592,7 +4592,7 @@ These callbacks are also specific to `LlmAgent` and trigger around the execution from typing import Dict, Any - GEMINI_2_FLASH="gemini-2.0-flash" + GEMINI_2_FLASH="gemini-2.5-flash" def get_capital_city(country: str) -> str: """Retrieves the capital city of a given country.""" @@ -4712,7 +4712,7 @@ These callbacks are also specific to `LlmAgent` and trigger around the execution from typing import Dict, Any from copy import deepcopy - GEMINI_2_FLASH="gemini-2.0-flash" + GEMINI_2_FLASH="gemini-2.5-flash" # --- Define a Simple Tool Function (Same as before) --- def get_capital_city(country: str) -> str: @@ -5703,7 +5703,7 @@ def get_current_time(city: str) -> dict: root_agent = Agent( name="weather_time_agent", - model="gemini-2.0-flash", + model="gemini-2.5-flash", description=( "Agent to answer questions about the time and weather in a city." ), @@ -8214,7 +8214,7 @@ application entirely on your machine and is recommended for internal development root_agent = Agent( name="weather_time_agent", - model="gemini-2.0-flash", + model="gemini-2.5-flash", description=( "Agent to answer questions about the time and weather in a city." ), @@ -8387,7 +8387,7 @@ agent will be unable to function. ```py root_agent = Agent( name="weather_time_agent", - model="replace-me-with-model-id", #e.g. gemini-2.0-flash-live-001 + model="replace-me-with-model-id", #e.g. gemini-2.5-flash-live-001 ... ``` @@ -8628,7 +8628,7 @@ Create the **ScienceTeacherAgent.java** file under the `src/main/java/agents/` d !!!note "Troubleshooting" - The model `gemini-2.0-flash-exp` will be deprecated in the future. If you see any issues on using it, try using `gemini-2.0-flash-live-001` instead + The model `gemini-2.5-flash-exp` will be deprecated in the future. If you see any issues on using it, try using `gemini-2.5-flash-live-001` instead We will use `Dev UI` to run this agent later. For the tool to automatically recognize the agent, its Java class has to comply with the following two rules: @@ -8903,7 +8903,7 @@ root_agent = Agent( # The Large Language Model (LLM) that agent will use. # Please fill in the latest model id that supports live from # https://google.github.io/adk-docs/get-started/streaming/quickstart-streaming/#supported-models - model="...", # for example: model="gemini-2.0-flash-live-001" or model="gemini-2.0-flash-live-preview-04-09" + model="...", # for example: model="gemini-2.5-flash-live-001" or model="gemini-2.5-flash-live-preview-04-09" # A short description of the agent's purpose. description="Agent to answer questions using Google Search.", # Instructions to set the agent's behavior. @@ -9505,7 +9505,7 @@ def get_weather(city: str) -> dict: # Create an agent with tools agent = Agent( name="weather_agent", - model="gemini-2.0-flash-exp", + model="gemini-2.5-flash-exp", description="Agent to answer questions using weather tools.", instruction="You must use the available tools to find an answer.", tools=[get_weather] @@ -9640,7 +9640,7 @@ def get_weather(city: str) -> dict: # Create an agent with tools agent = Agent( name="weather_agent", - model="gemini-2.0-flash-exp", + model="gemini-2.5-flash-exp", description="Agent to answer questions using weather tools.", instruction="You must use the available tools to find an answer.", tools=[get_weather] @@ -10403,7 +10403,7 @@ When modifications to the tools to add guardrails aren't possible, the [**`Befor # Hypothetical Agent setup root_agent = LlmAgent( # Use specific agent type - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='root_agent', instruction="...", before_tool_callback=validate_tool_params, # Assign the callback @@ -10650,7 +10650,7 @@ This example demonstrates the basic flow using the `InMemory` services for simpl # --- Constants --- APP_NAME = "memory_example_app" USER_ID = "mem_user" - MODEL = "gemini-2.0-flash" # Use a valid model + MODEL = "gemini-2.5-flash" # Use a valid model # --- Agent Definitions --- # Agent 1: Simple agent to capture information @@ -11031,7 +11031,7 @@ This is the simplest method for saving an agent's final text response directly i # Define agent with output_key greeting_agent = LlmAgent( name="Greeter", - model="gemini-2.0-flash", # Use a valid model + model="gemini-2.5-flash", # Use a valid model instruction="Generate a short, friendly greeting.", output_key="last_greeting" # Save response to state['last_greeting'] ) @@ -11339,15 +11339,15 @@ from google.adk.tools import google_search # Import the tool root_agent = Agent( name="google_search_agent", - model="gemini-2.0-flash-exp", # if this model does not work, try below - #model="gemini-2.0-flash-live-001", + model="gemini-2.5-flash-exp", # if this model does not work, try below + #model="gemini-2.5-flash-live-001", description="Agent to answer questions using Google Search.", instruction="Answer the question using the Google Search tool.", tools=[google_search], ) ``` -**Note:** To enable both text and audio/video input, the model must support the generateContent (for text) and bidiGenerateContent methods. Verify these capabilities by referring to the [List Models Documentation](https://ai.google.dev/api/models#method:-models.list). This quickstart utilizes the gemini-2.0-flash-exp model for demonstration purposes. +**Note:** To enable both text and audio/video input, the model must support the generateContent (for text) and bidiGenerateContent methods. Verify these capabilities by referring to the [List Models Documentation](https://ai.google.dev/api/models#method:-models.list). This quickstart utilizes the gemini-2.5-flash-exp model for demonstration purposes. Notice how easily you integrated [grounding with Google Search](https://ai.google.dev/gemini-api/docs/grounding?lang=python#configure-search) capabilities. The `Agent` class and the `google_search` tool handle the complex interactions with the LLM and grounding with the search API, allowing you to focus on the agent's *purpose* and *behavior*. @@ -11406,7 +11406,7 @@ These console logs are important in case you develop your own streaming applicat 6\. **Troubleshooting tips** - **When `ws://` doesn't work:** If you see any errors on the Chrome DevTools with regard to `ws://` connection, try replacing `ws://` with `wss://` on `app/static/js/app.js` at line 28. This may happen when you are running the sample on a cloud environment and using a proxy connection to connect from your browser. -- **When `gemini-2.0-flash-exp` model doesn't work:** If you see any errors on the app server console with regard to `gemini-2.0-flash-exp` model availability, try replacing it with `gemini-2.0-flash-live-001` on `app/google_search_agent/agent.py` at line 6. +- **When `gemini-2.5-flash-exp` model doesn't work:** If you see any errors on the app server console with regard to `gemini-2.5-flash-exp` model availability, try replacing it with `gemini-2.5-flash-live-001` on `app/google_search_agent/agent.py` at line 6. ## 4. Server code overview {#4.-server-side-code-overview} @@ -12092,7 +12092,7 @@ These console logs are important in case you develop your own streaming applicat 6\. **Troubleshooting tips** - **When your browser can't connect to the server via SSH proxy:** SSH proxy used in various cloud services may not work with SSE. Please try without SSH proxy, such as using a local laptop, or try the [WebSocket](custom-streaming-ws.md) version. -- **When `gemini-2.0-flash-exp` model doesn't work:** If you see any errors on the app server console with regard to `gemini-2.0-flash-exp` model availability, try replacing it with `gemini-2.0-flash-live-001` on `app/google_search_agent/agent.py` at line 6. +- **When `gemini-2.5-flash-exp` model doesn't work:** If you see any errors on the app server console with regard to `gemini-2.5-flash-exp` model availability, try replacing it with `gemini-2.5-flash-live-001` on `app/google_search_agent/agent.py` at line 6. ## 4. Agent definition @@ -12105,8 +12105,8 @@ from google.adk.tools import google_search # Import the tool root_agent = Agent( name="google_search_agent", - model="gemini-2.0-flash-exp", # if this model does not work, try below - #model="gemini-2.0-flash-live-001", + model="gemini-2.5-flash-exp", # if this model does not work, try below + #model="gemini-2.5-flash-live-001", description="Agent to answer questions using Google Search.", instruction="Answer the question using the Google Search tool.", tools=[google_search], @@ -13213,7 +13213,7 @@ async def monitor_video_stream( # Call the model to generate content based on the provided image and prompt response = client.models.generate_content( - model="gemini-2.0-flash-exp", + model="gemini-2.5-flash-exp", contents=contents, config=genai_types.GenerateContentConfig( system_instruction=( @@ -13247,7 +13247,7 @@ def stop_streaming(function_name: str): root_agent = Agent( - model="gemini-2.0-flash-exp", + model="gemini-2.5-flash-exp", name="video_streaming_agent", instruction=""" You are a monitoring agent. You can do video monitoring and stock price monitoring @@ -13872,7 +13872,7 @@ except Exception as e: # --- Agent Configuration --- # Configure and create the main LLM Agent. root_agent = LlmAgent( - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='enterprise_assistant', instruction='Help user integrate with multiple enterprise systems, including retrieving user information which may require authentication.', tools=userinfo_toolset.get_tools(), @@ -14341,7 +14341,7 @@ Search. The `google_search` tool is only compatible with Gemini 2 models. root_agent = Agent( name="basic_search_agent", - model="gemini-2.0-flash", + model="gemini-2.5-flash", description="Agent to answer questions using Google Search.", instruction="I can answer your questions by searching the internet. Just ask me anything!", # google_search is a pre-built tool which allows the agent to perform Google searches. @@ -14410,7 +14410,7 @@ like calculations, data manipulation, or running small scripts. APP_NAME = "calculator" USER_ID = "user1234" SESSION_ID = "session_code_exec_async" - GEMINI_MODEL = "gemini-2.0-flash" + GEMINI_MODEL = "gemini-2.5-flash" # Agent Definition code_agent = LlmAgent( @@ -14537,7 +14537,7 @@ APP_NAME_VSEARCH = "vertex_search_app" USER_ID_VSEARCH = "user_vsearch_1" SESSION_ID_VSEARCH = "session_vsearch_1" AGENT_NAME_VSEARCH = "doc_qa_agent" -GEMINI_2_FLASH = "gemini-2.0-flash" +GEMINI_2_FLASH = "gemini-2.5-flash" # Tool Instantiation # You MUST provide your datastore ID here. @@ -14659,7 +14659,7 @@ AGENT_NAME = "bigquery_agent" APP_NAME = "bigquery_app" USER_ID = "user1234" SESSION_ID = "1234" -GEMINI_MODEL = "gemini-2.0-flash" +GEMINI_MODEL = "gemini-2.5-flash" # Define a tool configuration to block any write operations tool_config = BigQueryToolConfig(write_mode=WriteMode.BLOCKED) @@ -14734,7 +14734,7 @@ to use built-in tools with other tools by using multiple agents: search_agent = Agent( - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='SearchAgent', instruction=""" You're a specialist in Google Search @@ -14742,7 +14742,7 @@ to use built-in tools with other tools by using multiple agents: tools=[google_search], ) coding_agent = Agent( - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='CodeAgent', instruction=""" You're a specialist in Code Execution @@ -14751,7 +14751,7 @@ to use built-in tools with other tools by using multiple agents: ) root_agent = Agent( name="RootAgent", - model="gemini-2.0-flash", + model="gemini-2.5-flash", description="Root Agent", tools=[agent_tool.AgentTool(agent=search_agent), agent_tool.AgentTool(agent=coding_agent)], ) @@ -14777,7 +14777,7 @@ to use built-in tools with other tools by using multiple agents: ```py root_agent = Agent( name="RootAgent", - model="gemini-2.0-flash", + model="gemini-2.5-flash", description="Root Agent", tools=[custom_function], executor=[BuiltInCodeExecutor] # <-- not supported when used with tools @@ -14799,7 +14799,7 @@ is **not** currently supported: ```py search_agent = Agent( - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='SearchAgent', instruction=""" You're a specialist in Google Search @@ -14807,7 +14807,7 @@ is **not** currently supported: tools=[google_search], ) coding_agent = Agent( - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='CodeAgent', instruction=""" You're a specialist in Code Execution @@ -14816,7 +14816,7 @@ is **not** currently supported: ) root_agent = Agent( name="RootAgent", - model="gemini-2.0-flash", + model="gemini-2.5-flash", description="Root Agent", sub_agents=[ search_agent, @@ -14921,7 +14921,7 @@ The docstring (or comments above) your function serve as the tool's description stock_price_agent = Agent( - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='stock_agent', instruction= 'You are an agent who retrieves stock prices. If a ticker symbol is provided, fetch the current price. If only a company name is given, first perform a Google search to find the correct ticker symbol before retrieving the stock price. If the provided ticker symbol is invalid or data cannot be retrieved, inform the user that the stock price could not be found.', description='This agent specializes in retrieving real-time stock prices. Given a stock ticker symbol (e.g., AAPL, GOOG, MSFT) or the stock name, use the tools and reliable data sources to provide the most up-to-date price.', @@ -15122,7 +15122,7 @@ Agent client received an event with long running function calls and check the st # 3. Use the tool in an Agent file_processor_agent = Agent( # Use a model compatible with function calling - model="gemini-2.0-flash", + model="gemini-2.5-flash", name='reimbursement_agent', instruction=""" You are an agent whose job is to handle the reimbursement process for @@ -15299,14 +15299,14 @@ The `AgentTool` class provides the following attributes for customizing its beha SESSION_ID="1234" summary_agent = Agent( - model="gemini-2.0-flash", + model="gemini-2.5-flash", name="summary_agent", instruction="""You are an expert summarizer. Please read the following text and provide a concise summary.""", description="Agent to summarize text", ) root_agent = Agent( - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='root_agent', instruction="""You are a helpful assistant. When the user provides a text, use the 'summarize' tool to generate a summary. Always forward the user's message exactly as received to the 'summarize' tool, without modifying or summarizing it yourself. Present the response from the tool to the user.""", tools=[AgentTool(agent=summary_agent)] @@ -15475,7 +15475,7 @@ you only need to follow a subset of these steps. from .tools import sample_toolset root_agent = LlmAgent( - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='enterprise_assistant', instruction='Help user, leverage the tools you have access to', tools=sample_toolset.get_tools(), @@ -15651,7 +15651,7 @@ Connect your agent to enterprise applications using from .tools import connector_tool root_agent = LlmAgent( - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='connector_agent', instruction="Help user, leverage the tools you have access to", tools=[connector_tool], @@ -15706,7 +15706,7 @@ workflow as a tool for your agent or create a new one. from .tools import integration_tool, connector_tool root_agent = LlmAgent( - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='integration_agent', instruction="Help user, leverage the tools you have access to", tools=[integration_tool], @@ -15896,7 +15896,7 @@ The following example showcases how an agent can use tools by **referencing thei APP_NAME="weather_sentiment_agent" USER_ID="user1234" SESSION_ID="1234" - MODEL_ID="gemini-2.0-flash" + MODEL_ID="gemini-2.5-flash" # Tool 1 def get_weather_report(city: str) -> dict: @@ -16096,14 +16096,14 @@ The `tool_context.actions` attribute (`ToolContext.actions()` in Java) holds an escalation_tool = FunctionTool(func=check_and_transfer) main_agent = Agent( - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='main_agent', instruction="""You are the first point of contact for customer support of an analytics tool. Answer general queries. If the user indicates urgency, use the 'check_and_transfer' tool.""", tools=[check_and_transfer] ) support_agent = Agent( - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='support_agent', instruction="""You are the dedicated support agent. Mentioned you are a support handler and please help the user with their urgent issue.""" ) @@ -16435,7 +16435,7 @@ math_toolset_instance = SimpleMathToolset(prefix="calculator_") # 5. Define an agent that uses both the individual tool and the toolset calculator_agent = LlmAgent( name="CalculatorAgent", - model="gemini-2.0-flash", # Replace with your desired model + model="gemini-2.5-flash", # Replace with your desired model instruction="You are a helpful calculator and greeter. " "Use 'greet_user' for greetings. " "Use 'calculator_add_numbers' to add and 'calculator_subtract_numbers' to subtract. " @@ -16528,7 +16528,7 @@ TARGET_FOLDER_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "/ # If you created ./adk_agent_samples/mcp_agent/your_folder, root_agent = LlmAgent( - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='filesystem_assistant_agent', instruction='Help the user manage their files. You can list files, read files, etc.', tools=[ @@ -16626,7 +16626,7 @@ if not google_maps_api_key: # You might want to raise an error or exit if the key is crucial and not found. root_agent = LlmAgent( - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='maps_assistant_agent', instruction='Help the user with mapping, directions, and finding places using Google Maps tools.', tools=[ @@ -16858,7 +16858,7 @@ if PATH_TO_YOUR_MCP_SERVER_SCRIPT == "/path/to/your/my_adk_mcp_server.py": # Optionally, raise an error if the path is critical root_agent = LlmAgent( - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='web_reader_mcp_client_agent', instruction="Use the 'load_web_page' tool to fetch content from a URL provided by the user.", tools=[ @@ -16961,7 +16961,7 @@ async def get_agent_async(): # Use in an agent root_agent = LlmAgent( - model='gemini-2.0-flash', # Adjust model name if needed based on availability + model='gemini-2.5-flash', # Adjust model name if needed based on availability name='enterprise_assistant', instruction='Help user accessing their file systems', tools=[toolset], # Provide the MCP tools to the ADK agent @@ -17107,7 +17107,7 @@ Follow these steps to integrate an OpenAPI spec into your agent: my_agent = LlmAgent( name="api_interacting_agent", - model="gemini-2.0-flash", # Or your preferred model + model="gemini-2.5-flash", # Or your preferred model tools=[toolset], # Pass the toolset # ... other agent config ... ) @@ -17157,7 +17157,7 @@ This example demonstrates generating tools from a simple Pet Store OpenAPI spec USER_ID_OPENAPI = "user_openapi_1" SESSION_ID_OPENAPI = f"session_openapi_{uuid.uuid4()}" # Unique session ID AGENT_NAME_OPENAPI = "petstore_manager_agent" - GEMINI_MODEL = "gemini-2.0-flash" + GEMINI_MODEL = "gemini-2.5-flash" # --- Sample OpenAPI Specification (JSON String) --- # A basic Pet Store API example using httpbin.org as a mock server @@ -17417,7 +17417,7 @@ ADK provides the `LangchainTool` wrapper to integrate tools from the LangChain e # Define the ADK agent, including the wrapped tool my_agent = Agent( name="langchain_tool_agent", - model="gemini-2.0-flash", + model="gemini-2.5-flash", description="Agent to answer questions using TavilySearch.", instruction="I can answer your questions by searching the internet. Just ask me anything!", tools=[adk_tavily_tool] # Add the wrapped tool here @@ -17473,7 +17473,7 @@ adk_tavily_tool = LangchainTool(tool=tavily_search) # Define Agent with the wrapped tool my_agent = Agent( name="langchain_tool_agent", - model="gemini-2.0-flash", + model="gemini-2.5-flash", description="Agent to answer questions using TavilySearch.", instruction="I can answer your questions by searching the internet. Just ask me anything!", tools=[adk_tavily_tool] # Add the wrapped tool here @@ -17557,7 +17557,7 @@ ADK provides the `CrewaiTool` wrapper to integrate tools from the CrewAI library # Define the ADK agent my_agent = Agent( name="crewai_search_agent", - model="gemini-2.0-flash", + model="gemini-2.5-flash", description="Agent to find recent news using the Serper search tool.", instruction="I can find the latest news for you. What topic are you interested in?", tools=[adk_serper_tool] # Add the wrapped tool here @@ -17614,7 +17614,7 @@ adk_serper_tool = CrewaiTool( serper_agent = Agent( name="basic_search_agent", - model="gemini-2.0-flash", + model="gemini-2.5-flash", description="Agent to answer questions using Google Search.", instruction="I can answer your questions by searching the internet. Just ask me anything!", # Add the Serper tool @@ -17810,7 +17810,7 @@ os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "False" # --- Define Model Constants for easier use --- # More supported models can be referenced here: https://ai.google.dev/gemini-api/docs/models#model-variations -MODEL_GEMINI_2_0_FLASH = "gemini-2.0-flash" +MODEL_GEMINI_2_5_FLASH = "gemini-2.5-flash" # More supported models can be referenced here: https://docs.litellm.ai/docs/providers/openai#openai-chat-completion-models MODEL_GPT_4O = "openai/gpt-4.1" # You can also try: gpt-4.1-mini, gpt-4o etc. @@ -17891,7 +17891,7 @@ Now, let's create the **Agent** itself. An `Agent` in ADK orchestrates the inter We configure it with several key parameters: * `name`: A unique identifier for this agent (e.g., "weather\_agent\_v1"). -* `model`: Specifies which LLM to use (e.g., `MODEL_GEMINI_2_0_FLASH`). We'll start with a specific Gemini model. +* `model`: Specifies which LLM to use (e.g., `MODEL_GEMINI_2_5_FLASH`). We'll start with a specific Gemini model. * `description`: A concise summary of the agent's overall purpose. This becomes crucial later when other agents need to decide whether to delegate tasks to *this* agent. * `instruction`: Detailed guidance for the LLM on how to behave, its persona, its goals, and specifically *how and when* to utilize its assigned `tools`. * `tools`: A list containing the actual Python tool functions the agent is allowed to use (e.g., `[get_weather]`). @@ -17904,7 +17904,7 @@ We configure it with several key parameters: ```python # @title Define the Weather Agent # Use one of the model constants defined earlier -AGENT_MODEL = MODEL_GEMINI_2_0_FLASH # Starting with Gemini +AGENT_MODEL = MODEL_GEMINI_2_5_FLASH # Starting with Gemini weather_agent = Agent( name="weather_agent_v1", @@ -18365,14 +18365,14 @@ Now, create the `Agent` instances for our specialists. Notice their highly focus # If you want to use models other than Gemini, Ensure LiteLlm is imported and API keys are set (from Step 0/2) # from google.adk.models.lite_llm import LiteLlm # MODEL_GPT_4O, MODEL_CLAUDE_SONNET etc. should be defined -# Or else, continue to use: model = MODEL_GEMINI_2_0_FLASH +# Or else, continue to use: model = MODEL_GEMINI_2_5_FLASH # --- Greeting Agent --- greeting_agent = None try: greeting_agent = Agent( # Using a potentially different/cheaper model for a simple task - model = MODEL_GEMINI_2_0_FLASH, + model = MODEL_GEMINI_2_5_FLASH, # model=LiteLlm(model=MODEL_GPT_4O), # If you would like to experiment with other models name="greeting_agent", instruction="You are the Greeting Agent. Your ONLY task is to provide a friendly greeting to the user. " @@ -18391,7 +18391,7 @@ farewell_agent = None try: farewell_agent = Agent( # Can use the same or a different model - model = MODEL_GEMINI_2_0_FLASH, + model = MODEL_GEMINI_2_5_FLASH, # model=LiteLlm(model=MODEL_GPT_4O), # If you would like to experiment with other models name="farewell_agent", instruction="You are the Farewell Agent. Your ONLY task is to provide a polite goodbye message. " @@ -18430,7 +18430,7 @@ runner_root = None # Initialize runner if greeting_agent and farewell_agent and 'get_weather' in globals(): # Let's use a capable Gemini model for the root agent to handle orchestration - root_agent_model = MODEL_GEMINI_2_0_FLASH + root_agent_model = MODEL_GEMINI_2_5_FLASH weather_agent_team = Agent( name="weather_agent_v2", # Give it a new version name @@ -18735,13 +18735,13 @@ from google.adk.agents import Agent from google.adk.models.lite_llm import LiteLlm from google.adk.runners import Runner # Ensure tools 'say_hello', 'say_goodbye' are defined (from Step 3) -# Ensure model constants MODEL_GPT_4O, MODEL_GEMINI_2_0_FLASH etc. are defined +# Ensure model constants MODEL_GPT_4O, MODEL_GEMINI_2_5_FLASH etc. are defined # --- Redefine Greeting Agent (from Step 3) --- greeting_agent = None try: greeting_agent = Agent( - model=MODEL_GEMINI_2_0_FLASH, + model=MODEL_GEMINI_2_5_FLASH, name="greeting_agent", instruction="You are the Greeting Agent. Your ONLY task is to provide a friendly greeting using the 'say_hello' tool. Do nothing else.", description="Handles simple greetings and hellos using the 'say_hello' tool.", @@ -18755,7 +18755,7 @@ except Exception as e: farewell_agent = None try: farewell_agent = Agent( - model=MODEL_GEMINI_2_0_FLASH, + model=MODEL_GEMINI_2_5_FLASH, name="farewell_agent", instruction="You are the Farewell Agent. Your ONLY task is to provide a polite goodbye message using the 'say_goodbye' tool. Do not perform any other actions.", description="Handles simple farewells and goodbyes using the 'say_goodbye' tool.", @@ -18772,7 +18772,7 @@ runner_root_stateful = None # Initialize runner # Check prerequisites before creating the root agent if greeting_agent and farewell_agent and 'get_weather_stateful' in globals(): - root_agent_model = MODEL_GEMINI_2_0_FLASH # Choose orchestration model + root_agent_model = MODEL_GEMINI_2_5_FLASH # Choose orchestration model root_agent_stateful = Agent( name="weather_agent_v4_stateful", # New version name @@ -19062,7 +19062,7 @@ greeting_agent = None try: # Use a defined model constant greeting_agent = Agent( - model=MODEL_GEMINI_2_0_FLASH, + model=MODEL_GEMINI_2_5_FLASH, name="greeting_agent", # Keep original name for consistency instruction="You are the Greeting Agent. Your ONLY task is to provide a friendly greeting using the 'say_hello' tool. Do nothing else.", description="Handles simple greetings and hellos using the 'say_hello' tool.", @@ -19076,7 +19076,7 @@ farewell_agent = None try: # Use a defined model constant farewell_agent = Agent( - model=MODEL_GEMINI_2_0_FLASH, + model=MODEL_GEMINI_2_5_FLASH, name="farewell_agent", # Keep original name instruction="You are the Farewell Agent. Your ONLY task is to provide a polite goodbye message using the 'say_goodbye' tool. Do not perform any other actions.", description="Handles simple farewells and goodbyes using the 'say_goodbye' tool.", @@ -19095,7 +19095,7 @@ runner_root_model_guardrail = None if greeting_agent and farewell_agent and 'get_weather_stateful' in globals() and 'block_keyword_guardrail' in globals(): # Use a defined model constant - root_agent_model = MODEL_GEMINI_2_0_FLASH + root_agent_model = MODEL_GEMINI_2_5_FLASH root_agent_model_guardrail = Agent( name="weather_agent_v5_model_guardrail", # New version name for clarity @@ -19354,7 +19354,7 @@ greeting_agent = None try: # Use a defined model constant greeting_agent = Agent( - model=MODEL_GEMINI_2_0_FLASH, + model=MODEL_GEMINI_2_5_FLASH, name="greeting_agent", # Keep original name for consistency instruction="You are the Greeting Agent. Your ONLY task is to provide a friendly greeting using the 'say_hello' tool. Do nothing else.", description="Handles simple greetings and hellos using the 'say_hello' tool.", @@ -19368,7 +19368,7 @@ farewell_agent = None try: # Use a defined model constant farewell_agent = Agent( - model=MODEL_GEMINI_2_0_FLASH, + model=MODEL_GEMINI_2_5_FLASH, name="farewell_agent", # Keep original name instruction="You are the Farewell Agent. Your ONLY task is to provide a polite goodbye message using the 'say_goodbye' tool. Do not perform any other actions.", description="Handles simple farewells and goodbyes using the 'say_goodbye' tool.", @@ -19388,7 +19388,7 @@ if ('greeting_agent' in globals() and greeting_agent and 'block_keyword_guardrail' in globals() and 'block_paris_tool_guardrail' in globals()): - root_agent_model = MODEL_GEMINI_2_0_FLASH + root_agent_model = MODEL_GEMINI_2_5_FLASH root_agent_tool_guardrail = Agent( name="weather_agent_v6_tool_guardrail", # New version name diff --git a/llms.txt b/llms.txt index 0ff16cbb..97e83563 100644 --- a/llms.txt +++ b/llms.txt @@ -92,7 +92,7 @@ from google.adk.tools import google_search root_agent = Agent( name="search_assistant", - model="gemini-2.0-flash", # Or your preferred Gemini model + model="gemini-2.5-flash", # Or your preferred Gemini model instruction="You are a helpful assistant. Answer user questions using Google Search when needed.", description="An assistant that can search the web.", tools=[google_search] @@ -109,13 +109,13 @@ Define a multi-agent system with coordinator agent, greeter agent, and task exec from google.adk.agents import LlmAgent, BaseAgent # Define individual agents -greeter = LlmAgent(name="greeter", model="gemini-2.0-flash", ...) -task_executor = LlmAgent(name="task_executor", model="gemini-2.0-flash", ...) +greeter = LlmAgent(name="greeter", model="gemini-2.5-flash", ...) +task_executor = LlmAgent(name="task_executor", model="gemini-2.5-flash", ...) # Create parent agent and assign children via sub_agents coordinator = LlmAgent( name="Coordinator", - model="gemini-2.0-flash", + model="gemini-2.5-flash", description="I coordinate greetings and tasks.", sub_agents=[ # Assign sub_agents here greeter, diff --git a/src/google/adk/a2a/converters/event_converter.py b/src/google/adk/a2a/converters/event_converter.py index 5e941a78..df824763 100644 --- a/src/google/adk/a2a/converters/event_converter.py +++ b/src/google/adk/a2a/converters/event_converter.py @@ -14,6 +14,7 @@ from __future__ import annotations +from collections.abc import Callable from datetime import datetime from datetime import timezone import logging @@ -57,6 +58,34 @@ DEFAULT_ERROR_MESSAGE = "An error occurred during processing" logger = logging.getLogger("google_adk." + __name__) +AdkEventToA2AEventsConverter = Callable[ + [ + Event, + InvocationContext, + Optional[str], + Optional[str], + GenAIPartToA2APartConverter, + ], + List[A2AEvent], +] +"""A callable that converts an ADK Event into a list of A2A events. + +This interface allows for custom logic to map ADK's event structure to the +event structure expected by the A2A server. + +Args: + event: The source ADK Event to convert. + invocation_context: The context of the ADK agent invocation. + task_id: The ID of the A2A task being processed. + context_id: The context ID from the A2A request. + part_converter: A function to convert GenAI content parts to A2A + parts. + +Returns: + A list of A2A events. +""" + + def _serialize_metadata_value(value: Any) -> str: """Safely serializes metadata values to string format. diff --git a/src/google/adk/a2a/converters/request_converter.py b/src/google/adk/a2a/converters/request_converter.py index 78a6d78e..92166c3f 100644 --- a/src/google/adk/a2a/converters/request_converter.py +++ b/src/google/adk/a2a/converters/request_converter.py @@ -14,8 +14,12 @@ from __future__ import annotations +from collections.abc import Callable import sys from typing import Any +from typing import Optional + +from pydantic import BaseModel try: from a2a.server.agent_execution import RequestContext @@ -35,6 +39,39 @@ from .part_converter import A2APartToGenAIPartConverter from .part_converter import convert_a2a_part_to_genai_part +@a2a_experimental +class AgentRunRequest(BaseModel): + """Data model for arguments passed to the ADK runner.""" + + user_id: Optional[str] = None + session_id: Optional[str] = None + invocation_id: Optional[str] = None + new_message: Optional[genai_types.Content] = None + state_delta: Optional[dict[str, Any]] = None + run_config: Optional[RunConfig] = None + + +A2ARequestToAgentRunRequestConverter = Callable[ + [ + RequestContext, + A2APartToGenAIPartConverter, + ], + AgentRunRequest, +] +"""A callable that converts an A2A RequestContext to RunnerRequest for ADK runner. + +This interface allows for custom logic to map an incoming A2A RequestContext to the +structured arguments expected by the ADK runner's `run_async` method. + +Args: + request: The incoming request context from the A2A server. + part_converter: A function to convert A2A content parts to GenAI parts. + +Returns: + An RunnerRequest object containing the keyword arguments for ADK runner's run_async method. +""" + + def _get_user_id(request: RequestContext) -> str: # Get user from call context if available (auth is enabled on a2a server) if ( @@ -49,20 +86,32 @@ def _get_user_id(request: RequestContext) -> str: @a2a_experimental -def convert_a2a_request_to_adk_run_args( +def convert_a2a_request_to_agent_run_request( request: RequestContext, part_converter: A2APartToGenAIPartConverter = convert_a2a_part_to_genai_part, -) -> dict[str, Any]: +) -> AgentRunRequest: + """Converts an A2A RequestContext to an AgentRunRequest model. + + Args: + request: The incoming request context from the A2A server. + part_converter: A function to convert A2A content parts to GenAI parts. + + Returns: + A AgentRunRequest object ready to be used as arguments for the ADK runner. + + Raises: + ValueError: If the request message is None. + """ if not request.message: raise ValueError('Request message cannot be None') - return { - 'user_id': _get_user_id(request), - 'session_id': request.context_id, - 'new_message': genai_types.Content( + return AgentRunRequest( + user_id=_get_user_id(request), + session_id=request.context_id, + new_message=genai_types.Content( role='user', parts=[part_converter(part) for part in request.message.parts], ), - 'run_config': RunConfig(), - } + run_config=RunConfig(), + ) diff --git a/src/google/adk/a2a/executor/a2a_agent_executor.py b/src/google/adk/a2a/executor/a2a_agent_executor.py index 4cb92843..608a8188 100644 --- a/src/google/adk/a2a/executor/a2a_agent_executor.py +++ b/src/google/adk/a2a/executor/a2a_agent_executor.py @@ -18,7 +18,6 @@ from datetime import datetime from datetime import timezone import inspect import logging -from typing import Any from typing import Awaitable from typing import Callable from typing import Optional @@ -52,12 +51,15 @@ from google.adk.runners import Runner from pydantic import BaseModel from typing_extensions import override +from ..converters.event_converter import AdkEventToA2AEventsConverter from ..converters.event_converter import convert_event_to_a2a_events from ..converters.part_converter import A2APartToGenAIPartConverter from ..converters.part_converter import convert_a2a_part_to_genai_part from ..converters.part_converter import convert_genai_part_to_a2a_part from ..converters.part_converter import GenAIPartToA2APartConverter -from ..converters.request_converter import convert_a2a_request_to_adk_run_args +from ..converters.request_converter import A2ARequestToAgentRunRequestConverter +from ..converters.request_converter import AgentRunRequest +from ..converters.request_converter import convert_a2a_request_to_agent_run_request from ..converters.utils import _get_adk_metadata_key from ..experimental import a2a_experimental from .task_result_aggregator import TaskResultAggregator @@ -75,6 +77,10 @@ class A2aAgentExecutorConfig(BaseModel): gen_ai_part_converter: GenAIPartToA2APartConverter = ( convert_genai_part_to_a2a_part ) + request_converter: A2ARequestToAgentRunRequestConverter = ( + convert_a2a_request_to_agent_run_request + ) + event_converter: AdkEventToA2AEventsConverter = convert_event_to_a2a_events @a2a_experimental @@ -192,19 +198,20 @@ class A2aAgentExecutor(AgentExecutor): # Resolve the runner instance runner = await self._resolve_runner() - # Convert the a2a request to ADK run args - run_args = convert_a2a_request_to_adk_run_args( - context, self._config.a2a_part_converter + # Convert the a2a request to AgentRunRequest + run_request = self._config.request_converter( + context, + self._config.a2a_part_converter, ) # ensure the session exists - session = await self._prepare_session(context, run_args, runner) + session = await self._prepare_session(context, run_request, runner) # create invocation context invocation_context = runner._new_invocation_context( session=session, - new_message=run_args['new_message'], - run_config=run_args['run_config'], + new_message=run_request.new_message, + run_config=run_request.run_config, ) # publish the task working event @@ -219,16 +226,16 @@ class A2aAgentExecutor(AgentExecutor): final=False, metadata={ _get_adk_metadata_key('app_name'): runner.app_name, - _get_adk_metadata_key('user_id'): run_args['user_id'], - _get_adk_metadata_key('session_id'): run_args['session_id'], + _get_adk_metadata_key('user_id'): run_request.user_id, + _get_adk_metadata_key('session_id'): run_request.session_id, }, ) ) task_result_aggregator = TaskResultAggregator() - async with Aclosing(runner.run_async(**run_args)) as agen: + async with Aclosing(runner.run_async(**vars(run_request))) as agen: async for adk_event in agen: - for a2a_event in convert_event_to_a2a_events( + for a2a_event in self._config.event_converter( adk_event, invocation_context, context.task_id, @@ -284,12 +291,15 @@ class A2aAgentExecutor(AgentExecutor): ) async def _prepare_session( - self, context: RequestContext, run_args: dict[str, Any], runner: Runner + self, + context: RequestContext, + run_request: AgentRunRequest, + runner: Runner, ): - session_id = run_args['session_id'] + session_id = run_request.session_id # create a new session if not exists - user_id = run_args['user_id'] + user_id = run_request.user_id session = await runner.session_service.get_session( app_name=runner.app_name, user_id=user_id, @@ -302,7 +312,7 @@ class A2aAgentExecutor(AgentExecutor): state={}, session_id=session_id, ) - # Update run_args with the new session_id - run_args['session_id'] = session.id + # Update run_request with the new session_id + run_request.session_id = session.id return session diff --git a/src/google/adk/agents/base_agent.py b/src/google/adk/agents/base_agent.py index 9f473de8..88cd05f2 100644 --- a/src/google/adk/agents/base_agent.py +++ b/src/google/adk/agents/base_agent.py @@ -286,7 +286,7 @@ class BaseAgent(BaseModel): with tracer.start_as_current_span(f'invoke_agent {self.name}') as span: ctx = self._create_invocation_context(parent_context) tracing.trace_agent_invocation(span, self, ctx) - if event := await self.__handle_before_agent_callback(ctx): + if event := await self._handle_before_agent_callback(ctx): yield event if ctx.end_invocation: return @@ -298,7 +298,7 @@ class BaseAgent(BaseModel): if ctx.end_invocation: return - if event := await self.__handle_after_agent_callback(ctx): + if event := await self._handle_after_agent_callback(ctx): yield event async with Aclosing(_run_with_trace()) as agen: @@ -324,7 +324,7 @@ class BaseAgent(BaseModel): with tracer.start_as_current_span(f'invoke_agent {self.name}') as span: ctx = self._create_invocation_context(parent_context) tracing.trace_agent_invocation(span, self, ctx) - if event := await self.__handle_before_agent_callback(ctx): + if event := await self._handle_before_agent_callback(ctx): yield event if ctx.end_invocation: return @@ -333,7 +333,7 @@ class BaseAgent(BaseModel): async for event in agen: yield event - if event := await self.__handle_after_agent_callback(ctx): + if event := await self._handle_after_agent_callback(ctx): yield event async with Aclosing(_run_with_trace()) as agen: @@ -438,7 +438,7 @@ class BaseAgent(BaseModel): return self.after_agent_callback return [self.after_agent_callback] - async def __handle_before_agent_callback( + async def _handle_before_agent_callback( self, ctx: InvocationContext ) -> Optional[Event]: """Runs the before_agent_callback if it exists. @@ -496,7 +496,7 @@ class BaseAgent(BaseModel): return None - async def __handle_after_agent_callback( + async def _handle_after_agent_callback( self, invocation_context: InvocationContext ) -> Optional[Event]: """Runs the after_agent_callback if it exists. diff --git a/src/google/adk/artifacts/artifact_util.py b/src/google/adk/artifacts/artifact_util.py new file mode 100644 index 00000000..15cdd4de --- /dev/null +++ b/src/google/adk/artifacts/artifact_util.py @@ -0,0 +1,116 @@ +# 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. +"""Utility functions for handling artifact URIs.""" + +from __future__ import annotations + +import re +from typing import NamedTuple +from typing import Optional + +from google.genai import types + + +class ParsedArtifactUri(NamedTuple): + """The result of parsing an artifact URI.""" + + app_name: str + user_id: str + session_id: Optional[str] + filename: str + version: int + + +_SESSION_SCOPED_ARTIFACT_URI_RE = re.compile( + r"artifact://apps/([^/]+)/users/([^/]+)/sessions/([^/]+)/artifacts/([^/]+)/versions/(\d+)" +) +_USER_SCOPED_ARTIFACT_URI_RE = re.compile( + r"artifact://apps/([^/]+)/users/([^/]+)/artifacts/([^/]+)/versions/(\d+)" +) + + +def parse_artifact_uri(uri: str) -> Optional[ParsedArtifactUri]: + """Parses an artifact URI. + + Args: + uri: The artifact URI to parse. + + Returns: + A ParsedArtifactUri if parsing is successful, None otherwise. + """ + if not uri or not uri.startswith("artifact://"): + return None + + match = _SESSION_SCOPED_ARTIFACT_URI_RE.match(uri) + if match: + return ParsedArtifactUri( + app_name=match.group(1), + user_id=match.group(2), + session_id=match.group(3), + filename=match.group(4), + version=int(match.group(5)), + ) + + match = _USER_SCOPED_ARTIFACT_URI_RE.match(uri) + if match: + return ParsedArtifactUri( + app_name=match.group(1), + user_id=match.group(2), + session_id=None, + filename=match.group(3), + version=int(match.group(4)), + ) + + return None + + +def get_artifact_uri( + app_name: str, + user_id: str, + filename: str, + version: int, + session_id: Optional[str] = None, +) -> str: + """Constructs an artifact URI. + + Args: + app_name: The name of the application. + user_id: The ID of the user. + filename: The name of the artifact file. + version: The version of the artifact. + session_id: The ID of the session. + + Returns: + The constructed artifact URI. + """ + if session_id: + return f"artifact://apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{filename}/versions/{version}" + else: + return f"artifact://apps/{app_name}/users/{user_id}/artifacts/{filename}/versions/{version}" + + +def is_artifact_ref(artifact: types.Part) -> bool: + """Checks if an artifact part is an artifact reference. + + Args: + artifact: The artifact part to check. + + Returns: + True if the artifact part is an artifact reference, False otherwise. + """ + return bool( + artifact.file_data + and artifact.file_data.file_uri + and artifact.file_data.file_uri.startswith("artifact://") + ) diff --git a/src/google/adk/artifacts/gcs_artifact_service.py b/src/google/adk/artifacts/gcs_artifact_service.py index d8a5b2fd..b834584b 100644 --- a/src/google/adk/artifacts/gcs_artifact_service.py +++ b/src/google/adk/artifacts/gcs_artifact_service.py @@ -212,6 +212,11 @@ class GcsArtifactService(BaseArtifactService): blob.upload_from_string( data=artifact.text, ) + elif artifact.file_data: + raise NotImplementedError( + "Saving artifact with file_data is not supported yet in" + " GcsArtifactService." + ) else: raise ValueError("Artifact must have either inline_data or text.") diff --git a/src/google/adk/artifacts/in_memory_artifact_service.py b/src/google/adk/artifacts/in_memory_artifact_service.py index ab556f2b..246e8a85 100644 --- a/src/google/adk/artifacts/in_memory_artifact_service.py +++ b/src/google/adk/artifacts/in_memory_artifact_service.py @@ -13,10 +13,12 @@ # limitations under the License. from __future__ import annotations +import dataclasses import logging from typing import Any from typing import Optional +from google.adk.artifacts import artifact_util from google.genai import types from pydantic import BaseModel from pydantic import Field @@ -28,6 +30,19 @@ from .base_artifact_service import BaseArtifactService logger = logging.getLogger("google_adk." + __name__) +@dataclasses.dataclass +class _ArtifactEntry: + """Represents a single version of an artifact stored in memory. + + Attributes: + data: The actual data of the artifact. + artifact_version: Metadata about this specific version of the artifact. + """ + + data: types.Part + artifact_version: ArtifactVersion + + class InMemoryArtifactService(BaseArtifactService, BaseModel): """An in-memory implementation of the artifact service. @@ -35,7 +50,7 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel): testing and development only. """ - artifacts: dict[str, list[types.Part]] = Field(default_factory=dict) + artifacts: dict[str, list[_ArtifactEntry]] = Field(default_factory=dict) def _file_has_user_namespace(self, filename: str) -> bool: """Checks if the filename has a user namespace. @@ -87,15 +102,42 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel): session_id: Optional[str] = None, custom_metadata: Optional[dict[str, Any]] = None, ) -> int: - # TODO: b/447451270 - Support saving artifact with custom metadata. - if custom_metadata: - raise NotImplementedError("custom_metadata is not supported yet.") - path = self._artifact_path(app_name, user_id, filename, session_id) if path not in self.artifacts: self.artifacts[path] = [] version = len(self.artifacts[path]) - self.artifacts[path].append(artifact) + if self._file_has_user_namespace(filename): + canonical_uri = f"memory://apps/{app_name}/users/{user_id}/artifacts/{filename}/versions/{version}" + else: + canonical_uri = f"memory://apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{filename}/versions/{version}" + + artifact_version = ArtifactVersion( + version=version, + canonical_uri=canonical_uri, + ) + if custom_metadata: + artifact_version.custom_metadata = custom_metadata + + if artifact.inline_data is not None: + artifact_version.mime_type = artifact.inline_data.mime_type + elif artifact.text is not None: + artifact_version.mime_type = "text/plain" + elif artifact.file_data is not None: + if artifact_util.is_artifact_ref(artifact): + if not artifact_util.parse_artifact_uri(artifact.file_data.file_uri): + raise ValueError( + f"Invalid artifact reference URI: {artifact.file_data.file_uri}" + ) + # If it's a valid artifact URI, we store the artifact part as-is. + # And we don't know the mime type until we load it. + else: + artifact_version.mime_type = artifact.file_data.mime_type + else: + raise ValueError("Not supported artifact type.") + + self.artifacts[path].append( + _ArtifactEntry(data=artifact, artifact_version=artifact_version) + ) return version @override @@ -114,7 +156,41 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel): return None if version is None: version = -1 - return versions[version] + + try: + artifact_entry = versions[version] + except IndexError: + return None + + if artifact_entry is None: + return None + + # Resolve artifact reference if needed. + artifact_data = artifact_entry.data + if artifact_util.is_artifact_ref(artifact_data): + parsed_uri = artifact_util.parse_artifact_uri( + artifact_data.file_data.file_uri + ) + if not parsed_uri: + raise ValueError( + "Invalid artifact reference URI:" + f" {artifact_data.file_data.file_uri}" + ) + return await self.load_artifact( + app_name=parsed_uri.app_name, + user_id=parsed_uri.user_id, + filename=parsed_uri.filename, + session_id=parsed_uri.session_id, + version=parsed_uri.version, + ) + + if ( + artifact_data == types.Part() + or artifact_data == types.Part(text="") + or (artifact_data.inline_data and not artifact_data.inline_data.data) + ): + return None + return artifact_data @override async def list_artifact_keys( @@ -172,8 +248,11 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel): filename: str, session_id: Optional[str] = None, ) -> list[ArtifactVersion]: - # TODO: b/447451270 - Support list_artifact_versions. - raise NotImplementedError("list_artifact_versions is not implemented yet.") + path = self._artifact_path(app_name, user_id, filename, session_id) + entries = self.artifacts.get(path) + if not entries: + return [] + return [entry.artifact_version for entry in entries] @override async def get_artifact_version( @@ -185,5 +264,14 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel): session_id: Optional[str] = None, version: Optional[int] = None, ) -> Optional[ArtifactVersion]: - # TODO: b/447451270 - Support get_artifact_version. - raise NotImplementedError("get_artifact_version is not implemented yet.") + path = self._artifact_path(app_name, user_id, filename, session_id) + entries = self.artifacts.get(path) + if not entries: + return None + + if version is None: + version = -1 + try: + return entries[version].artifact_version + except IndexError: + return None diff --git a/src/google/adk/evaluation/agent_evaluator.py b/src/google/adk/evaluation/agent_evaluator.py index 049d2bf7..3788a2e6 100644 --- a/src/google/adk/evaluation/agent_evaluator.py +++ b/src/google/adk/evaluation/agent_evaluator.py @@ -118,7 +118,7 @@ class AgentEvaluator: Args: agent_module: The path to python module that contains the definition of the agent. There is convention in place here, where the code is going to - look for 'root_agent' in the loaded module. + look for 'root_agent' or `get_agent_async` in the loaded module. eval_set: The eval set. criteria: Evauation criterias, a dictionary of metric names to their respective thresholds. This field is deprecated. @@ -144,7 +144,7 @@ class AgentEvaluator: if eval_config is None: raise ValueError("`eval_config` is required.") - agent_for_eval = AgentEvaluator._get_agent_for_eval( + agent_for_eval = await AgentEvaluator._get_agent_for_eval( module_name=agent_module, agent_name=agent_name ) eval_metrics = get_eval_metrics_from_config(eval_config) @@ -200,7 +200,7 @@ class AgentEvaluator: Args: agent_module: The path to python module that contains the definition of the agent. There is convention in place here, where the code is going to - look for 'root_agent' in the loaded module. + look for 'root_agent' or 'get_agent_async' in the loaded module. eval_dataset_file_path_or_dir: The eval data set. This can be either a string representing full path to the file containing eval dataset, or a directory that is recursively explored for all files that have a @@ -466,12 +466,26 @@ class AgentEvaluator: return "\n".join([str(t) for t in tool_calls]) @staticmethod - def _get_agent_for_eval( + async def _get_agent_for_eval( module_name: str, agent_name: Optional[str] = None ) -> BaseAgent: module_path = f"{module_name}" agent_module = importlib.import_module(module_path) - root_agent = agent_module.agent.root_agent + print(dir(agent_module)) + if hasattr(agent_module, "agent"): + if hasattr(agent_module.agent, "root_agent"): + root_agent = agent_module.agent.root_agent + elif hasattr(agent_module.agent, "get_agent_async"): + root_agent, _ = await agent_module.agent.get_agent_async() + else: + raise ValueError( + f"Module {module_name} does not have a root_agent or" + " get_agent_async method." + ) + else: + raise ValueError( + f"Module {module_name} does not have a member named `agent`." + ) agent_for_eval = root_agent if agent_name: diff --git a/src/google/adk/events/event_actions.py b/src/google/adk/events/event_actions.py index af3168fe..da55c873 100644 --- a/src/google/adk/events/event_actions.py +++ b/src/google/adk/events/event_actions.py @@ -104,3 +104,6 @@ class EventActions(BaseModel): agent_state: Optional[dict[str, Any]] = None """The agent state at the current event.""" + + rewind_before_invocation_id: Optional[str] = None + """The invocation id to rewind to. This is only set for rewind event.""" diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index f5b242b7..531a5034 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -438,6 +438,10 @@ class BaseLlmFlow(ABC): from ...agents.llm_agent import LlmAgent agent = invocation_context.agent + if not isinstance(agent, LlmAgent): + raise TypeError( + f'Expected agent to be an LlmAgent, but got {type(agent)}' + ) # Runs processors. for processor in self.request_processors: @@ -468,7 +472,7 @@ class BaseLlmFlow(ABC): tools = await _convert_tool_union_to_tools( tool_union, ReadonlyContext(invocation_context), - llm_request.model, + agent.model, multiple_tools, ) for tool in tools: diff --git a/src/google/adk/flows/llm_flows/contents.py b/src/google/adk/flows/llm_flows/contents.py index 1ccb230c..3c4da720 100644 --- a/src/google/adk/flows/llm_flows/contents.py +++ b/src/google/adk/flows/llm_flows/contents.py @@ -310,11 +310,30 @@ def _get_contents( accumulated_input_transcription = '' accumulated_output_transcription = '' + # Filter out events that are annulled by a rewind. + # By iterating backward, when a rewind event is found, we skip all events + # from that point back to the `rewind_before_invocation_id`, thus removing + # them from the history used for the LLM request. + rewind_filtered_events = [] + i = len(events) - 1 + while i >= 0: + event = events[i] + if event.actions and event.actions.rewind_before_invocation_id: + rewind_invocation_id = event.actions.rewind_before_invocation_id + for j in range(0, i, 1): + if events[j].invocation_id == rewind_invocation_id: + i = j + break + else: + rewind_filtered_events.append(event) + i -= 1 + rewind_filtered_events.reverse() + # Parse the events, leaving the contents and the function calls and # responses from the current agent. raw_filtered_events = [] has_compaction_events = False - for event in events: + for event in rewind_filtered_events: if _contains_empty_content(event): continue if not _is_event_belongs_to_branch(current_branch, event): diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py index b7508aee..4380322b 100644 --- a/src/google/adk/flows/llm_flows/functions.py +++ b/src/google/adk/flows/llm_flows/functions.py @@ -275,21 +275,37 @@ async def _execute_single_function_call_async( tool_confirmation: Optional[ToolConfirmation] = None, ) -> Optional[Event]: """Execute a single function call with thread safety for state modifications.""" - tool, tool_context = _get_tool_and_context( - invocation_context, - function_call, - tools_dict, - tool_confirmation, + # Do not use "args" as the variable name, because it is a reserved keyword + # in python debugger. + # Make a deep copy to avoid being modified. + function_args = ( + copy.deepcopy(function_call.args) if function_call.args else {} ) - with tracer.start_as_current_span(f'execute_tool {tool.name}'): - # Do not use "args" as the variable name, because it is a reserved keyword - # in python debugger. - # Make a deep copy to avoid being modified. - function_args = ( - copy.deepcopy(function_call.args) if function_call.args else {} - ) + tool_context = _create_tool_context( + invocation_context, function_call, tool_confirmation + ) + try: + tool = _get_tool(function_call, tools_dict) + except ValueError as tool_error: + tool = BaseTool(name=function_call.name, description='Tool not found') + error_response = ( + await invocation_context.plugin_manager.run_on_tool_error_callback( + tool=tool, + tool_args=function_args, + tool_context=tool_context, + error=tool_error, + ) + ) + if error_response is not None: + return __build_response_event( + tool, error_response, tool_context, invocation_context + ) + else: + raise tool_error + + with tracer.start_as_current_span(f'execute_tool {tool.name}'): # Step 1: Check if plugin before_tool_callback overrides the function # response. function_response = ( @@ -639,25 +655,46 @@ async def _process_function_live_helper( return function_response +def _get_tool( + function_call: types.FunctionCall, tools_dict: dict[str, BaseTool] +): + """Returns the tool corresponding to the function call.""" + if function_call.name not in tools_dict: + raise ValueError( + f'Function {function_call.name} is not found in the tools_dict:' + f' {tools_dict.keys()}.' + ) + + return tools_dict[function_call.name] + + +def _create_tool_context( + invocation_context: InvocationContext, + function_call: types.FunctionCall, + tool_confirmation: Optional[ToolConfirmation] = None, +): + """Creates a ToolContext object.""" + return ToolContext( + invocation_context=invocation_context, + function_call_id=function_call.id, + tool_confirmation=tool_confirmation, + ) + + def _get_tool_and_context( invocation_context: InvocationContext, function_call: types.FunctionCall, tools_dict: dict[str, BaseTool], tool_confirmation: Optional[ToolConfirmation] = None, ): - if function_call.name not in tools_dict: - raise ValueError( - f'Function {function_call.name} is not found in the tools_dict.' - ) - - tool_context = ToolContext( - invocation_context=invocation_context, - function_call_id=function_call.id, - tool_confirmation=tool_confirmation, + """Returns the tool and tool context corresponding to the function call.""" + tool = _get_tool(function_call, tools_dict) + tool_context = _create_tool_context( + invocation_context, + function_call, + tool_confirmation, ) - tool = tools_dict[function_call.name] - return (tool, tool_context) diff --git a/src/google/adk/models/gemini_llm_connection.py b/src/google/adk/models/gemini_llm_connection.py index 509128b0..f1470c0a 100644 --- a/src/google/adk/models/gemini_llm_connection.py +++ b/src/google/adk/models/gemini_llm_connection.py @@ -205,7 +205,7 @@ class GeminiLlmConnection(BaseLlmConnection): ] yield LlmResponse(content=types.Content(role='model', parts=parts)) if message.session_resumption_update: - logger.info('Redeived session reassumption message: %s', message) + logger.info('Received session resumption message: %s', message) yield ( LlmResponse( live_session_resumption_update=message.session_resumption_update diff --git a/src/google/adk/plugins/reflect_retry_tool_plugin.py b/src/google/adk/plugins/reflect_retry_tool_plugin.py index cc501a93..a3a0cc25 100644 --- a/src/google/adk/plugins/reflect_retry_tool_plugin.py +++ b/src/google/adk/plugins/reflect_retry_tool_plugin.py @@ -142,8 +142,20 @@ class ReflectAndRetryToolPlugin(BasePlugin): tool_args: dict[str, Any], tool_context: ToolContext, result: Any, - ) -> Optional[dict]: - """Handles successful tool calls or extracts and processes errors.""" + ) -> Optional[dict[str, Any]]: + """Handles successful tool calls or extracts and processes errors. + + Args: + tool: The tool that was called. + tool_args: The arguments passed to the tool. + tool_context: The context of the tool call. + result: The result of the tool call. + + Returns: + An optional dictionary containing reflection guidance if an error is + detected, or None if the tool call was successful or the + response is already a reflection message. + """ if ( isinstance(result, dict) and result.get("response_type") == REFLECT_AND_RETRY_RESPONSE_TYPE @@ -157,7 +169,8 @@ class ReflectAndRetryToolPlugin(BasePlugin): if error: return await self._handle_tool_error(tool, tool_args, tool_context, error) - # On success, reset the failure count for this specific tool within its scope. + # On success, reset the failure count for this specific tool within + # its scope. await self._reset_failures_for_tool(tool_context, tool.name) return None @@ -168,7 +181,7 @@ class ReflectAndRetryToolPlugin(BasePlugin): tool_args: dict[str, Any], tool_context: ToolContext, result: Any, - ) -> Optional[Any]: + ) -> Optional[dict[str, Any]]: """Extracts an error from a successful tool result and triggers retry logic. This is useful when tool call finishes successfully but the result contains @@ -176,6 +189,15 @@ class ReflectAndRetryToolPlugin(BasePlugin): By overriding this method, you can trigger retry logic on these successful results that contain errors. + + Args: + tool: The tool that was called. + tool_args: The arguments passed to the tool. + tool_context: The context of the tool call. + result: The result of the tool call. + + Returns: + The extracted error if any, or None if no error was detected. """ return None @@ -186,8 +208,18 @@ class ReflectAndRetryToolPlugin(BasePlugin): tool_args: dict[str, Any], tool_context: ToolContext, error: Exception, - ) -> Optional[dict]: - """Handles tool exceptions by providing reflection guidance.""" + ) -> Optional[dict[str, Any]]: + """Handles tool exceptions by providing reflection guidance. + + Args: + tool: The tool that was called. + tool_args: The arguments passed to the tool. + tool_context: The context of the tool call. + error: The exception raised by the tool. + + Returns: + An optional dictionary containing reflection guidance for the error. + """ return await self._handle_tool_error(tool, tool_args, tool_context, error) async def _handle_tool_error( @@ -196,8 +228,18 @@ class ReflectAndRetryToolPlugin(BasePlugin): tool_args: dict[str, Any], tool_context: ToolContext, error: Any, - ) -> Optional[dict]: - """Central, thread-safe logic for processing tool errors.""" + ) -> Optional[dict[str, Any]]: + """Central, thread-safe logic for processing tool errors. + + Args: + tool: The tool that was called. + tool_args: The arguments passed to the tool. + tool_context: The context of the tool call. + error: The error to be handled. + + Returns: + An optional dictionary containing reflection guidance for the error. + """ if self.max_retries == 0: if self.throw_exception_if_retry_exceeded: raise error @@ -285,6 +327,7 @@ This is retry attempt **{retry_count} of {self.max_retries}**. Analyze the error 2. **State or Preconditions**: Did a previous step fail or not produce the necessary state/resource for this tool to succeed? 3. **Alternative Approach**: Is this the right tool for the job? Could another tool or a different sequence of steps achieve the goal? 4. **Simplify the Task**: Can you break the problem down into smaller, simpler steps? +5. **Wrong Function Name**: Does the error indicates the tool is not found? Please check again and only use available tools. Formulate a new plan based on your analysis and try a corrected or different approach. """ diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index 26116899..f68008af 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -28,6 +28,7 @@ from typing import Optional import warnings from google.adk.apps.compaction import _run_compaction_for_sliding_window +from google.adk.artifacts import artifact_util from google.genai import types from .agents.active_streaming_tool import ActiveStreamingTool @@ -427,6 +428,146 @@ class Runner: async for event in agen: yield event + async def rewind_async( + self, + *, + user_id: str, + session_id: str, + rewind_before_invocation_id: str, + ) -> None: + """Rewinds the session to before the specified invocation.""" + session = await self.session_service.get_session( + app_name=self.app_name, user_id=user_id, session_id=session_id + ) + if not session: + raise ValueError(f'Session not found: {session_id}') + + rewind_event_index = -1 + for i, event in enumerate(session.events): + if event.invocation_id == rewind_before_invocation_id: + rewind_event_index = i + break + + if rewind_event_index == -1: + raise ValueError( + f'Invocation ID not found: {rewind_before_invocation_id}' + ) + + # Compute state delta to reverse changes + state_delta = await self._compute_state_delta_for_rewind( + session, rewind_event_index + ) + + # Compute artifact delta to reverse changes + artifact_delta = await self._compute_artifact_delta_for_rewind( + session, rewind_event_index + ) + + # Create rewind event + rewind_event = Event( + invocation_id=new_invocation_context_id(), + author='user', + actions=EventActions( + rewind_before_invocation_id=rewind_before_invocation_id, + state_delta=state_delta, + artifact_delta=artifact_delta, + ), + ) + + logger.info('Rewinding session to invocation: %s', rewind_event) + + await self.session_service.append_event(session=session, event=rewind_event) + + async def _compute_state_delta_for_rewind( + self, session: Session, rewind_event_index: int + ) -> dict[str, Any]: + """Computes the state delta to reverse changes.""" + state_at_rewind_point: dict[str, Any] = {} + for i in range(rewind_event_index): + if session.events[i].actions.state_delta: + for k, v in session.events[i].actions.state_delta.items(): + if k.startswith('app:') or k.startswith('user:'): + continue + if v is None: + state_at_rewind_point.pop(k, None) + else: + state_at_rewind_point[k] = v + + current_state = session.state + rewind_state_delta = {} + + # 1. Add/update keys in rewind_state_delta to match state_at_rewind_point. + for key, value_at_rewind in state_at_rewind_point.items(): + if key not in current_state or current_state[key] != value_at_rewind: + rewind_state_delta[key] = value_at_rewind + + # 2. Set keys to None in rewind_state_delta if they are in current_state + # but not in state_at_rewind_point. These keys were added after the + # rewind point and need to be removed. + for key in current_state: + if key.startswith('app:') or key.startswith('user:'): + continue + if key not in state_at_rewind_point: + rewind_state_delta[key] = None + + return rewind_state_delta + + async def _compute_artifact_delta_for_rewind( + self, session: Session, rewind_event_index: int + ) -> dict[str, int]: + """Computes the artifact delta to reverse changes.""" + if not self.artifact_service: + return {} + + versions_at_rewind_point: dict[str, int] = {} + for i in range(rewind_event_index): + event = session.events[i] + if event.actions.artifact_delta: + versions_at_rewind_point.update(event.actions.artifact_delta) + + current_versions: dict[str, int] = {} + for event in session.events: + if event.actions.artifact_delta: + current_versions.update(event.actions.artifact_delta) + + rewind_artifact_delta = {} + for filename, vn in current_versions.items(): + if filename.startswith('user:'): + # User artifacts are not restored on rewind. + continue + vt = versions_at_rewind_point.get(filename) + if vt == vn: + continue + + rewind_artifact_delta[filename] = vn + 1 + if vt is None: + # Artifact did not exist at rewind point. Mark it as inaccessible. + artifact = types.Part( + inline_data=types.Blob( + mime_type='application/octet-stream', data=b'' + ) + ) + else: + # Artifact version changed after rewind point. Restore to version at + # rewind point. + artifact_uri = artifact_util.get_artifact_uri( + app_name=self.app_name, + user_id=session.user_id, + session_id=session.id, + filename=filename, + version=vt, + ) + artifact = types.Part(file_data=types.FileData(file_uri=artifact_uri)) + await self.artifact_service.save_artifact( + app_name=self.app_name, + user_id=session.user_id, + session_id=session.id, + filename=filename, + artifact=artifact, + ) + + return rewind_artifact_delta + async def _run_compaction_default(self, session: Session): """Runs compaction for other types of compactors. @@ -1083,7 +1224,7 @@ class InMemoryRunner(Runner): self, agent: Optional[BaseAgent] = None, *, - app_name: Optional[str] = 'InMemoryRunner', + app_name: Optional[str] = None, plugins: Optional[list[BasePlugin]] = None, app: Optional[App] = None, ): @@ -1094,6 +1235,8 @@ class InMemoryRunner(Runner): app_name: The application name of the runner. Defaults to 'InMemoryRunner'. """ + if app is None and app_name is None: + app_name = 'InMemoryRunner' super().__init__( app_name=app_name, agent=agent, diff --git a/src/google/adk/sessions/_session_util.py b/src/google/adk/sessions/_session_util.py index 2cc65949..68340c8b 100644 --- a/src/google/adk/sessions/_session_util.py +++ b/src/google/adk/sessions/_session_util.py @@ -16,23 +16,16 @@ from __future__ import annotations from typing import Any from typing import Optional +from typing import Type +from typing import TypeVar -from google.genai import types +M = TypeVar("M") -def decode_content( - content: Optional[dict[str, Any]], -) -> Optional[types.Content]: - """Decodes a content object from a JSON dictionary.""" - if not content: +def decode_model( + data: Optional[dict[str, Any]], model_cls: Type[M] +) -> Optional[M]: + """Decodes a pydantic model object from a JSON dictionary.""" + if data is None: return None - return types.Content.model_validate(content) - - -def decode_grounding_metadata( - grounding_metadata: Optional[dict[str, Any]], -) -> Optional[types.GroundingMetadata]: - """Decodes a grounding metadata object from a JSON dictionary.""" - if not grounding_metadata: - return None - return types.GroundingMetadata.model_validate(grounding_metadata) + return model_cls.model_validate(data) diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index 9eca2753..04af695b 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/src/google/adk/sessions/database_session_service.py @@ -23,6 +23,7 @@ from typing import Any from typing import Optional import uuid +from google.genai import types from sqlalchemy import Boolean from sqlalchemy import delete from sqlalchemy import Dialect @@ -252,6 +253,12 @@ class StorageEvent(Base): custom_metadata: Mapped[dict[str, Any]] = mapped_column( DynamicJSON, nullable=True ) + usage_metadata: Mapped[dict[str, Any]] = mapped_column( + DynamicJSON, nullable=True + ) + citation_metadata: Mapped[dict[str, Any]] = mapped_column( + DynamicJSON, nullable=True + ) partial: Mapped[bool] = mapped_column(Boolean, nullable=True) turn_complete: Mapped[bool] = mapped_column(Boolean, nullable=True) @@ -318,6 +325,14 @@ class StorageEvent(Base): ) if event.custom_metadata: storage_event.custom_metadata = event.custom_metadata + if event.usage_metadata: + storage_event.usage_metadata = event.usage_metadata.model_dump( + exclude_none=True, mode="json" + ) + if event.citation_metadata: + storage_event.citation_metadata = event.citation_metadata.model_dump( + exclude_none=True, mode="json" + ) return storage_event def to_event(self) -> Event: @@ -328,17 +343,23 @@ class StorageEvent(Base): branch=self.branch, actions=self.actions, timestamp=self.timestamp.timestamp(), - content=_session_util.decode_content(self.content), long_running_tool_ids=self.long_running_tool_ids, partial=self.partial, turn_complete=self.turn_complete, error_code=self.error_code, error_message=self.error_message, interrupted=self.interrupted, - grounding_metadata=_session_util.decode_grounding_metadata( - self.grounding_metadata - ), custom_metadata=self.custom_metadata, + content=_session_util.decode_model(self.content, types.Content), + grounding_metadata=_session_util.decode_model( + self.grounding_metadata, types.GroundingMetadata + ), + usage_metadata=_session_util.decode_model( + self.usage_metadata, types.GenerateContentResponseUsageMetadata + ), + citation_metadata=_session_util.decode_model( + self.citation_metadata, types.CitationMetadata + ), ) diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index 72ff0d6c..6901f4cb 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/src/google/adk/sessions/vertex_ai_session_service.py @@ -376,8 +376,9 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event: interrupted = getattr(event_metadata, 'interrupted', None) branch = getattr(event_metadata, 'branch', None) custom_metadata = getattr(event_metadata, 'custom_metadata', None) - grounding_metadata = _session_util.decode_grounding_metadata( - getattr(event_metadata, 'grounding_metadata', None) + grounding_metadata = _session_util.decode_model( + getattr(event_metadata, 'grounding_metadata', None), + types.GroundingMetadata, ) else: long_running_tool_ids = None @@ -393,8 +394,8 @@ def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event: invocation_id=api_event_obj.invocation_id, author=api_event_obj.author, actions=event_actions, - content=_session_util.decode_content( - getattr(api_event_obj, 'content', None) + content=_session_util.decode_model( + getattr(api_event_obj, 'content', None), types.Content ), timestamp=api_event_obj.timestamp.timestamp(), error_code=getattr(api_event_obj, 'error_code', None), diff --git a/src/google/adk/telemetry/tracing.py b/src/google/adk/telemetry/tracing.py index 12fbca1d..cc0a2a23 100644 --- a/src/google/adk/telemetry/tracing.py +++ b/src/google/adk/telemetry/tracing.py @@ -24,6 +24,7 @@ from __future__ import annotations import json +import os from typing import Any from typing import TYPE_CHECKING @@ -33,6 +34,9 @@ from opentelemetry import trace from .. import version from ..events.event import Event +# By default some ADK spans include attributes with potential PII data. +# This env, when set to false, allows to disable populating those attributes. +ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS = 'ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS' # TODO: Replace with constant from opentelemetry.semconv when it reaches version 1.37 in g3. GEN_AI_AGENT_DESCRIPTION = 'gen_ai.agent.description' GEN_AI_AGENT_NAME = 'gen_ai.agent.name' @@ -138,10 +142,13 @@ def trace_tool_call( span.set_attribute('gcp.vertex.agent.llm_request', '{}') span.set_attribute('gcp.vertex.agent.llm_response', '{}') - span.set_attribute( - 'gcp.vertex.agent.tool_call_args', - _safe_json_serialize(args), - ) + if _should_add_request_response_to_spans(): + span.set_attribute( + 'gcp.vertex.agent.tool_call_args', + _safe_json_serialize(args), + ) + else: + span.set_attribute('gcp.vertex.agent.tool_call_args', {}) # Tracing tool response tool_call_id = '' @@ -163,10 +170,13 @@ def trace_tool_call( if not isinstance(tool_response, dict): tool_response = {'result': tool_response} span.set_attribute('gcp.vertex.agent.event_id', function_response_event.id) - span.set_attribute( - 'gcp.vertex.agent.tool_response', - _safe_json_serialize(tool_response), - ) + if _should_add_request_response_to_spans(): + span.set_attribute( + 'gcp.vertex.agent.tool_response', + _safe_json_serialize(tool_response), + ) + else: + span.set_attribute('gcp.vertex.agent.tool_response', {}) def trace_merged_tool_calls( @@ -200,10 +210,13 @@ def trace_merged_tool_calls( except Exception: # pylint: disable=broad-exception-caught function_response_event_json = '' - span.set_attribute( - 'gcp.vertex.agent.tool_response', - function_response_event_json, - ) + if _should_add_request_response_to_spans(): + span.set_attribute( + 'gcp.vertex.agent.tool_response', + function_response_event_json, + ) + else: + span.set_attribute('gcp.vertex.agent.tool_response', {}) # Setting empty llm request and response (as UI expect these) while not # applicable for tool_response. span.set_attribute('gcp.vertex.agent.llm_request', '{}') @@ -243,10 +256,13 @@ def trace_call_llm( ) span.set_attribute('gcp.vertex.agent.event_id', event_id) # Consider removing once GenAI SDK provides a way to record this info. - span.set_attribute( - 'gcp.vertex.agent.llm_request', - _safe_json_serialize(_build_llm_request_for_trace(llm_request)), - ) + if _should_add_request_response_to_spans(): + span.set_attribute( + 'gcp.vertex.agent.llm_request', + _safe_json_serialize(_build_llm_request_for_trace(llm_request)), + ) + else: + span.set_attribute('gcp.vertex.agent.llm_request', {}) # Consider removing once GenAI SDK provides a way to record this info. if llm_request.config: if llm_request.config.top_p: @@ -265,10 +281,13 @@ def trace_call_llm( except Exception: # pylint: disable=broad-exception-caught llm_response_json = '' - span.set_attribute( - 'gcp.vertex.agent.llm_response', - llm_response_json, - ) + if _should_add_request_response_to_spans(): + span.set_attribute( + 'gcp.vertex.agent.llm_response', + llm_response_json, + ) + else: + span.set_attribute('gcp.vertex.agent.llm_response', {}) if llm_response.usage_metadata is not None: span.set_attribute( @@ -309,15 +328,18 @@ def trace_send_data( span.set_attribute('gcp.vertex.agent.event_id', event_id) # Once instrumentation is added to the GenAI SDK, consider whether this # information still needs to be recorded by the Agent Development Kit. - span.set_attribute( - 'gcp.vertex.agent.data', - _safe_json_serialize([ - types.Content(role=content.role, parts=content.parts).model_dump( - exclude_none=True - ) - for content in data - ]), - ) + if _should_add_request_response_to_spans(): + span.set_attribute( + 'gcp.vertex.agent.data', + _safe_json_serialize([ + types.Content(role=content.role, parts=content.parts).model_dump( + exclude_none=True + ) + for content in data + ]), + ) + else: + span.set_attribute('gcp.vertex.agent.data', {}) def _build_llm_request_for_trace(llm_request: LlmRequest) -> dict[str, Any]: @@ -350,3 +372,14 @@ def _build_llm_request_for_trace(llm_request: LlmRequest) -> dict[str, Any]: ) ) return result + + +# Defaults to true for now to preserve backward compatibility. +# Once prompt and response logging is well established in ADK, we might start +# a deprecation of request/response content in spans by switching the default +# to false. +def _should_add_request_response_to_spans() -> bool: + disabled_via_env_var = os.getenv( + ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS, 'true' + ).lower() in ('false', '0') + return not disabled_via_env_var diff --git a/src/google/adk/tools/_gemini_schema_util.py b/src/google/adk/tools/_gemini_schema_util.py index c8776343..df76f3c8 100644 --- a/src/google/adk/tools/_gemini_schema_util.py +++ b/src/google/adk/tools/_gemini_schema_util.py @@ -142,7 +142,9 @@ def _sanitize_schema_formats_for_gemini( ) -> dict[str, Any]: """Filters the schema to only include fields that are supported by JSONSchema.""" supported_fields: set[str] = set(_ExtendedJSONSchema.model_fields.keys()) - schema_field_names: set[str] = {"items"} # 'additional_properties' to come + # Gemini rejects schemas that include `additionalProperties`, so drop it. + supported_fields.discard("additional_properties") + schema_field_names: set[str] = {"items"} list_schema_field_names: set[str] = { "any_of", # 'one_of', 'all_of', 'not' to come } diff --git a/src/google/adk/tools/mcp_tool/mcp_tool.py b/src/google/adk/tools/mcp_tool/mcp_tool.py index 94a7e63f..545f81ee 100644 --- a/src/google/adk/tools/mcp_tool/mcp_tool.py +++ b/src/google/adk/tools/mcp_tool/mcp_tool.py @@ -17,8 +17,10 @@ from __future__ import annotations import base64 import inspect import logging +import sys from typing import Any from typing import Callable +from typing import Dict from typing import Optional from typing import Union import warnings @@ -27,6 +29,7 @@ from fastapi.openapi.models import APIKeyIn from google.genai.types import FunctionDeclaration from typing_extensions import override +from ...agents.readonly_context import ReadonlyContext from .._gemini_schema_util import _to_gemini_schema from .mcp_session_manager import MCPSessionManager from .mcp_session_manager import retry_on_closed_resource @@ -36,8 +39,6 @@ from .mcp_session_manager import retry_on_closed_resource try: from mcp.types import Tool as McpBaseTool except ImportError as e: - import sys - if sys.version_info < (3, 10): raise ImportError( "MCP Tool requires Python 3.10 or above. Please upgrade your Python" @@ -75,6 +76,9 @@ class McpTool(BaseAuthenticatedTool): auth_scheme: Optional[AuthScheme] = None, auth_credential: Optional[AuthCredential] = None, require_confirmation: Union[bool, Callable[..., bool]] = False, + header_provider: Optional[ + Callable[[ReadonlyContext], Dict[str, str]] + ] = None, ): """Initializes an MCPTool. @@ -106,6 +110,7 @@ class McpTool(BaseAuthenticatedTool): self._mcp_tool = mcp_tool self._mcp_session_manager = mcp_session_manager self._require_confirmation = require_confirmation + self._header_provider = header_provider @override def _get_declaration(self) -> FunctionDeclaration: @@ -192,10 +197,24 @@ class McpTool(BaseAuthenticatedTool): Any: The response from the tool. """ # Extract headers from credential for session pooling - headers = await self._get_headers(tool_context, credential) + auth_headers = await self._get_headers(tool_context, credential) + dynamic_headers = None + if self._header_provider: + dynamic_headers = self._header_provider( + ReadonlyContext(tool_context._invocation_context) + ) + + headers: Dict[str, str] = {} + if auth_headers: + headers.update(auth_headers) + if dynamic_headers: + headers.update(dynamic_headers) + final_headers = headers if headers else None # Get the session from the session manager - session = await self._mcp_session_manager.create_session(headers=headers) + session = await self._mcp_session_manager.create_session( + headers=final_headers + ) response = await session.call_tool(self._mcp_tool.name, arguments=args) return response diff --git a/src/google/adk/tools/mcp_tool/mcp_toolset.py b/src/google/adk/tools/mcp_tool/mcp_toolset.py index 42b1137e..06c57a45 100644 --- a/src/google/adk/tools/mcp_tool/mcp_toolset.py +++ b/src/google/adk/tools/mcp_tool/mcp_toolset.py @@ -16,6 +16,8 @@ from __future__ import annotations import logging import sys +from typing import Any +from typing import AsyncIterator from typing import Callable from typing import Dict from typing import List @@ -107,6 +109,9 @@ class McpToolset(BaseToolset): auth_scheme: Optional[AuthScheme] = None, auth_credential: Optional[AuthCredential] = None, require_confirmation: Union[bool, Callable[..., bool]] = False, + header_provider: Optional[ + Callable[[ReadonlyContext], Dict[str, str]] + ] = None, ): """Initializes the MCPToolset. @@ -130,6 +135,8 @@ class McpToolset(BaseToolset): require_confirmation: Whether tools in this toolset require confirmation. Can be a single boolean or a callable to apply to all tools. + header_provider: A callable that takes a ReadonlyContext and returns a + dictionary of headers to be used for the MCP session. """ super().__init__(tool_filter=tool_filter, tool_name_prefix=tool_name_prefix) @@ -138,6 +145,7 @@ class McpToolset(BaseToolset): self._connection_params = connection_params self._errlog = errlog + self._header_provider = header_provider # Create the session manager that will handle the MCP connection self._mcp_session_manager = MCPSessionManager( @@ -162,8 +170,13 @@ class McpToolset(BaseToolset): Returns: List[BaseTool]: A list of tools available under the specified context. """ + headers = ( + self._header_provider(readonly_context) + if self._header_provider and readonly_context + else None + ) # Get session from session manager - session = await self._mcp_session_manager.create_session() + session = await self._mcp_session_manager.create_session(headers=headers) # Fetch available tools from the MCP server tools_response: ListToolsResult = await session.list_tools() @@ -177,6 +190,7 @@ class McpToolset(BaseToolset): auth_scheme=self._auth_scheme, auth_credential=self._auth_credential, require_confirmation=self._require_confirmation, + header_provider=self._header_provider, ) if self._is_tool_selected(mcp_tool, readonly_context): diff --git a/tests/integration/fixture/hello_world_agent_async/__init__.py b/tests/integration/fixture/hello_world_agent_async/__init__.py new file mode 100644 index 00000000..c48963cd --- /dev/null +++ b/tests/integration/fixture/hello_world_agent_async/__init__.py @@ -0,0 +1,15 @@ +# 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 . import agent diff --git a/tests/integration/fixture/hello_world_agent_async/agent.py b/tests/integration/fixture/hello_world_agent_async/agent.py new file mode 100644 index 00000000..b105065c --- /dev/null +++ b/tests/integration/fixture/hello_world_agent_async/agent.py @@ -0,0 +1,104 @@ +# 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. + +# Hello world agent from agent 1.0 revised to be defined with get_agent_async +# instead of root_agent - https://colab.sandbox.google.com/drive/1Zq-nqmgK0nCERCv8jKIaoeTTgbNn6oSo?resourcekey=0-GYaz9pFT4wY8CI8Cvjy5GA#scrollTo=u3X3XwDOaCv9 +import contextlib +import random +from typing import Optional + +from google.adk import Agent +from google.adk.agents import llm_agent +from google.genai import types + + +def roll_die(sides: int) -> int: + """Roll a die and return the rolled result. + + Args: + sides: The integer number of sides the die has. + + Returns: + An integer of the result of rolling the die. + """ + return random.randint(1, sides) + + +def check_prime(nums: list[int]) -> list[str]: + """Check if a given list of numbers are prime. + + Args: + nums: The list of numbers to check. + + Returns: + A str indicating which number is prime. + """ + primes = set() + for number in nums: + number = int(number) + if number <= 1: + continue + is_prime = True + for i in range(2, int(number**0.5) + 1): + if number % i == 0: + is_prime = False + break + if is_prime: + primes.add(number) + return ( + 'No prime numbers found.' + if not primes + else f"{', '.join(str(num) for num in primes)} are prime numbers." + ) + + +async def get_agent_async() -> ( + tuple[llm_agent.LlmAgent, Optional[contextlib.AsyncExitStack]] +): + """Returns the root agent.""" + root_agent = Agent( + model='gemini-2.0-flash-001', + name='data_processing_agent', + instruction=""" + You roll dice and answer questions about the outcome of the dice rolls. + You can roll dice of different sizes. + You can use multiple tools in parallel by calling functions in parallel(in one request and in one round). + The only things you do are roll dice for the user and discuss the outcomes. + It is ok to discuss previous dice roles, and comment on the dice rolls. + When you are asked to roll a die, you must call the roll_die tool with the number of sides. Be sure to pass in an integer. Do not pass in a string. + You should never roll a die on your own. + When checking prime numbers, call the check_prime tool with a list of integers. Be sure to pass in a list of integers. You should never pass in a string. + You should not check prime numbers before calling the tool. + When you are asked to roll a die and check prime numbers, you should always make the following two function calls: + 1. You should first call the roll_die tool to get a roll. Wait for the function response before calling the check_prime tool. + 2. After you get the function response from roll_die tool, you should call the check_prime tool with the roll_die result. + 2.1 If user asks you to check primes based on previous rolls, make sure you include the previous rolls in the list. + 3. When you respond, you must include the roll_die result from step 1. + You should always perform the previous 3 steps when asking for a roll and checking prime numbers. + You should not rely on the previous history on prime results. + """, + tools=[ + roll_die, + check_prime, + ], + generate_content_config=types.GenerateContentConfig( + safety_settings=[ + types.SafetySetting( # avoid false alarm about rolling dice. + category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, + threshold=types.HarmBlockThreshold.OFF, + ), + ] + ), + ) + return root_agent, None diff --git a/tests/integration/fixture/hello_world_agent_async/roll_die.test.json b/tests/integration/fixture/hello_world_agent_async/roll_die.test.json new file mode 100644 index 00000000..7e787d40 --- /dev/null +++ b/tests/integration/fixture/hello_world_agent_async/roll_die.test.json @@ -0,0 +1,55 @@ +{ + "eval_set_id": "56540925-a5ff-49fe-a4e1-589fe78066f2", + "name": "56540925-a5ff-49fe-a4e1-589fe78066f2", + "description": null, + "eval_cases": [ + { + "eval_id": "tests/integration/fixture/hello_world_agent_async/roll_die.test.json", + "conversation": [ + { + "invocation_id": "b01f67f0-9f23-44d6-bbe4-36ea235cb9fb", + "user_content": { + "parts": [ + { + "video_metadata": null, + "thought": null, + "code_execution_result": null, + "executable_code": null, + "file_data": null, + "function_call": null, + "function_response": null, + "inline_data": null, + "text": "Hi who are you?" + } + ], + "role": "user" + }, + "final_response": { + "parts": [ + { + "video_metadata": null, + "thought": null, + "code_execution_result": null, + "executable_code": null, + "file_data": null, + "function_call": null, + "function_response": null, + "inline_data": null, + "text": "I am a data processing agent. I can roll dice and check if the results are prime numbers. What would you like me to do? \n" + } + ], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 1747341775.8937013 + } + ], + "session_input": null, + "creation_timestamp": 1747341775.8937826 + } + ], + "creation_timestamp": 1747341775.8937957 +} \ No newline at end of file diff --git a/tests/integration/fixture/hello_world_agent_async/test_config.json b/tests/integration/fixture/hello_world_agent_async/test_config.json new file mode 100644 index 00000000..c7fba6a4 --- /dev/null +++ b/tests/integration/fixture/hello_world_agent_async/test_config.json @@ -0,0 +1,6 @@ +{ + "criteria": { + "tool_trajectory_avg_score": 1.0, + "response_match_score": 0.5 + } +} diff --git a/tests/integration/test_single_agent.py b/tests/integration/test_single_agent.py index 183005ed..49f3b7ba 100644 --- a/tests/integration/test_single_agent.py +++ b/tests/integration/test_single_agent.py @@ -23,3 +23,14 @@ async def test_eval_agent(): eval_dataset_file_path_or_dir="tests/integration/fixture/home_automation_agent/simple_test.test.json", num_runs=4, ) + + +@pytest.mark.asyncio +async def test_eval_agent_async(): + await AgentEvaluator.evaluate( + agent_module="tests.integration.fixture.hello_world_agent_async", + eval_dataset_file_path_or_dir=( + "tests/integration/fixture/hello_world_agent_async/roll_die.test.json" + ), + num_runs=4, + ) diff --git a/tests/unittests/a2a/converters/test_request_converter.py b/tests/unittests/a2a/converters/test_request_converter.py index 115b2312..360afbd3 100644 --- a/tests/unittests/a2a/converters/test_request_converter.py +++ b/tests/unittests/a2a/converters/test_request_converter.py @@ -27,7 +27,7 @@ pytestmark = pytest.mark.skipif( try: from a2a.server.agent_execution import RequestContext from google.adk.a2a.converters.request_converter import _get_user_id - from google.adk.a2a.converters.request_converter import convert_a2a_request_to_adk_run_args + from google.adk.a2a.converters.request_converter import convert_a2a_request_to_agent_run_request from google.adk.runners import RunConfig from google.genai import types as genai_types except ImportError as e: @@ -143,11 +143,11 @@ class TestGetUserId: assert result == "A2A_USER_None" -class TestConvertA2aRequestToAdkRunArgs: - """Test cases for convert_a2a_request_to_adk_run_args function.""" +class TestConvertA2aRequestToAgentRunRequest: + """Test cases for convert_a2a_request_to_agent_run_request function.""" def test_convert_a2a_request_basic(self): - """Test basic conversion of A2A request to ADK run args.""" + """Test basic conversion of A2A request to ADK AgentRunRequest.""" # Arrange mock_part1 = Mock() mock_part2 = Mock() @@ -173,16 +173,18 @@ class TestConvertA2aRequestToAdkRunArgs: mock_convert_part.side_effect = [mock_genai_part1, mock_genai_part2] # Act - result = convert_a2a_request_to_adk_run_args(request, mock_convert_part) + result = convert_a2a_request_to_agent_run_request( + request, mock_convert_part + ) # Assert assert result is not None - assert result["user_id"] == "test_user" - assert result["session_id"] == "test_context_123" - assert isinstance(result["new_message"], genai_types.Content) - assert result["new_message"].role == "user" - assert result["new_message"].parts == [mock_genai_part1, mock_genai_part2] - assert isinstance(result["run_config"], RunConfig) + assert result.user_id == "test_user" + assert result.session_id == "test_context_123" + assert isinstance(result.new_message, genai_types.Content) + assert result.new_message.role == "user" + assert result.new_message.parts == [mock_genai_part1, mock_genai_part2] + assert isinstance(result.run_config, RunConfig) # Verify calls assert mock_convert_part.call_count == 2 @@ -197,7 +199,7 @@ class TestConvertA2aRequestToAdkRunArgs: # Act & Assert with pytest.raises(ValueError, match="Request message cannot be None"): - convert_a2a_request_to_adk_run_args(request) + convert_a2a_request_to_agent_run_request(request) def test_convert_a2a_request_empty_parts(self): """Test conversion with empty parts list.""" @@ -212,16 +214,18 @@ class TestConvertA2aRequestToAdkRunArgs: request.call_context = None # Act - result = convert_a2a_request_to_adk_run_args(request, mock_convert_part) + result = convert_a2a_request_to_agent_run_request( + request, mock_convert_part + ) # Assert assert result is not None - assert result["user_id"] == "A2A_USER_test_context_123" - assert result["session_id"] == "test_context_123" - assert isinstance(result["new_message"], genai_types.Content) - assert result["new_message"].role == "user" - assert result["new_message"].parts == [] - assert isinstance(result["run_config"], RunConfig) + assert result.user_id == "A2A_USER_test_context_123" + assert result.session_id == "test_context_123" + assert isinstance(result.new_message, genai_types.Content) + assert result.new_message.role == "user" + assert result.new_message.parts == [] + assert isinstance(result.run_config, RunConfig) # Verify convert_part wasn't called mock_convert_part.assert_not_called() @@ -244,16 +248,18 @@ class TestConvertA2aRequestToAdkRunArgs: mock_convert_part.return_value = mock_genai_part # Act - result = convert_a2a_request_to_adk_run_args(request, mock_convert_part) + result = convert_a2a_request_to_agent_run_request( + request, mock_convert_part + ) # Assert assert result is not None - assert result["user_id"] == "A2A_USER_None" - assert result["session_id"] is None - assert isinstance(result["new_message"], genai_types.Content) - assert result["new_message"].role == "user" - assert result["new_message"].parts == [mock_genai_part] - assert isinstance(result["run_config"], RunConfig) + assert result.user_id == "A2A_USER_None" + assert result.session_id is None + assert isinstance(result.new_message, genai_types.Content) + assert result.new_message.role == "user" + assert result.new_message.parts == [mock_genai_part] + assert isinstance(result.run_config, RunConfig) def test_convert_a2a_request_no_auth(self): """Test conversion when no authentication is available.""" @@ -273,16 +279,18 @@ class TestConvertA2aRequestToAdkRunArgs: mock_convert_part.return_value = mock_genai_part # Act - result = convert_a2a_request_to_adk_run_args(request, mock_convert_part) + result = convert_a2a_request_to_agent_run_request( + request, mock_convert_part + ) # Assert assert result is not None - assert result["user_id"] == "A2A_USER_session_123" - assert result["session_id"] == "session_123" - assert isinstance(result["new_message"], genai_types.Content) - assert result["new_message"].role == "user" - assert result["new_message"].parts == [mock_genai_part] - assert isinstance(result["run_config"], RunConfig) + assert result.user_id == "A2A_USER_session_123" + assert result.session_id == "session_123" + assert isinstance(result.new_message, genai_types.Content) + assert result.new_message.role == "user" + assert result.new_message.parts == [mock_genai_part] + assert isinstance(result.run_config, RunConfig) class TestIntegration: @@ -312,16 +320,18 @@ class TestIntegration: mock_convert_part.return_value = mock_genai_part # Act - result = convert_a2a_request_to_adk_run_args(request, mock_convert_part) + result = convert_a2a_request_to_agent_run_request( + request, mock_convert_part + ) # Assert assert result is not None - assert result["user_id"] == "auth_user" # Should use authenticated user - assert result["session_id"] == "mysession" - assert isinstance(result["new_message"], genai_types.Content) - assert result["new_message"].role == "user" - assert result["new_message"].parts == [mock_genai_part] - assert isinstance(result["run_config"], RunConfig) + assert result.user_id == "auth_user" # Should use authenticated user + assert result.session_id == "mysession" + assert isinstance(result.new_message, genai_types.Content) + assert result.new_message.role == "user" + assert result.new_message.parts == [mock_genai_part] + assert isinstance(result.run_config, RunConfig) def test_end_to_end_conversion_with_fallback_user(self): """Test end-to-end conversion with fallback user ID.""" @@ -341,15 +351,17 @@ class TestIntegration: mock_convert_part.return_value = mock_genai_part # Act - result = convert_a2a_request_to_adk_run_args(request, mock_convert_part) + result = convert_a2a_request_to_agent_run_request( + request, mock_convert_part + ) # Assert assert result is not None assert ( - result["user_id"] == "A2A_USER_test_session_456" + result.user_id == "A2A_USER_test_session_456" ) # Should fallback to context ID - assert result["session_id"] == "test_session_456" - assert isinstance(result["new_message"], genai_types.Content) - assert result["new_message"].role == "user" - assert result["new_message"].parts == [mock_genai_part] - assert isinstance(result["run_config"], RunConfig) + assert result.session_id == "test_session_456" + assert isinstance(result.new_message, genai_types.Content) + assert result.new_message.role == "user" + assert result.new_message.parts == [mock_genai_part] + assert isinstance(result.run_config, RunConfig) diff --git a/tests/unittests/a2a/executor/test_a2a_agent_executor.py b/tests/unittests/a2a/executor/test_a2a_agent_executor.py index 11d6b3a7..4bcc7a91 100644 --- a/tests/unittests/a2a/executor/test_a2a_agent_executor.py +++ b/tests/unittests/a2a/executor/test_a2a_agent_executor.py @@ -31,10 +31,13 @@ try: from a2a.types import Message from a2a.types import TaskState from a2a.types import TextPart + from google.adk.a2a.converters.request_converter import AgentRunRequest from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutor from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutorConfig from google.adk.events.event import Event + from google.adk.runners import RunConfig from google.adk.runners import Runner + from google.genai.types import Content except ImportError as e: if sys.version_info < (3, 10): # Imports are not needed since tests will be skipped due to pytestmark. @@ -58,9 +61,13 @@ class TestA2aAgentExecutor: self.mock_a2a_part_converter = Mock() self.mock_gen_ai_part_converter = Mock() + self.mock_request_converter = Mock() + self.mock_event_converter = Mock() self.mock_config = A2aAgentExecutorConfig( a2a_part_converter=self.mock_a2a_part_converter, gen_ai_part_converter=self.mock_gen_ai_part_converter, + request_converter=self.mock_request_converter, + event_converter=self.mock_event_converter, ) self.executor = A2aAgentExecutor( runner=self.mock_runner, config=self.mock_config @@ -84,71 +91,73 @@ class TestA2aAgentExecutor: async def test_execute_success_new_task(self): """Test successful execution of a new task.""" # Setup - with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_a2a_request_to_adk_run_args" - ) as mock_convert: - mock_convert.return_value = { - "user_id": "test-user", - "session_id": "test-session", - "new_message": Mock(), - "run_config": Mock(), - } + self.mock_request_converter.return_value = AgentRunRequest( + user_id="test-user", + session_id="test-session", + new_message=Mock(spec=Content), + run_config=Mock(spec=RunConfig), + ) + # Mock session service + mock_session = Mock() + mock_session.id = "test-session" + self.mock_runner.session_service.get_session = AsyncMock( + return_value=mock_session + ) - # Mock session service - mock_session = Mock() - mock_session.id = "test-session" - self.mock_runner.session_service.get_session = AsyncMock( - return_value=mock_session - ) + # Mock invocation context + mock_invocation_context = Mock() + self.mock_runner._new_invocation_context.return_value = ( + mock_invocation_context + ) - # Mock invocation context - mock_invocation_context = Mock() - self.mock_runner._new_invocation_context.return_value = ( - mock_invocation_context - ) + # Mock agent run with proper async generator + mock_event = Mock(spec=Event) - # Mock agent run with proper async generator - mock_event = Mock(spec=Event) + # Configure run_async to return the async generator when awaited + async def mock_run_async(**kwargs): + async for item in self._create_async_generator([mock_event]): + yield item - # Configure run_async to return the async generator when awaited - async def mock_run_async(**kwargs): - async for item in self._create_async_generator([mock_event]): - yield item + self.mock_runner.run_async = mock_run_async + self.mock_event_converter.return_value = [] - self.mock_runner.run_async = mock_run_async + # Execute + await self.executor.execute(self.mock_context, self.mock_event_queue) - with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_event_to_a2a_events" - ) as mock_convert_events: - mock_convert_events.return_value = [] + # Verify request converter was called with proper arguments + self.mock_request_converter.assert_called_once_with( + self.mock_context, self.mock_a2a_part_converter + ) - # Execute - await self.executor.execute(self.mock_context, self.mock_event_queue) + # Verify event converter was called with proper arguments + self.mock_event_converter.assert_called_once_with( + mock_event, + mock_invocation_context, + self.mock_context.task_id, + self.mock_context.context_id, + self.mock_gen_ai_part_converter, + ) - # Verify task submitted event was enqueued - assert self.mock_event_queue.enqueue_event.call_count >= 3 - submitted_event = self.mock_event_queue.enqueue_event.call_args_list[0][ - 0 - ][0] - assert submitted_event.status.state == TaskState.submitted - assert submitted_event.final == False + # Verify task submitted event was enqueued + assert self.mock_event_queue.enqueue_event.call_count >= 3 + submitted_event = self.mock_event_queue.enqueue_event.call_args_list[0][0][ + 0 + ] + assert submitted_event.status.state == TaskState.submitted + assert submitted_event.final == False - # Verify working event was enqueued - working_event = self.mock_event_queue.enqueue_event.call_args_list[1][ - 0 - ][0] - assert working_event.status.state == TaskState.working - assert working_event.final == False + # Verify working event was enqueued + working_event = self.mock_event_queue.enqueue_event.call_args_list[1][0][0] + assert working_event.status.state == TaskState.working + assert working_event.final == False - # Verify final event was enqueued with proper message field - final_event = self.mock_event_queue.enqueue_event.call_args_list[-1][0][ - 0 - ] - assert final_event.final == True - # The TaskResultAggregator is created with default state (working), and since no messages - # are processed, it will publish a status event with the current state - assert hasattr(final_event.status, "message") - assert final_event.status.state == TaskState.working + # Verify final event was enqueued with proper message field + final_event = self.mock_event_queue.enqueue_event.call_args_list[-1][0][0] + assert final_event.final == True + # The TaskResultAggregator is created with default state (working), and since no messages + # are processed, it will publish a status event with the current state + assert hasattr(final_event.status, "message") + assert final_event.status.state == TaskState.working @pytest.mark.asyncio async def test_execute_no_message_error(self): @@ -164,73 +173,76 @@ class TestA2aAgentExecutor: self.mock_context.current_task = Mock() self.mock_context.task_id = "existing-task-id" - with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_a2a_request_to_adk_run_args" - ) as mock_convert: - mock_convert.return_value = { - "user_id": "test-user", - "session_id": "test-session", - "new_message": Mock(), - "run_config": Mock(), - } + self.mock_request_converter.return_value = AgentRunRequest( + user_id="test-user", + session_id="test-session", + new_message=Mock(spec=Content), + run_config=Mock(spec=RunConfig), + ) - # Mock session service - mock_session = Mock() - mock_session.id = "test-session" - self.mock_runner.session_service.get_session = AsyncMock( - return_value=mock_session - ) + # Mock session service + mock_session = Mock() + mock_session.id = "test-session" + self.mock_runner.session_service.get_session = AsyncMock( + return_value=mock_session + ) - # Mock invocation context - mock_invocation_context = Mock() - self.mock_runner._new_invocation_context.return_value = ( - mock_invocation_context - ) + # Mock invocation context + mock_invocation_context = Mock() + self.mock_runner._new_invocation_context.return_value = ( + mock_invocation_context + ) - # Mock agent run with proper async generator - mock_event = Mock(spec=Event) + # Mock agent run with proper async generator + mock_event = Mock(spec=Event) - # Configure run_async to return the async generator when awaited - async def mock_run_async(**kwargs): - async for item in self._create_async_generator([mock_event]): - yield item + # Configure run_async to return the async generator when awaited + async def mock_run_async(**kwargs): + async for item in self._create_async_generator([mock_event]): + yield item - self.mock_runner.run_async = mock_run_async + self.mock_runner.run_async = mock_run_async + self.mock_event_converter.return_value = [] - with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_event_to_a2a_events" - ) as mock_convert_events: - mock_convert_events.return_value = [] + # Execute + await self.executor.execute(self.mock_context, self.mock_event_queue) - # Execute - await self.executor.execute(self.mock_context, self.mock_event_queue) + # Verify request converter was called with proper arguments + self.mock_request_converter.assert_called_once_with( + self.mock_context, self.mock_a2a_part_converter + ) - # Verify no submitted event (first call should be working event) - working_event = self.mock_event_queue.enqueue_event.call_args_list[0][ - 0 - ][0] - assert working_event.status.state == TaskState.working - assert working_event.final == False + # Verify event converter was called with proper arguments + self.mock_event_converter.assert_called_once_with( + mock_event, + mock_invocation_context, + self.mock_context.task_id, + self.mock_context.context_id, + self.mock_gen_ai_part_converter, + ) - # Verify final event was enqueued with proper message field - final_event = self.mock_event_queue.enqueue_event.call_args_list[-1][0][ - 0 - ] - assert final_event.final == True - # The TaskResultAggregator is created with default state (working), and since no messages - # are processed, it will publish a status event with the current state - assert hasattr(final_event.status, "message") - assert final_event.status.state == TaskState.working + # Verify no submitted event (first call should be working event) + working_event = self.mock_event_queue.enqueue_event.call_args_list[0][0][0] + assert working_event.status.state == TaskState.working + assert working_event.final == False + + # Verify final event was enqueued with proper message field + final_event = self.mock_event_queue.enqueue_event.call_args_list[-1][0][0] + assert final_event.final == True + # The TaskResultAggregator is created with default state (working), and since no messages + # are processed, it will publish a status event with the current state + assert hasattr(final_event.status, "message") + assert final_event.status.state == TaskState.working @pytest.mark.asyncio async def test_prepare_session_new_session(self): """Test session preparation when session doesn't exist.""" - run_args = { - "user_id": "test-user", - "session_id": None, - "new_message": Mock(), - "run_config": Mock(), - } + run_args = AgentRunRequest( + user_id="test-user", + session_id=None, + new_message=Mock(spec=Content), + run_config=Mock(spec=RunConfig), + ) # Mock session service self.mock_runner.session_service.get_session = AsyncMock(return_value=None) @@ -247,18 +259,18 @@ class TestA2aAgentExecutor: # Verify session was created assert result == mock_session - assert run_args["session_id"] is not None + assert run_args.session_id is not None self.mock_runner.session_service.create_session.assert_called_once() @pytest.mark.asyncio async def test_prepare_session_existing_session(self): """Test session preparation when session exists.""" - run_args = { - "user_id": "test-user", - "session_id": "existing-session", - "new_message": Mock(), - "run_config": Mock(), - } + run_args = AgentRunRequest( + user_id="test-user", + session_id="existing-session", + new_message=Mock(spec=Content), + run_config=Mock(spec=RunConfig), + ) # Mock session service mock_session = Mock() @@ -397,63 +409,55 @@ class TestA2aAgentExecutor: executor = A2aAgentExecutor(runner=create_runner, config=self.mock_config) - with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_a2a_request_to_adk_run_args" - ) as mock_convert: - mock_convert.return_value = { - "user_id": "test-user", - "session_id": "test-session", - "new_message": Mock(), - "run_config": Mock(), - } + self.mock_request_converter.return_value = AgentRunRequest( + user_id="test-user", + session_id="test-session", + new_message=Mock(spec=Content), + run_config=Mock(spec=RunConfig), + ) - # Mock session service - mock_session = Mock() - mock_session.id = "test-session" - self.mock_runner.session_service.get_session = AsyncMock( - return_value=mock_session - ) + # Mock session service + mock_session = Mock() + mock_session.id = "test-session" + self.mock_runner.session_service.get_session = AsyncMock( + return_value=mock_session + ) - # Mock invocation context - mock_invocation_context = Mock() - self.mock_runner._new_invocation_context.return_value = ( - mock_invocation_context - ) + # Mock invocation context + mock_invocation_context = Mock() + self.mock_runner._new_invocation_context.return_value = ( + mock_invocation_context + ) - # Mock agent run with proper async generator - mock_event = Mock(spec=Event) + # Mock agent run with proper async generator + mock_event = Mock(spec=Event) - async def mock_run_async(**kwargs): - async for item in self._create_async_generator([mock_event]): - yield item + async def mock_run_async(**kwargs): + async for item in self._create_async_generator([mock_event]): + yield item - self.mock_runner.run_async = mock_run_async + self.mock_runner.run_async = mock_run_async - with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_event_to_a2a_events" - ) as mock_convert_events: - mock_convert_events.return_value = [] + self.mock_event_converter.return_value = [] - # Execute - await executor.execute(self.mock_context, self.mock_event_queue) + # Execute + await executor.execute(self.mock_context, self.mock_event_queue) - # Verify task submitted event was enqueued - assert self.mock_event_queue.enqueue_event.call_count >= 3 - submitted_event = self.mock_event_queue.enqueue_event.call_args_list[0][ - 0 - ][0] - assert submitted_event.status.state == TaskState.submitted - assert submitted_event.final == False + # Verify task submitted event was enqueued + assert self.mock_event_queue.enqueue_event.call_count >= 3 + submitted_event = self.mock_event_queue.enqueue_event.call_args_list[0][0][ + 0 + ] + assert submitted_event.status.state == TaskState.submitted + assert submitted_event.final == False - # Verify final event was enqueued with proper message field - final_event = self.mock_event_queue.enqueue_event.call_args_list[-1][0][ - 0 - ] - assert final_event.final == True - # The TaskResultAggregator is created with default state (working), and since no messages - # are processed, it will publish a status event with the current state - assert hasattr(final_event.status, "message") - assert final_event.status.state == TaskState.working + # Verify final event was enqueued with proper message field + final_event = self.mock_event_queue.enqueue_event.call_args_list[-1][0][0] + assert final_event.final == True + # The TaskResultAggregator is created with default state (working), and since no messages + # are processed, it will publish a status event with the current state + assert hasattr(final_event.status, "message") + assert final_event.status.state == TaskState.working @pytest.mark.asyncio async def test_execute_with_async_callable_runner(self): @@ -464,63 +468,55 @@ class TestA2aAgentExecutor: executor = A2aAgentExecutor(runner=create_runner, config=self.mock_config) - with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_a2a_request_to_adk_run_args" - ) as mock_convert: - mock_convert.return_value = { - "user_id": "test-user", - "session_id": "test-session", - "new_message": Mock(), - "run_config": Mock(), - } + self.mock_request_converter.return_value = AgentRunRequest( + user_id="test-user", + session_id="test-session", + new_message=Mock(spec=Content), + run_config=Mock(spec=RunConfig), + ) - # Mock session service - mock_session = Mock() - mock_session.id = "test-session" - self.mock_runner.session_service.get_session = AsyncMock( - return_value=mock_session - ) + # Mock session service + mock_session = Mock() + mock_session.id = "test-session" + self.mock_runner.session_service.get_session = AsyncMock( + return_value=mock_session + ) - # Mock invocation context - mock_invocation_context = Mock() - self.mock_runner._new_invocation_context.return_value = ( - mock_invocation_context - ) + # Mock invocation context + mock_invocation_context = Mock() + self.mock_runner._new_invocation_context.return_value = ( + mock_invocation_context + ) - # Mock agent run with proper async generator - mock_event = Mock(spec=Event) + # Mock agent run with proper async generator + mock_event = Mock(spec=Event) - async def mock_run_async(**kwargs): - async for item in self._create_async_generator([mock_event]): - yield item + async def mock_run_async(**kwargs): + async for item in self._create_async_generator([mock_event]): + yield item - self.mock_runner.run_async = mock_run_async + self.mock_runner.run_async = mock_run_async - with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_event_to_a2a_events" - ) as mock_convert_events: - mock_convert_events.return_value = [] + self.mock_event_converter.return_value = [] - # Execute - await executor.execute(self.mock_context, self.mock_event_queue) + # Execute + await executor.execute(self.mock_context, self.mock_event_queue) - # Verify task submitted event was enqueued - assert self.mock_event_queue.enqueue_event.call_count >= 3 - submitted_event = self.mock_event_queue.enqueue_event.call_args_list[0][ - 0 - ][0] - assert submitted_event.status.state == TaskState.submitted - assert submitted_event.final == False + # Verify task submitted event was enqueued + assert self.mock_event_queue.enqueue_event.call_count >= 3 + submitted_event = self.mock_event_queue.enqueue_event.call_args_list[0][0][ + 0 + ] + assert submitted_event.status.state == TaskState.submitted + assert submitted_event.final == False - # Verify final event was enqueued with proper message field - final_event = self.mock_event_queue.enqueue_event.call_args_list[-1][0][ - 0 - ] - assert final_event.final == True - # The TaskResultAggregator is created with default state (working), and since no messages - # are processed, it will publish a status event with the current state - assert hasattr(final_event.status, "message") - assert final_event.status.state == TaskState.working + # Verify final event was enqueued with proper message field + final_event = self.mock_event_queue.enqueue_event.call_args_list[-1][0][0] + assert final_event.final == True + # The TaskResultAggregator is created with default state (working), and since no messages + # are processed, it will publish a status event with the current state + assert hasattr(final_event.status, "message") + assert final_event.status.state == TaskState.working @pytest.mark.asyncio async def test_handle_request_integration(self): @@ -529,83 +525,75 @@ class TestA2aAgentExecutor: self.mock_context.task_id = "test-task-id" # Setup detailed mocks + self.mock_request_converter.return_value = AgentRunRequest( + user_id="test-user", + session_id="test-session", + new_message=Mock(spec=Content), + run_config=Mock(spec=RunConfig), + ) + + # Mock session service + mock_session = Mock() + mock_session.id = "test-session" + self.mock_runner.session_service.get_session = AsyncMock( + return_value=mock_session + ) + + # Mock invocation context + mock_invocation_context = Mock() + self.mock_runner._new_invocation_context.return_value = ( + mock_invocation_context + ) + + # Mock agent run with multiple events using proper async generator + mock_events = [Mock(spec=Event), Mock(spec=Event)] + + # Configure run_async to return the async generator when awaited + async def mock_run_async(**kwargs): + async for item in self._create_async_generator(mock_events): + yield item + + self.mock_runner.run_async = mock_run_async + + self.mock_event_converter.return_value = [Mock()] + with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_a2a_request_to_adk_run_args" - ) as mock_convert: - mock_convert.return_value = { - "user_id": "test-user", - "session_id": "test-session", - "new_message": Mock(), - "run_config": Mock(), - } + "google.adk.a2a.executor.a2a_agent_executor.TaskResultAggregator" + ) as mock_aggregator_class: + mock_aggregator = Mock() + mock_aggregator.task_state = TaskState.working + # Mock the task_status_message property to return None by default + mock_aggregator.task_status_message = None + mock_aggregator_class.return_value = mock_aggregator - # Mock session service - mock_session = Mock() - mock_session.id = "test-session" - self.mock_runner.session_service.get_session = AsyncMock( - return_value=mock_session + # Execute + await self.executor._handle_request( + self.mock_context, self.mock_event_queue ) - # Mock invocation context - mock_invocation_context = Mock() - self.mock_runner._new_invocation_context.return_value = ( - mock_invocation_context - ) + # Verify working event was enqueued + working_events = [ + call[0][0] + for call in self.mock_event_queue.enqueue_event.call_args_list + if hasattr(call[0][0], "status") + and call[0][0].status.state == TaskState.working + ] + assert len(working_events) >= 1 - # Mock agent run with multiple events using proper async generator - mock_events = [Mock(spec=Event), Mock(spec=Event)] + # Verify aggregator processed events + assert mock_aggregator.process_event.call_count == len(mock_events) - # Configure run_async to return the async generator when awaited - async def mock_run_async(**kwargs): - async for item in self._create_async_generator(mock_events): - yield item - - self.mock_runner.run_async = mock_run_async - - with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_event_to_a2a_events" - ) as mock_convert_events: - mock_convert_events.return_value = [Mock()] - - with patch( - "google.adk.a2a.executor.a2a_agent_executor.TaskResultAggregator" - ) as mock_aggregator_class: - mock_aggregator = Mock() - mock_aggregator.task_state = TaskState.working - # Mock the task_status_message property to return None by default - mock_aggregator.task_status_message = None - mock_aggregator_class.return_value = mock_aggregator - - # Execute - await self.executor._handle_request( - self.mock_context, self.mock_event_queue - ) - - # Verify working event was enqueued - working_events = [ - call[0][0] - for call in self.mock_event_queue.enqueue_event.call_args_list - if hasattr(call[0][0], "status") - and call[0][0].status.state == TaskState.working - ] - assert len(working_events) >= 1 - - # Verify aggregator processed events - assert mock_aggregator.process_event.call_count == len(mock_events) - - # Verify final event has message field from aggregator and state is completed when aggregator state is working - final_events = [ - call[0][0] - for call in self.mock_event_queue.enqueue_event.call_args_list - if hasattr(call[0][0], "final") and call[0][0].final == True - ] - assert len(final_events) >= 1 - final_event = final_events[-1] # Get the last final event - assert ( - final_event.status.message == mock_aggregator.task_status_message - ) - # When aggregator state is working but no message, final event should be working - assert final_event.status.state == TaskState.working + # Verify final event has message field from aggregator and state is completed when aggregator state is working + final_events = [ + call[0][0] + for call in self.mock_event_queue.enqueue_event.call_args_list + if hasattr(call[0][0], "final") and call[0][0].final == True + ] + assert len(final_events) >= 1 + final_event = final_events[-1] # Get the last final event + assert final_event.status.message == mock_aggregator.task_status_message + # When aggregator state is working but no message, final event should be working + assert final_event.status.state == TaskState.working @pytest.mark.asyncio async def test_cancel_with_task_id(self): @@ -637,31 +625,26 @@ class TestA2aAgentExecutor: None # Make sure it goes through submitted event creation ) - with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_a2a_request_to_adk_run_args" - ) as mock_convert: - mock_convert.side_effect = Exception("Test error") + self.mock_request_converter.side_effect = Exception("Test error") - # Execute (should not raise since we catch the exception) - await self.executor.execute(self.mock_context, self.mock_event_queue) + # Execute (should not raise since we catch the exception) + await self.executor.execute(self.mock_context, self.mock_event_queue) - # Verify both submitted and failure events were enqueued - # First call should be submitted event, last should be failure event - assert self.mock_event_queue.enqueue_event.call_count >= 2 + # Verify both submitted and failure events were enqueued + # First call should be submitted event, last should be failure event + assert self.mock_event_queue.enqueue_event.call_count >= 2 - # Check submitted event (first) - submitted_event = self.mock_event_queue.enqueue_event.call_args_list[0][ - 0 - ][0] - assert submitted_event.status.state == TaskState.submitted - assert submitted_event.final == False + # Check submitted event (first) + submitted_event = self.mock_event_queue.enqueue_event.call_args_list[0][0][ + 0 + ] + assert submitted_event.status.state == TaskState.submitted + assert submitted_event.final == False - # Check failure event (last) - failure_event = self.mock_event_queue.enqueue_event.call_args_list[-1][0][ - 0 - ] - assert failure_event.status.state == TaskState.failed - assert failure_event.final == True + # Check failure event (last) + failure_event = self.mock_event_queue.enqueue_event.call_args_list[-1][0][0] + assert failure_event.status.state == TaskState.failed + assert failure_event.final == True @pytest.mark.asyncio async def test_handle_request_with_aggregator_message(self): @@ -680,69 +663,63 @@ class TestA2aAgentExecutor: test_message.parts = [Mock(spec=TextPart)] # Setup detailed mocks + self.mock_request_converter.return_value = AgentRunRequest( + user_id="test-user", + session_id="test-session", + new_message=Mock(spec=Content), + run_config=Mock(spec=RunConfig), + ) + + # Mock session service + mock_session = Mock() + mock_session.id = "test-session" + self.mock_runner.session_service.get_session = AsyncMock( + return_value=mock_session + ) + + # Mock invocation context + mock_invocation_context = Mock() + self.mock_runner._new_invocation_context.return_value = ( + mock_invocation_context + ) + + # Mock agent run with multiple events using proper async generator + mock_events = [Mock(spec=Event), Mock(spec=Event)] + + # Configure run_async to return the async generator when awaited + async def mock_run_async(**kwargs): + async for item in self._create_async_generator(mock_events): + yield item + + self.mock_runner.run_async = mock_run_async + + self.mock_event_converter.return_value = [Mock()] + with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_a2a_request_to_adk_run_args" - ) as mock_convert: - mock_convert.return_value = { - "user_id": "test-user", - "session_id": "test-session", - "new_message": Mock(), - "run_config": Mock(), - } + "google.adk.a2a.executor.a2a_agent_executor.TaskResultAggregator" + ) as mock_aggregator_class: + mock_aggregator = Mock() + mock_aggregator.task_state = TaskState.completed + # Mock the task_status_message property to return a test message + mock_aggregator.task_status_message = test_message + mock_aggregator_class.return_value = mock_aggregator - # Mock session service - mock_session = Mock() - mock_session.id = "test-session" - self.mock_runner.session_service.get_session = AsyncMock( - return_value=mock_session + # Execute + await self.executor._handle_request( + self.mock_context, self.mock_event_queue ) - # Mock invocation context - mock_invocation_context = Mock() - self.mock_runner._new_invocation_context.return_value = ( - mock_invocation_context - ) - - # Mock agent run with multiple events using proper async generator - mock_events = [Mock(spec=Event), Mock(spec=Event)] - - # Configure run_async to return the async generator when awaited - async def mock_run_async(**kwargs): - async for item in self._create_async_generator(mock_events): - yield item - - self.mock_runner.run_async = mock_run_async - - with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_event_to_a2a_events" - ) as mock_convert_events: - mock_convert_events.return_value = [Mock()] - - with patch( - "google.adk.a2a.executor.a2a_agent_executor.TaskResultAggregator" - ) as mock_aggregator_class: - mock_aggregator = Mock() - mock_aggregator.task_state = TaskState.completed - # Mock the task_status_message property to return a test message - mock_aggregator.task_status_message = test_message - mock_aggregator_class.return_value = mock_aggregator - - # Execute - await self.executor._handle_request( - self.mock_context, self.mock_event_queue - ) - - # Verify final event has message field from aggregator - final_events = [ - call[0][0] - for call in self.mock_event_queue.enqueue_event.call_args_list - if hasattr(call[0][0], "final") and call[0][0].final == True - ] - assert len(final_events) >= 1 - final_event = final_events[-1] # Get the last final event - assert final_event.status.message == test_message - # When aggregator state is completed (not working), final event should be completed - assert final_event.status.state == TaskState.completed + # Verify final event has message field from aggregator + final_events = [ + call[0][0] + for call in self.mock_event_queue.enqueue_event.call_args_list + if hasattr(call[0][0], "final") and call[0][0].final == True + ] + assert len(final_events) >= 1 + final_event = final_events[-1] # Get the last final event + assert final_event.status.message == test_message + # When aggregator state is completed (not working), final event should be completed + assert final_event.status.state == TaskState.completed @pytest.mark.asyncio async def test_handle_request_with_non_working_aggregator_state(self): @@ -761,69 +738,63 @@ class TestA2aAgentExecutor: test_message.parts = [Mock(spec=TextPart)] # Setup detailed mocks + self.mock_request_converter.return_value = AgentRunRequest( + user_id="test-user", + session_id="test-session", + new_message=Mock(spec=Content), + run_config=Mock(spec=RunConfig), + ) + + # Mock session service + mock_session = Mock() + mock_session.id = "test-session" + self.mock_runner.session_service.get_session = AsyncMock( + return_value=mock_session + ) + + # Mock invocation context + mock_invocation_context = Mock() + self.mock_runner._new_invocation_context.return_value = ( + mock_invocation_context + ) + + # Mock agent run with multiple events using proper async generator + mock_events = [Mock(spec=Event), Mock(spec=Event)] + + # Configure run_async to return the async generator when awaited + async def mock_run_async(**kwargs): + async for item in self._create_async_generator(mock_events): + yield item + + self.mock_runner.run_async = mock_run_async + + self.mock_event_converter.return_value = [Mock()] + with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_a2a_request_to_adk_run_args" - ) as mock_convert: - mock_convert.return_value = { - "user_id": "test-user", - "session_id": "test-session", - "new_message": Mock(), - "run_config": Mock(), - } + "google.adk.a2a.executor.a2a_agent_executor.TaskResultAggregator" + ) as mock_aggregator_class: + mock_aggregator = Mock() + # Test with failed state - should preserve failed state + mock_aggregator.task_state = TaskState.failed + mock_aggregator.task_status_message = test_message + mock_aggregator_class.return_value = mock_aggregator - # Mock session service - mock_session = Mock() - mock_session.id = "test-session" - self.mock_runner.session_service.get_session = AsyncMock( - return_value=mock_session + # Execute + await self.executor._handle_request( + self.mock_context, self.mock_event_queue ) - # Mock invocation context - mock_invocation_context = Mock() - self.mock_runner._new_invocation_context.return_value = ( - mock_invocation_context - ) - - # Mock agent run with multiple events using proper async generator - mock_events = [Mock(spec=Event), Mock(spec=Event)] - - # Configure run_async to return the async generator when awaited - async def mock_run_async(**kwargs): - async for item in self._create_async_generator(mock_events): - yield item - - self.mock_runner.run_async = mock_run_async - - with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_event_to_a2a_events" - ) as mock_convert_events: - mock_convert_events.return_value = [Mock()] - - with patch( - "google.adk.a2a.executor.a2a_agent_executor.TaskResultAggregator" - ) as mock_aggregator_class: - mock_aggregator = Mock() - # Test with failed state - should preserve failed state - mock_aggregator.task_state = TaskState.failed - mock_aggregator.task_status_message = test_message - mock_aggregator_class.return_value = mock_aggregator - - # Execute - await self.executor._handle_request( - self.mock_context, self.mock_event_queue - ) - - # Verify final event preserves the non-working state - final_events = [ - call[0][0] - for call in self.mock_event_queue.enqueue_event.call_args_list - if hasattr(call[0][0], "final") and call[0][0].final == True - ] - assert len(final_events) >= 1 - final_event = final_events[-1] # Get the last final event - assert final_event.status.message == test_message - # When aggregator state is failed (not working), final event should keep failed state - assert final_event.status.state == TaskState.failed + # Verify final event preserves the non-working state + final_events = [ + call[0][0] + for call in self.mock_event_queue.enqueue_event.call_args_list + if hasattr(call[0][0], "final") and call[0][0].final == True + ] + assert len(final_events) >= 1 + final_event = final_events[-1] # Get the last final event + assert final_event.status.message == test_message + # When aggregator state is failed (not working), final event should keep failed state + assert final_event.status.state == TaskState.failed @pytest.mark.asyncio async def test_handle_request_with_working_state_publishes_artifact_and_completed( @@ -846,84 +817,77 @@ class TestA2aAgentExecutor: test_message.parts = [Part(root=TextPart(text="test content"))] # Setup detailed mocks + self.mock_request_converter.return_value = AgentRunRequest( + user_id="test-user", + session_id="test-session", + new_message=Mock(spec=Content), + run_config=Mock(spec=RunConfig), + ) + + # Mock session service + mock_session = Mock() + mock_session.id = "test-session" + self.mock_runner.session_service.get_session = AsyncMock( + return_value=mock_session + ) + + # Mock invocation context + mock_invocation_context = Mock() + self.mock_runner._new_invocation_context.return_value = ( + mock_invocation_context + ) + + # Mock agent run with multiple events using proper async generator + mock_events = [Mock(spec=Event), Mock(spec=Event)] + + # Configure run_async to return the async generator when awaited + async def mock_run_async(**kwargs): + async for item in self._create_async_generator(mock_events): + yield item + + self.mock_runner.run_async = mock_run_async + + self.mock_event_converter.return_value = [Mock()] + with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_a2a_request_to_adk_run_args" - ) as mock_convert: - mock_convert.return_value = { - "user_id": "test-user", - "session_id": "test-session", - "new_message": Mock(), - "run_config": Mock(), - } + "google.adk.a2a.executor.a2a_agent_executor.TaskResultAggregator" + ) as mock_aggregator_class: + mock_aggregator = Mock() + # Test with working state - should publish artifact update and completed status + mock_aggregator.task_state = TaskState.working + mock_aggregator.task_status_message = test_message + mock_aggregator_class.return_value = mock_aggregator - # Mock session service - mock_session = Mock() - mock_session.id = "test-session" - self.mock_runner.session_service.get_session = AsyncMock( - return_value=mock_session + # Execute + await self.executor._handle_request( + self.mock_context, self.mock_event_queue ) - # Mock invocation context - mock_invocation_context = Mock() - self.mock_runner._new_invocation_context.return_value = ( - mock_invocation_context - ) + # Verify artifact update event was published + artifact_events = [ + call[0][0] + for call in self.mock_event_queue.enqueue_event.call_args_list + if hasattr(call[0][0], "artifact") and call[0][0].last_chunk == True + ] + assert len(artifact_events) == 1 + artifact_event = artifact_events[0] + assert artifact_event.task_id == "test-task-id" + assert artifact_event.context_id == "test-context-id" + # Check that artifact parts correspond to message parts + assert len(artifact_event.artifact.parts) == len(test_message.parts) + assert artifact_event.artifact.parts == test_message.parts - # Mock agent run with multiple events using proper async generator - mock_events = [Mock(spec=Event), Mock(spec=Event)] - - # Configure run_async to return the async generator when awaited - async def mock_run_async(**kwargs): - async for item in self._create_async_generator(mock_events): - yield item - - self.mock_runner.run_async = mock_run_async - - with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_event_to_a2a_events" - ) as mock_convert_events: - mock_convert_events.return_value = [Mock()] - - with patch( - "google.adk.a2a.executor.a2a_agent_executor.TaskResultAggregator" - ) as mock_aggregator_class: - mock_aggregator = Mock() - # Test with working state - should publish artifact update and completed status - mock_aggregator.task_state = TaskState.working - mock_aggregator.task_status_message = test_message - mock_aggregator_class.return_value = mock_aggregator - - # Execute - await self.executor._handle_request( - self.mock_context, self.mock_event_queue - ) - - # Verify artifact update event was published - artifact_events = [ - call[0][0] - for call in self.mock_event_queue.enqueue_event.call_args_list - if hasattr(call[0][0], "artifact") - and call[0][0].last_chunk == True - ] - assert len(artifact_events) == 1 - artifact_event = artifact_events[0] - assert artifact_event.task_id == "test-task-id" - assert artifact_event.context_id == "test-context-id" - # Check that artifact parts correspond to message parts - assert len(artifact_event.artifact.parts) == len(test_message.parts) - assert artifact_event.artifact.parts == test_message.parts - - # Verify final status event was published with completed state - final_events = [ - call[0][0] - for call in self.mock_event_queue.enqueue_event.call_args_list - if hasattr(call[0][0], "final") and call[0][0].final == True - ] - assert len(final_events) >= 1 - final_event = final_events[-1] # Get the last final event - assert final_event.status.state == TaskState.completed - assert final_event.task_id == "test-task-id" - assert final_event.context_id == "test-context-id" + # Verify final status event was published with completed state + final_events = [ + call[0][0] + for call in self.mock_event_queue.enqueue_event.call_args_list + if hasattr(call[0][0], "final") and call[0][0].final == True + ] + assert len(final_events) >= 1 + final_event = final_events[-1] # Get the last final event + assert final_event.status.state == TaskState.completed + assert final_event.task_id == "test-task-id" + assert final_event.context_id == "test-context-id" @pytest.mark.asyncio async def test_handle_request_with_non_working_state_publishes_status_only( @@ -946,76 +910,69 @@ class TestA2aAgentExecutor: test_message.parts = [Part(root=TextPart(text="test content"))] # Setup detailed mocks + self.mock_request_converter.return_value = AgentRunRequest( + user_id="test-user", + session_id="test-session", + new_message=Mock(spec=Content), + run_config=Mock(spec=RunConfig), + ) + + # Mock session service + mock_session = Mock() + mock_session.id = "test-session" + self.mock_runner.session_service.get_session = AsyncMock( + return_value=mock_session + ) + + # Mock invocation context + mock_invocation_context = Mock() + self.mock_runner._new_invocation_context.return_value = ( + mock_invocation_context + ) + + # Mock agent run with multiple events using proper async generator + mock_events = [Mock(spec=Event), Mock(spec=Event)] + + # Configure run_async to return the async generator when awaited + async def mock_run_async(**kwargs): + async for item in self._create_async_generator(mock_events): + yield item + + self.mock_runner.run_async = mock_run_async + + self.mock_event_converter.return_value = [Mock()] + with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_a2a_request_to_adk_run_args" - ) as mock_convert: - mock_convert.return_value = { - "user_id": "test-user", - "session_id": "test-session", - "new_message": Mock(), - "run_config": Mock(), - } + "google.adk.a2a.executor.a2a_agent_executor.TaskResultAggregator" + ) as mock_aggregator_class: + mock_aggregator = Mock() + # Test with auth_required state - should publish only status event + mock_aggregator.task_state = TaskState.auth_required + mock_aggregator.task_status_message = test_message + mock_aggregator_class.return_value = mock_aggregator - # Mock session service - mock_session = Mock() - mock_session.id = "test-session" - self.mock_runner.session_service.get_session = AsyncMock( - return_value=mock_session + # Execute + await self.executor._handle_request( + self.mock_context, self.mock_event_queue ) - # Mock invocation context - mock_invocation_context = Mock() - self.mock_runner._new_invocation_context.return_value = ( - mock_invocation_context - ) + # Verify no artifact update event was published + artifact_events = [ + call[0][0] + for call in self.mock_event_queue.enqueue_event.call_args_list + if hasattr(call[0][0], "artifact") and call[0][0].last_chunk == True + ] + assert len(artifact_events) == 0 - # Mock agent run with multiple events using proper async generator - mock_events = [Mock(spec=Event), Mock(spec=Event)] - - # Configure run_async to return the async generator when awaited - async def mock_run_async(**kwargs): - async for item in self._create_async_generator(mock_events): - yield item - - self.mock_runner.run_async = mock_run_async - - with patch( - "google.adk.a2a.executor.a2a_agent_executor.convert_event_to_a2a_events" - ) as mock_convert_events: - mock_convert_events.return_value = [Mock()] - - with patch( - "google.adk.a2a.executor.a2a_agent_executor.TaskResultAggregator" - ) as mock_aggregator_class: - mock_aggregator = Mock() - # Test with auth_required state - should publish only status event - mock_aggregator.task_state = TaskState.auth_required - mock_aggregator.task_status_message = test_message - mock_aggregator_class.return_value = mock_aggregator - - # Execute - await self.executor._handle_request( - self.mock_context, self.mock_event_queue - ) - - # Verify no artifact update event was published - artifact_events = [ - call[0][0] - for call in self.mock_event_queue.enqueue_event.call_args_list - if hasattr(call[0][0], "artifact") - and call[0][0].last_chunk == True - ] - assert len(artifact_events) == 0 - - # Verify final status event was published with the actual state and message - final_events = [ - call[0][0] - for call in self.mock_event_queue.enqueue_event.call_args_list - if hasattr(call[0][0], "final") and call[0][0].final == True - ] - assert len(final_events) >= 1 - final_event = final_events[-1] # Get the last final event - assert final_event.status.state == TaskState.auth_required - assert final_event.status.message == test_message - assert final_event.task_id == "test-task-id" - assert final_event.context_id == "test-context-id" + # Verify final status event was published with the actual state and message + final_events = [ + call[0][0] + for call in self.mock_event_queue.enqueue_event.call_args_list + if hasattr(call[0][0], "final") and call[0][0].final == True + ] + assert len(final_events) >= 1 + final_event = final_events[-1] # Get the last final event + assert final_event.status.state == TaskState.auth_required + assert final_event.status.message == test_message + assert final_event.task_id == "test-task-id" + assert final_event.context_id == "test-context-id" diff --git a/tests/unittests/artifacts/test_artifact_service.py b/tests/unittests/artifacts/test_artifact_service.py index f9065959..920aad68 100644 --- a/tests/unittests/artifacts/test_artifact_service.py +++ b/tests/unittests/artifacts/test_artifact_service.py @@ -14,11 +14,14 @@ """Tests for the artifact service.""" +from datetime import datetime import enum from typing import Optional from typing import Union from unittest import mock +from unittest.mock import patch +from google.adk.artifacts.base_artifact_service import ArtifactVersion from google.adk.artifacts.gcs_artifact_service import GcsArtifactService from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService from google.genai import types @@ -26,6 +29,9 @@ import pytest Enum = enum.Enum +# Define a fixed datetime object to be returned by datetime.now() +FIXED_DATETIME = datetime(2025, 1, 1, 12, 0, 0) + class ArtifactServiceType(Enum): IN_MEMORY = "IN_MEMORY" @@ -195,6 +201,15 @@ async def test_save_load_delete(service_type): == artifact ) + # Attempt to load a version that doesn't exist + assert not await artifact_service.load_artifact( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename=filename, + version=3, + ) + await artifact_service.delete_artifact( app_name=app_name, user_id=user_id, @@ -322,3 +337,171 @@ async def test_list_keys_preserves_user_prefix(): # Should contain prefixed names and session file expected_keys = ["user:document.pdf", "user:image.png", "session_file.txt"] assert sorted(artifact_keys) == sorted(expected_keys) + + +@pytest.mark.asyncio +async def test_list_artifact_versions_and_get_artifact_version(): + """Tests listing artifact versions and getting a specific version.""" + artifact_service = InMemoryArtifactService() + app_name = "app0" + user_id = "user0" + session_id = "123" + filename = "filename" + versions = [ + types.Part.from_bytes( + data=i.to_bytes(2, byteorder="big"), mime_type="text/plain" + ) + for i in range(4) + ] + + with patch( + "google.adk.artifacts.base_artifact_service.datetime" + ) as mock_datetime: + mock_datetime.now.return_value = FIXED_DATETIME + + for i in range(4): + await artifact_service.save_artifact( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename=filename, + artifact=versions[i], + custom_metadata={"key": "value" + str(i)}, + ) + + artifact_versions = await artifact_service.list_artifact_versions( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename=filename, + ) + + expected_artifact_versions = [ + ArtifactVersion( + version=i, + canonical_uri=( + f"memory://apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{filename}/versions/{i}" + ), + custom_metadata={"key": "value" + str(i)}, + mime_type="text/plain", + create_time=FIXED_DATETIME.timestamp(), + ) + for i in range(4) + ] + assert artifact_versions == expected_artifact_versions + + # Get latest artifact version when version is not specified + assert ( + await artifact_service.get_artifact_version( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename=filename, + ) + == expected_artifact_versions[-1] + ) + + # Get artifact version by version number + assert ( + await artifact_service.get_artifact_version( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename=filename, + version=2, + ) + == expected_artifact_versions[2] + ) + + +@pytest.mark.asyncio +async def test_list_artifact_versions_with_user_prefix(): + """Tests listing artifact versions with user prefix.""" + artifact_service = InMemoryArtifactService() + app_name = "app0" + user_id = "user0" + session_id = "123" + user_scoped_filename = "user:document.pdf" + versions = [ + types.Part.from_bytes( + data=i.to_bytes(2, byteorder="big"), mime_type="text/plain" + ) + for i in range(4) + ] + + with patch( + "google.adk.artifacts.base_artifact_service.datetime" + ) as mock_datetime: + mock_datetime.now.return_value = FIXED_DATETIME + + for i in range(4): + # Save artifacts with "user:" prefix (cross-session artifacts) + await artifact_service.save_artifact( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename=user_scoped_filename, + artifact=versions[i], + custom_metadata={"key": "value" + str(i)}, + ) + + artifact_versions = await artifact_service.list_artifact_versions( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename=user_scoped_filename, + ) + + expected_artifact_versions = [ + ArtifactVersion( + version=i, + canonical_uri=( + f"memory://apps/{app_name}/users/{user_id}/artifacts/{user_scoped_filename}/versions/{i}" + ), + custom_metadata={"key": "value" + str(i)}, + mime_type="text/plain", + create_time=FIXED_DATETIME.timestamp(), + ) + for i in range(4) + ] + assert artifact_versions == expected_artifact_versions + + +@pytest.mark.asyncio +async def test_get_artifact_version_artifact_does_not_exist(): + """Tests getting an artifact version when artifact does not exist.""" + artifact_service = InMemoryArtifactService() + assert not await artifact_service.get_artifact_version( + app_name="test_app", + user_id="test_user", + session_id="session_id", + filename="filename", + ) + + +@pytest.mark.asyncio +async def test_get_artifact_version_out_of_index(): + """Tests loading an artifact with an out-of-index version.""" + artifact_service = InMemoryArtifactService() + app_name = "app0" + user_id = "user0" + session_id = "123" + filename = "filename" + artifact = types.Part.from_bytes(data=b"test_data", mime_type="text/plain") + + await artifact_service.save_artifact( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename=filename, + artifact=artifact, + ) + + # Attempt to get a version that doesn't exist + assert not await artifact_service.get_artifact_version( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename=filename, + version=3, + ) diff --git a/tests/unittests/artifacts/test_artifact_util.py b/tests/unittests/artifacts/test_artifact_util.py new file mode 100644 index 00000000..995bf015 --- /dev/null +++ b/tests/unittests/artifacts/test_artifact_util.py @@ -0,0 +1,109 @@ +# 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. + +"""Tests for artifact_util.""" + +from google.adk.artifacts import artifact_util +from google.genai import types +import pytest + + +def test_parse_session_scoped_artifact_uri(): + """Tests parsing a valid session-scoped artifact URI.""" + uri = "artifact://apps/app1/users/user1/sessions/session1/artifacts/file1/versions/123" + parsed = artifact_util.parse_artifact_uri(uri) + assert parsed is not None + assert parsed.app_name == "app1" + assert parsed.user_id == "user1" + assert parsed.session_id == "session1" + assert parsed.filename == "file1" + assert parsed.version == 123 + + +def test_parse_user_scoped_artifact_uri(): + """Tests parsing a valid user-scoped artifact URI.""" + uri = "artifact://apps/app2/users/user2/artifacts/file2/versions/456" + parsed = artifact_util.parse_artifact_uri(uri) + assert parsed is not None + assert parsed.app_name == "app2" + assert parsed.user_id == "user2" + assert parsed.session_id is None + assert parsed.filename == "file2" + assert parsed.version == 456 + + +@pytest.mark.parametrize( + "invalid_uri", + [ + "http://example.com", + "artifact://invalid", + "artifact://app1/user1/sessions/session1/artifacts/file1", + "artifact://apps/app1/users/user1/sessions/session1/artifacts/file1", + "artifact://apps/app1/users/user1/artifacts/file1", + ], +) +def test_parse_invalid_artifact_uri(invalid_uri): + """Tests parsing invalid artifact URIs.""" + assert artifact_util.parse_artifact_uri(invalid_uri) is None + + +def test_get_session_scoped_artifact_uri(): + """Tests constructing a session-scoped artifact URI.""" + uri = artifact_util.get_artifact_uri( + app_name="app1", + user_id="user1", + session_id="session1", + filename="file1", + version=123, + ) + assert ( + uri + == "artifact://apps/app1/users/user1/sessions/session1/artifacts/file1/versions/123" + ) + + +def test_get_user_scoped_artifact_uri(): + """Tests constructing a user-scoped artifact URI.""" + uri = artifact_util.get_artifact_uri( + app_name="app2", user_id="user2", filename="file2", version=456 + ) + assert uri == "artifact://apps/app2/users/user2/artifacts/file2/versions/456" + + +def test_is_artifact_ref_true(): + """Tests is_artifact_ref with a valid artifact reference.""" + artifact = types.Part( + file_data=types.FileData( + file_uri="artifact://apps/a/u/s/f/v/1", mime_type="text/plain" + ) + ) + assert artifact_util.is_artifact_ref(artifact) is True + + +@pytest.mark.parametrize( + "part", + [ + types.Part(text="hello"), + types.Part(inline_data=types.Blob(data=b"123", mime_type="text/plain")), + types.Part( + file_data=types.FileData( + file_uri="http://example.com", mime_type="text/plain" + ) + ), + types.Part(), + ], +) +def test_is_artifact_ref_false(part): + """Tests is_artifact_ref with non-reference parts.""" + assert artifact_util.is_artifact_ref(part) is False diff --git a/tests/unittests/flows/llm_flows/test_base_llm_flow.py b/tests/unittests/flows/llm_flows/test_base_llm_flow.py index 47b4b00f..81ef925a 100644 --- a/tests/unittests/flows/llm_flows/test_base_llm_flow.py +++ b/tests/unittests/flows/llm_flows/test_base_llm_flow.py @@ -14,13 +14,13 @@ """Unit tests for BaseLlmFlow toolset integration.""" -from typing import Optional +from unittest import mock from unittest.mock import AsyncMock -from google.adk.agents.callback_context import CallbackContext from google.adk.agents.llm_agent import Agent from google.adk.events.event import Event from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow +from google.adk.models.google_llm import Gemini from google.adk.models.llm_request import LlmRequest from google.adk.models.llm_response import LlmResponse from google.adk.plugins.base_plugin import BasePlugin @@ -95,7 +95,6 @@ async def test_preprocess_calls_toolset_process_llm_request(): async def test_preprocess_handles_mixed_tools_and_toolsets(): """Test that _preprocess_async properly handles both tools and toolsets.""" from google.adk.tools.base_tool import BaseTool - from google.adk.tools.function_tool import FunctionTool # Create a mock tool class _MockTool(BaseTool): @@ -200,6 +199,46 @@ async def test_preprocess_with_google_search_workaround(): assert {d.name for d in declarations} == {'_my_tool', 'google_search_agent'} +@pytest.mark.asyncio +async def test_preprocess_calls_convert_tool_union_to_tools(): + """Test that _preprocess_async calls _convert_tool_union_to_tools.""" + + class _MockTool: + process_llm_request = AsyncMock() + + mock_tool_instance = _MockTool() + + def _my_tool(sides: int) -> int: + """A simple tool.""" + return sides + + with mock.patch( + 'google.adk.agents.llm_agent._convert_tool_union_to_tools', + new_callable=AsyncMock, + ) as mock_convert: + mock_convert.return_value = [mock_tool_instance] + + model = Gemini(model='gemini-2') + agent = Agent( + name='test_agent', model=model, tools=[_my_tool, google_search] + ) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, user_content='test message' + ) + flow = BaseLlmFlowForTesting() + llm_request = LlmRequest(model='gemini-2') + + async for _ in flow._preprocess_async(invocation_context, llm_request): + pass + + mock_convert.assert_called_with( + google_search, + mock.ANY, # ReadonlyContext(invocation_context) + model, + True, # multiple_tools + ) + + # TODO(b/448114567): Remove the following # test_handle_after_model_callback_grounding tests once the workaround # is no longer needed. diff --git a/tests/unittests/flows/llm_flows/test_contents.py b/tests/unittests/flows/llm_flows/test_contents.py index 7283a0ce..9e77407b 100644 --- a/tests/unittests/flows/llm_flows/test_contents.py +++ b/tests/unittests/flows/llm_flows/test_contents.py @@ -324,6 +324,63 @@ async def test_confirmation_events_are_filtered(): ] +@pytest.mark.asyncio +async def test_rewind_events_are_filtered_out(): + """Test that events are filtered based on rewind action.""" + agent = Agent(model="gemini-2.5-flash", name="test_agent") + llm_request = LlmRequest(model="gemini-2.5-flash") + invocation_context = await testing_utils.create_invocation_context( + agent=agent + ) + + events = [ + Event( + invocation_id="inv1", + author="user", + content=types.UserContent("First message"), + ), + Event( + invocation_id="inv1", + author="test_agent", + content=types.ModelContent("First response"), + ), + Event( + invocation_id="inv2", + author="user", + content=types.UserContent("Second message"), + ), + Event( + invocation_id="inv2", + author="test_agent", + content=types.ModelContent("Second response"), + ), + Event( + invocation_id="rewind_inv", + author="test_agent", + actions=EventActions(rewind_before_invocation_id="inv2"), + ), + Event( + invocation_id="inv3", + author="user", + content=types.UserContent("Third message"), + ), + ] + invocation_context.session.events = events + + # Process the request + async for _ in contents.request_processor.run_async( + invocation_context, llm_request + ): + pass + + # Verify rewind correctly filters conversation history + assert llm_request.contents == [ + types.UserContent("First message"), + types.ModelContent("First response"), + types.UserContent("Third message"), + ] + + @pytest.mark.asyncio async def test_events_with_empty_content_are_skipped(): """Test that events with empty content (state-only changes) are skipped.""" diff --git a/tests/unittests/plugins/test_reflect_retry_tool_plugin.py b/tests/unittests/plugins/test_reflect_retry_tool_plugin.py index ac085020..1e15f338 100644 --- a/tests/unittests/plugins/test_reflect_retry_tool_plugin.py +++ b/tests/unittests/plugins/test_reflect_retry_tool_plugin.py @@ -16,10 +16,14 @@ from typing import Any from unittest import IsolatedAsyncioTestCase from unittest.mock import Mock +from google.adk.agents.llm_agent import LlmAgent from google.adk.plugins.reflect_retry_tool_plugin import REFLECT_AND_RETRY_RESPONSE_TYPE from google.adk.plugins.reflect_retry_tool_plugin import ReflectAndRetryToolPlugin from google.adk.tools.base_tool import BaseTool from google.adk.tools.tool_context import ToolContext +from google.genai import types + +from .. import testing_utils class MockTool(BaseTool): @@ -524,3 +528,56 @@ class TestReflectAndRetryToolPlugin(IsolatedAsyncioTestCase): result=custom_error, ) self.assertEqual(result4["retry_count"], 1) + + async def test_hallucinating_tool_name(self): + """Test that hallucinating tool name is handled correctly.""" + wrong_function_call = types.Part.from_function_call( + name="increase_by_one", args={"x": 1} + ) + correct_function_call = types.Part.from_function_call( + name="increase", args={"x": 1} + ) + responses: list[types.Content] = [ + wrong_function_call, + correct_function_call, + "response1", + ] + mock_model = testing_utils.MockModel.create(responses=responses) + + function_called = 0 + + def increase(x: int) -> int: + nonlocal function_called + function_called += 1 + return x + 1 + + agent = LlmAgent(name="root_agent", model=mock_model, tools=[increase]) + runner = testing_utils.TestInMemoryRunner( + agent=agent, plugins=[self.get_plugin()] + ) + + events = await runner.run_async_with_new_session("test") + + # Assert that the first event is a function call with the wrong name + assert events[0].content.parts[0].function_call.name == "increase_by_one" + + # Assert that the second event is a function response with the + # reflection_guidance + assert ( + events[1].content.parts[0].function_response.response["error_type"] + == "ValueError" + ) + assert ( + events[1].content.parts[0].function_response.response["retry_count"] + == 1 + ) + assert ( + "Wrong Function Name" + in events[1] + .content.parts[0] + .function_response.response["reflection_guidance"] + ) + + # Assert that the third event is a function call with the correct name + assert events[2].content.parts[0].function_call.name == "increase" + self.assertEqual(function_called, 1) diff --git a/tests/unittests/runners/test_runner_rewind.py b/tests/unittests/runners/test_runner_rewind.py new file mode 100644 index 00000000..ae325e5a --- /dev/null +++ b/tests/unittests/runners/test_runner_rewind.py @@ -0,0 +1,248 @@ +# 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. + +"""Tests for runner.rewind_async.""" + +from google.adk.agents.base_agent import BaseAgent +from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService +from google.adk.events.event import Event +from google.adk.events.event import EventActions +from google.adk.runners import Runner +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.genai import types +import pytest + + +class TestRunnerRewind: + """Tests for runner.rewind_async.""" + + runner: Runner + + def setup_method(self): + """Set up test fixtures.""" + root_agent = BaseAgent(name="test_agent") + session_service = InMemorySessionService() + artifact_service = InMemoryArtifactService() + self.runner = Runner( + app_name="test_app", + agent=root_agent, + session_service=session_service, + artifact_service=artifact_service, + ) + + @pytest.mark.asyncio + async def test_rewind_async_with_state_and_artifacts(self): + """Tests rewind_async rewinds state and artifacts.""" + runner = self.runner + user_id = "test_user" + session_id = "test_session" + + # 1. Setup session and initial artifacts + session = await runner.session_service.create_session( + app_name=runner.app_name, user_id=user_id, session_id=session_id + ) + + # invocation1 + await runner.artifact_service.save_artifact( + app_name=runner.app_name, + user_id=user_id, + session_id=session_id, + filename="f1", + artifact=types.Part.from_text(text="f1v0"), + ) + event1 = Event( + invocation_id="invocation1", + author="agent", + content=types.Content(parts=[types.Part.from_text(text="event1")]), + actions=EventActions( + state_delta={"k1": "v1"}, artifact_delta={"f1": 0} + ), + ) + await runner.session_service.append_event(session=session, event=event1) + + # invocation2 + await runner.artifact_service.save_artifact( + app_name=runner.app_name, + user_id=user_id, + session_id=session_id, + filename="f1", + artifact=types.Part.from_text(text="f1v1"), + ) + await runner.artifact_service.save_artifact( + app_name=runner.app_name, + user_id=user_id, + session_id=session_id, + filename="f2", + artifact=types.Part.from_text(text="f2v0"), + ) + event2 = Event( + invocation_id="invocation2", + author="agent", + content=types.Content(parts=[types.Part.from_text(text="event2")]), + actions=EventActions( + state_delta={"k1": "v2", "k2": "v2"}, + artifact_delta={"f1": 1, "f2": 0}, + ), + ) + await runner.session_service.append_event(session=session, event=event2) + + # invocation3 + event3 = Event( + invocation_id="invocation3", + author="agent", + content=types.Content(parts=[types.Part.from_text(text="event3")]), + actions=EventActions(state_delta={"k2": "v3"}), + ) + await runner.session_service.append_event(session=session, event=event3) + + session = await runner.session_service.get_session( + app_name=runner.app_name, user_id=user_id, session_id=session_id + ) + assert session.state == {"k1": "v2", "k2": "v3"} + assert await runner.artifact_service.load_artifact( + app_name=runner.app_name, + user_id=user_id, + session_id=session_id, + filename="f1", + ) == types.Part.from_text(text="f1v1") + assert await runner.artifact_service.load_artifact( + app_name=runner.app_name, + user_id=user_id, + session_id=session_id, + filename="f2", + ) == types.Part.from_text(text="f2v0") + + # 2. Rewind before invocation2 + await runner.rewind_async( + user_id=user_id, + session_id=session_id, + rewind_before_invocation_id="invocation2", + ) + + # 3. Verify state and artifacts are rewinded + session = await runner.session_service.get_session( + app_name=runner.app_name, user_id=user_id, session_id=session_id + ) + # After rewind before invocation2, only event1 state delta should apply. + assert session.state["k1"] == "v1" + assert not session.state["k2"] + # f1 should be rewinded to v0 + assert await runner.artifact_service.load_artifact( + app_name=runner.app_name, + user_id=user_id, + session_id=session_id, + filename="f1", + ) == types.Part.from_text(text="f1v0") + # f2 should not exist + assert ( + await runner.artifact_service.load_artifact( + app_name=runner.app_name, + user_id=user_id, + session_id=session_id, + filename="f2", + ) + is None + ) + + @pytest.mark.asyncio + async def test_rewind_async_not_first_invocation(self): + """Tests rewind_async rewinds state and artifacts to invocation2.""" + runner = self.runner + user_id = "test_user" + session_id = "test_session" + + # 1. Setup session and initial artifacts + session = await runner.session_service.create_session( + app_name=runner.app_name, user_id=user_id, session_id=session_id + ) + # invocation1 + await runner.artifact_service.save_artifact( + app_name=runner.app_name, + user_id=user_id, + session_id=session_id, + filename="f1", + artifact=types.Part.from_text(text="f1v0"), + ) + event1 = Event( + invocation_id="invocation1", + author="agent", + content=types.Content(parts=[types.Part.from_text(text="event1")]), + actions=EventActions( + state_delta={"k1": "v1"}, artifact_delta={"f1": 0} + ), + ) + await runner.session_service.append_event(session=session, event=event1) + + # invocation2 + await runner.artifact_service.save_artifact( + app_name=runner.app_name, + user_id=user_id, + session_id=session_id, + filename="f1", + artifact=types.Part.from_text(text="f1v1"), + ) + await runner.artifact_service.save_artifact( + app_name=runner.app_name, + user_id=user_id, + session_id=session_id, + filename="f2", + artifact=types.Part.from_text(text="f2v0"), + ) + event2 = Event( + invocation_id="invocation2", + author="agent", + content=types.Content(parts=[types.Part.from_text(text="event2")]), + actions=EventActions( + state_delta={"k1": "v2", "k2": "v2"}, + artifact_delta={"f1": 1, "f2": 0}, + ), + ) + await runner.session_service.append_event(session=session, event=event2) + + # invocation3 + event3 = Event( + invocation_id="invocation3", + author="agent", + content=types.Content(parts=[types.Part.from_text(text="event3")]), + actions=EventActions(state_delta={"k2": "v3"}), + ) + await runner.session_service.append_event(session=session, event=event3) + + # 2. Rewind before invocation3 + await runner.rewind_async( + user_id=user_id, + session_id=session_id, + rewind_before_invocation_id="invocation3", + ) + + # 3. Verify state and artifacts are rewinded + session = await runner.session_service.get_session( + app_name=runner.app_name, user_id=user_id, session_id=session_id + ) + # After rewind before invocation3, event1 and event2 state deltas should apply. + assert session.state == {"k1": "v2", "k2": "v2"} + # f1 should be v1 + assert await runner.artifact_service.load_artifact( + app_name=runner.app_name, + user_id=user_id, + session_id=session_id, + filename="f1", + ) == types.Part.from_text(text="f1v1") + # f2 should be v0 + assert await runner.artifact_service.load_artifact( + app_name=runner.app_name, + user_id=user_id, + session_id=session_id, + filename="f2", + ) == types.Part.from_text(text="f2v0") diff --git a/tests/unittests/sessions/test_session_service.py b/tests/unittests/sessions/test_session_service.py index 0c005d68..0dd4162e 100644 --- a/tests/unittests/sessions/test_session_service.py +++ b/tests/unittests/sessions/test_session_service.py @@ -390,6 +390,11 @@ async def test_append_event_complete(service_type): error_code='error_code', error_message='error_message', interrupted=True, + usage_metadata=types.GenerateContentResponseUsageMetadata( + prompt_token_count=1, candidates_token_count=1, total_token_count=2 + ), + citation_metadata=types.CitationMetadata(), + custom_metadata={'custom_key': 'custom_value'}, ) await session_service.append_event(session=session, event=event) diff --git a/tests/unittests/telemetry/test_spans.py b/tests/unittests/telemetry/test_spans.py index cfd2b745..38a8358f 100644 --- a/tests/unittests/telemetry/test_spans.py +++ b/tests/unittests/telemetry/test_spans.py @@ -13,6 +13,7 @@ # limitations under the License. import json +import os from typing import Any from typing import Dict from typing import Optional @@ -23,6 +24,7 @@ from google.adk.agents.llm_agent import LlmAgent from google.adk.models.llm_request import LlmRequest from google.adk.models.llm_response import LlmResponse from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.telemetry.tracing import ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS from google.adk.telemetry.tracing import trace_agent_invocation from google.adk.telemetry.tracing import trace_call_llm from google.adk.telemetry.tracing import trace_merged_tool_calls @@ -371,3 +373,147 @@ def test_trace_merged_tool_calls_sets_correct_attributes( expected_calls, any_order=True ) mock_event_fixture.model_dumps_json.assert_called_once_with(exclude_none=True) + + +@pytest.mark.asyncio +async def test_call_llm_disabling_request_response_content( + monkeypatch, mock_span_fixture +): + """Test trace_call_llm doesn't set request and response attributes if env is set to false""" + # Arrange + monkeypatch.setenv(ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS, 'false') + monkeypatch.setattr( + 'opentelemetry.trace.get_current_span', lambda: mock_span_fixture + ) + + agent = LlmAgent(name='test_agent') + invocation_context = await _create_invocation_context(agent) + llm_request = LlmRequest( + model='gemini-pro', + contents=[ + types.Content( + role='user', + parts=[types.Part(text='Hello, how are you?')], + ), + ], + ) + llm_response = LlmResponse( + turn_complete=True, + finish_reason=types.FinishReason.STOP, + ) + + # Act + trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response) + + # Assert + assert not any( + call_obj.args[0] == 'gcp.vertex.agent.llm_request' + and call_obj.args[1] != {} + for call_obj in mock_span_fixture.set_attribute.call_args_list + ), "Attribute 'gcp.vertex.agent.llm_request' was incorrectly set on the span." + + assert not any( + call_obj.args[0] == 'gcp.vertex.agent.llm_response' + and call_obj.args[1] != {} + for call_obj in mock_span_fixture.set_attribute.call_args_list + ), ( + "Attribute 'gcp.vertex.agent.llm_response' was incorrectly set on the" + ' span.' + ) + + +def test_trace_tool_call_disabling_request_response_content( + monkeypatch, + mock_span_fixture, + mock_tool_fixture, + mock_event_fixture, +): + """Test trace_tool_call doesn't set request and response attributes if env is set to false""" + # Arrange + monkeypatch.setenv(ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS, 'false') + monkeypatch.setattr( + 'opentelemetry.trace.get_current_span', lambda: mock_span_fixture + ) + + test_args: Dict[str, Any] = {'query': 'details', 'id_list': [1, 2, 3]} + test_tool_call_id: str = 'tool_call_id_002' + test_event_id: str = 'event_id_dict_002' + dict_function_response: Dict[str, Any] = { + 'data': 'structured_data', + 'count': 5, + } + + mock_event_fixture.id = test_event_id + mock_event_fixture.content = types.Content( + role='user', + parts=[ + types.Part( + function_response=types.FunctionResponse( + id=test_tool_call_id, + name='test_function_1', + response=dict_function_response, + ) + ), + ], + ) + + # Act + trace_tool_call( + tool=mock_tool_fixture, + args=test_args, + function_response_event=mock_event_fixture, + ) + + # Assert + assert not any( + call_obj.args[0] == 'gcp.vertex.agent.tool_call_args' + and call_obj.args[1] != {} + for call_obj in mock_span_fixture.set_attribute.call_args_list + ), ( + "Attribute 'gcp.vertex.agent.tool_call_args' was incorrectly set on the" + ' span.' + ) + + assert not any( + call_obj.args[0] == 'gcp.vertex.agent.tool_response' + and call_obj.args[1] != {} + for call_obj in mock_span_fixture.set_attribute.call_args_list + ), ( + "Attribute 'gcp.vertex.agent.tool_response' was incorrectly set on the" + ' span.' + ) + + +def test_trace_merged_tool_disabling_request_response_content( + monkeypatch, + mock_span_fixture, + mock_event_fixture, +): + """Test trace_merged_tool doesn't set request and response attributes if env is set to false""" + # Arrange + monkeypatch.setenv(ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS, 'false') + monkeypatch.setattr( + 'opentelemetry.trace.get_current_span', lambda: mock_span_fixture + ) + + test_response_event_id = 'merged_evt_id_001' + custom_event_json_output = ( + '{"custom_event_payload": true, "details": "merged_details"}' + ) + mock_event_fixture.model_dumps_json.return_value = custom_event_json_output + + # Act + trace_merged_tool_calls( + response_event_id=test_response_event_id, + function_response_event=mock_event_fixture, + ) + + # Assert + assert not any( + call_obj.args[0] == 'gcp.vertex.agent.tool_response' + and call_obj.args[1] != {} + for call_obj in mock_span_fixture.set_attribute.call_args_list + ), ( + "Attribute 'gcp.vertex.agent.tool_response' was incorrectly set on the" + ' span.' + ) diff --git a/tests/unittests/tools/mcp_tool/test_mcp_tool.py b/tests/unittests/tools/mcp_tool/test_mcp_tool.py index 77854477..37c2df89 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_tool.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_tool.py @@ -640,3 +640,74 @@ class TestMCPTool: with pytest.raises(TypeError): MCPTool(mcp_tool=self.mock_mcp_tool) # Missing session manager + + @pytest.mark.asyncio + async def test_run_async_impl_with_header_provider_no_auth(self): + """Test running tool with header_provider but no auth.""" + expected_headers = {"X-Tenant-ID": "test-tenant"} + header_provider = Mock(return_value=expected_headers) + tool = MCPTool( + mcp_tool=self.mock_mcp_tool, + mcp_session_manager=self.mock_session_manager, + header_provider=header_provider, + ) + + expected_response = {"result": "success"} + self.mock_session.call_tool = AsyncMock(return_value=expected_response) + + tool_context = Mock(spec=ToolContext) + tool_context._invocation_context = Mock() + args = {"param1": "test_value"} + + result = await tool._run_async_impl( + args=args, tool_context=tool_context, credential=None + ) + + assert result == expected_response + header_provider.assert_called_once() + self.mock_session_manager.create_session.assert_called_once_with( + headers=expected_headers + ) + self.mock_session.call_tool.assert_called_once_with( + "test_tool", arguments=args + ) + + @pytest.mark.asyncio + async def test_run_async_impl_with_header_provider_and_oauth2(self): + """Test running tool with header_provider and OAuth2 auth.""" + dynamic_headers = {"X-Tenant-ID": "test-tenant"} + header_provider = Mock(return_value=dynamic_headers) + tool = MCPTool( + mcp_tool=self.mock_mcp_tool, + mcp_session_manager=self.mock_session_manager, + header_provider=header_provider, + ) + + oauth2_auth = OAuth2Auth(access_token="test_access_token") + credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, oauth2=oauth2_auth + ) + + expected_response = {"result": "success"} + self.mock_session.call_tool = AsyncMock(return_value=expected_response) + + tool_context = Mock(spec=ToolContext) + tool_context._invocation_context = Mock() + args = {"param1": "test_value"} + + result = await tool._run_async_impl( + args=args, tool_context=tool_context, credential=credential + ) + + assert result == expected_response + header_provider.assert_called_once() + self.mock_session_manager.create_session.assert_called_once() + call_args = self.mock_session_manager.create_session.call_args + headers = call_args[1]["headers"] + assert headers == { + "Authorization": "Bearer test_access_token", + "X-Tenant-ID": "test-tenant", + } + self.mock_session.call_tool.assert_called_once_with( + "test_tool", arguments=args + ) diff --git a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py index d5e6ae24..37b23d99 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py @@ -29,6 +29,7 @@ pytestmark = pytest.mark.skipif( # Import dependencies with version checking try: + from google.adk.agents.readonly_context import ReadonlyContext from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams @@ -55,6 +56,7 @@ except ImportError as e: StreamableHTTPConnectionParams = DummyClass MCPTool = DummyClass MCPToolset = DummyClass + ReadonlyContext = DummyClass else: raise e @@ -245,6 +247,31 @@ class TestMCPToolset: assert tools[0].name == "read_file" assert tools[1].name == "write_file" + @pytest.mark.asyncio + async def test_get_tools_with_header_provider(self): + """Test get_tools with a header_provider.""" + mock_tools = [MockMCPTool("tool1"), MockMCPTool("tool2")] + self.mock_session.list_tools = AsyncMock( + return_value=MockListToolsResult(mock_tools) + ) + mock_readonly_context = Mock(spec=ReadonlyContext) + expected_headers = {"X-Tenant-ID": "test-tenant"} + header_provider = Mock(return_value=expected_headers) + + toolset = MCPToolset( + connection_params=self.mock_stdio_params, + header_provider=header_provider, + ) + toolset._mcp_session_manager = self.mock_session_manager + + tools = await toolset.get_tools(readonly_context=mock_readonly_context) + + assert len(tools) == 2 + header_provider.assert_called_once_with(mock_readonly_context) + self.mock_session_manager.create_session.assert_called_once_with( + headers=expected_headers + ) + @pytest.mark.asyncio async def test_close_success(self): """Test successful cleanup.""" diff --git a/tests/unittests/tools/test_agent_tool.py b/tests/unittests/tools/test_agent_tool.py index 0b476966..1743a30f 100644 --- a/tests/unittests/tools/test_agent_tool.py +++ b/tests/unittests/tools/test_agent_tool.py @@ -178,7 +178,8 @@ def test_update_state(): assert runner.session.state['state_1'] == 'changed_value' -def test_update_artifacts(): +@mark.asyncio +async def test_update_artifacts(): """The agent tool can read and write artifacts.""" async def before_tool_agent(callback_context: CallbackContext): @@ -219,12 +220,21 @@ def test_update_artifacts(): runner = testing_utils.InMemoryRunner(root_agent) runner.run('test1') - artifacts_path = f'test_app/test_user/{runner.session_id}' - assert runner.runner.artifact_service.artifacts == { - f'{artifacts_path}/artifact_1': [Part.from_text(text='test')], - f'{artifacts_path}/artifact_2': [Part.from_text(text='test 2')], - f'{artifacts_path}/artifact_3': [Part.from_text(text='test 2 3')], - } + async def load_artifact(filename: str): + return await runner.runner.artifact_service.load_artifact( + app_name='test_app', + user_id='test_user', + session_id=runner.session_id, + filename=filename, + ) + + assert await runner.runner.artifact_service.list_artifact_keys( + app_name='test_app', user_id='test_user', session_id=runner.session_id + ) == ['artifact_1', 'artifact_2', 'artifact_3'] + + assert await load_artifact('artifact_1') == Part.from_text(text='test') + assert await load_artifact('artifact_2') == Part.from_text(text='test 2') + assert await load_artifact('artifact_3') == Part.from_text(text='test 2 3') @mark.parametrize(