encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-03-30 10:57:20 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat: introduces a new AgentEngineSandboxCodeExecutor class that supports executes agent generated code

Browse Source 
The AgentEngineSandboxCodeExecutor uses the Vertex AI Code Execution Sandbox API to execute code

PiperOrigin-RevId: 821197794
This commit is contained in:
Google Team Member
2025-10-18 20:23:27 -07:00
committed by Copybara-Service
parent af74eba695
commit a5b742b360
8 changed files with 3 additions and 455 deletions
Download Patch File Download Diff File Expand all files Collapse all files
contributing/samples/agent_engine_code_execution/README
-18
View File
@@ -1,18 +0,0 @@
# OAuth Sample
## Introduction
This sample data science agent uses Agent Engine Code Execution Sandbox to execute LLM generated code.
## How to use
* 1. Follow https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/code-execution/overview to create a code execution sandbox environment.
* 2. Replace the SANDBOX_RESOURCE_NAME with the one you just created. If you dont want to create a new sandbox environment directly, the Agent Engine Code Execution Sandbox will create one for you by default using the AGENT_ENGINE_RESOURCE_NAME you specified, however, please ensure to clean up sandboxes after use, otherwise, it will consume quotas.
## Sample prompt
* Can you write a function that calculates the sum from 1 to 100.
* The dataset is given as below. Store,Date,Weekly_Sales,Holiday_Flag,Temperature,Fuel_Price,CPI,Unemployment Store 1,2023-06-01,1000,0,70,3.0,200,5 Store 2,2023-06-02,1200,1,80,3.5,210,6 Store 3,2023-06-03,1400,0,90,4.0,220,7 Store 4,2023-06-04,1600,1,70,4.5,230,8 Store 5,2023-06-05,1800,0,80,5.0,240,9 Store 6,2023-06-06,2000,1,90,5.5,250,10 Store 7,2023-06-07,2200,0,90,6.0,260,11 Plot a scatter plot showcasing the relationship between Weekly Sales and Temperature for each store, distinguishing stores with a Holiday Flag.
contributing/samples/agent_engine_code_execution/__init__.py
-15
View File
@@ -1,15 +0,0 @@
# 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
contributing/samples/agent_engine_code_execution/agent.py
-95
View File
@@ -1,95 +0,0 @@
# 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.
"""Data science agent."""
from google.adk.agents.llm_agent import Agent
from google.adk.code_executors.agent_engine_sandbox_code_executor import AgentEngineSandboxCodeExecutor
def base_system_instruction():
"""Returns: data science agent system instruction."""
return """
# Guidelines
**Objective:** Assist the user in achieving their data analysis goals within the context of a Python Colab notebook, **with emphasis on avoiding assumptions and ensuring accuracy.** Reaching that goal can involve multiple steps. When you need to generate code, you **don't** need to solve the goal in one go. Only generate the next step at a time.
**Code Execution:** All code snippets provided will be executed within the Colab environment.
**Statefulness:** All code snippets are executed and the variables stays in the environment. You NEVER need to re-initialize variables. You NEVER need to reload files. You NEVER need to re-import libraries.
**Output Visibility:** Always print the output of code execution to visualize results, especially for data exploration and analysis. For example:
- To look a the shape of a pandas.DataFrame do:
```tool_code
print(df.shape)
```
The output will be presented to you as:
```tool_outputs
(49, 7)
```
- To display the result of a numerical computation:
```tool_code
x = 10 ** 9 - 12 ** 5
print(f'{{x=}}')
```
The output will be presented to you as:
```tool_outputs
x=999751168
```
- You **never** generate ```tool_outputs yourself.
- You can then use this output to decide on next steps.
- Print just variables (e.g., `print(f'{{variable=}}')`.
**No Assumptions:** **Crucially, avoid making assumptions about the nature of the data or column names.** Base findings solely on the data itself. Always use the information obtained from `explore_df` to guide your analysis.
**Available files:** Only use the files that are available as specified in the list of available files.
**Data in prompt:** Some queries contain the input data directly in the prompt. You have to parse that data into a pandas DataFrame. ALWAYS parse all the data. NEVER edit the data that are given to you.
**Answerability:** Some queries may not be answerable with the available data. In those cases, inform the user why you cannot process their query and suggest what type of data would be needed to fulfill their request.
"""
root_agent = Agent(
model="gemini-2.0-flash-001",
name="agent_engine_code_execution_agent",
instruction=base_system_instruction() + """
You need to assist the user with their queries by looking at the data and the context in the conversation.
You final answer should summarize the code and code execution relevant to the user query.
You should include all pieces of data to answer the user query, such as the table from code execution results.
If you cannot answer the question directly, you should follow the guidelines above to generate the next step.
If the question can be answered directly with writing any code, you should do that.
If you doesn't have enough data to answer the question, you should ask for clarification from the user.
You should NEVER install any package on your own like `pip install ...`.
When plotting trends, you should make sure to sort and order the data by the x-axis.
""",
code_executor=AgentEngineSandboxCodeExecutor(
# Replace with your sandbox resource name if you already have one.
sandbox_resource_name="SANDBOX_RESOURCE_NAME",
# "projects/vertex-agent-loadtest/locations/us-central1/reasoningEngines/6842889780301135872/sandboxEnvironments/6545148628569161728",
# Replace with agent engine resource name used for creating sandbox if
# sandbox_resource_name is not set.
agent_engine_resource_name="AGENT_ENGINE_RESOURCE_NAME",
),
)
pyproject.toml
+1 -1
View File
@@ -32,7 +32,7 @@ dependencies = [
"click>=8.1.8, <9.0.0", # For CLI tools
"fastapi>=0.115.0, <1.0.0", # FastAPI framework
"google-api-python-client>=2.157.0, <3.0.0", # Google API client discovery
"google-cloud-aiplatform[agent_engines]>=1.121.0, <2.0.0",# For VertexAI integrations, e.g. example store.
"google-cloud-aiplatform[agent_engines]>=1.112.0, <2.0.0",# For VertexAI integrations, e.g. example store.
"google-cloud-bigtable>=2.32.0", # For Bigtable database
"google-cloud-discoveryengine>=0.13.12, <0.14.0", # For Discovery Engine Search Tool
"google-cloud-secret-manager>=2.22.0, <3.0.0", # Fetching secrets in RestAPI Tool
src/google/adk/code_executors/__init__.py
-13
View File
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import logging
from .base_code_executor import BaseCodeExecutor
@@ -31,7 +29,6 @@ __all__ = [
'VertexAiCodeExecutor',
'ContainerCodeExecutor',
'GkeCodeExecutor',
'AgentEngineSandboxCodeExecutor',
]
@@ -66,14 +63,4 @@ def __getattr__(name: str):
'GkeCodeExecutor requires additional dependencies. '
'Please install with: pip install "google-adk[extensions]"'
) from e
elif name == 'AgentEngineSandboxCodeExecutor':
try:
from .agent_engine_sandbox_code_executor import AgentEngineSandboxCodeExecutor
return AgentEngineSandboxCodeExecutor
except ImportError as e:
raise ImportError(
'AgentEngineSandboxCodeExecutor requires additional dependencies. '
'Please install with: pip install "google-adk[extensions]"'
) from e
raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
src/google/adk/code_executors/agent_engine_sandbox_code_executor.py
-187
View File
@@ -1,187 +0,0 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import json
import logging
import mimetypes
import re
from typing import Optional
from typing_extensions import override
import vertexai
from vertexai import types
from ..agents.invocation_context import InvocationContext
from ..utils.feature_decorator import experimental
from .base_code_executor import BaseCodeExecutor
from .code_execution_utils import CodeExecutionInput
from .code_execution_utils import CodeExecutionResult
from .code_execution_utils import File
logger = logging.getLogger('google_adk.' + __name__)
@experimental
class AgentEngineSandboxCodeExecutor(BaseCodeExecutor):
"""A code executor that uses Agent Engine Code Execution Sandbox to execute code.
Attributes:
sandbox_resource_name: If set, load the existing resource name of the code
interpreter extension instead of creating a new one. Format:
projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789
"""
sandbox_resource_name: str = None
def __init__(
self,
sandbox_resource_name: Optional[str] = None,
agent_engine_resource_name: Optional[str] = None,
**data,
):
"""Initializes the AgentEngineSandboxCodeExecutor.
Args:
sandbox_resource_name: If set, load the existing resource name of code
execution sandbox, if not set, create a new one. Format:
projects/123/locations/us-central1/reasoningEngines/456/
sandboxEnvironments/789
agent_engine_resource_name: The resource name of the agent engine to use
to create the code execution sandbox. Format:
projects/123/locations/us-central1/reasoningEngines/456, when both
sandbox_resource_name and agent_engine_resource_name are set,
agent_engine_resource_name will be ignored.
**data: Additional keyword arguments to be passed to the base class.
"""
super().__init__(**data)
sandbox_resource_name_pattern = r'^projects/([a-zA-Z0-9-_]+)/locations/([a-zA-Z0-9-_]+)/reasoningEngines/(\d+)/sandboxEnvironments/(\d+)$'
agent_engine_resource_name_pattern = r'^projects/([a-zA-Z0-9-_]+)/locations/([a-zA-Z0-9-_]+)/reasoningEngines/(\d+)$'
if sandbox_resource_name is not None:
self.sandbox_resource_name = sandbox_resource_name
self._project_id, self._location = (
self._get_project_id_and_location_from_resource_name(
sandbox_resource_name, sandbox_resource_name_pattern
)
)
elif agent_engine_resource_name is not None:
self._project_id, self._location = (
self._get_project_id_and_location_from_resource_name(
agent_engine_resource_name, agent_engine_resource_name_pattern
)
)
# @TODO - Add TTL for sandbox creation after it is available
# in SDK.
operation = self._get_api_client().agent_engines.sandboxes.create(
spec={'code_execution_environment': {}},
name=agent_engine_resource_name,
config=types.CreateAgentEngineSandboxConfig(
display_name='default_sandbox'
),
)
self.sandbox_resource_name = operation.response.name
else:
raise ValueError(
'Either sandbox_resource_name or agent_engine_resource_name must be'
' set.'
)
@override
def execute_code(
self,
invocation_context: InvocationContext,
code_execution_input: CodeExecutionInput,
) -> CodeExecutionResult:
# Execute the code.
input_data = {
'code': code_execution_input.code,
}
if code_execution_input.input_files:
input_data['files'] = [
{
'name': f.name,
'contents': f.content,
'mimeType': f.mime_type,
}
for f in code_execution_input.input_files
]
code_execution_response = (
self._get_api_client().agent_engines.sandboxes.execute_code(
name=self.sandbox_resource_name,
input_data=input_data,
)
)
saved_files = []
stdout = ''
stderr = ''
for output in code_execution_response.outputs:
if output.mime_type == 'application/json' and (
output.metadata is None
or output.metadata.attributes is None
or 'file_name' not in output.metadata.attributes
):
json_output_data = json.loads(output.data.decode('utf-8'))
stdout = json_output_data.get('stdout', '')
stderr = json_output_data.get('stderr', '')
else:
file_name = ''
if (
output.metadata is not None
and output.metadata.attributes is not None
):
file_name = output.metadata.attributes.get('file_name', b'').decode(
'utf-8'
)
mime_type = output.mime_type
if not mime_type:
mime_type, _ = mimetypes.guess_type(file_name)
saved_files.append(
File(
name=file_name,
content=output.data,
mime_type=mime_type,
)
)
# Collect the final result.
return CodeExecutionResult(
stdout=stdout,
stderr=stderr,
output_files=saved_files,
)
def _get_api_client(self):
"""Instantiates an API client for the given project and location.
It needs to be instantiated inside each request so that the event loop
management can be properly propagated.
Returns:
An API client for the given project and location.
"""
return vertexai.Client(project=self._project_id, location=self._location)
def _get_project_id_and_location_from_resource_name(
self, resource_name: str, pattern: str
) -> tuple[str, str]:
"""Extracts the project ID and location from the resource name."""
match = re.fullmatch(pattern, resource_name)
if not match:
raise ValueError(f'resource name {resource_name} is not valid.')
return match.groups()[0], match.groups()[1]
src/google/adk/code_executors/code_execution_utils.py
+2 -4
View File
@@ -14,8 +14,6 @@
"""Utility functions for code execution."""
from __future__ import annotations
import base64
import binascii
import copy
@@ -36,9 +34,9 @@ class File:
The name of the file with file extension (e.g., "file.csv").
"""
content: str | bytes
content: str
"""
The base64-encoded bytes of the file content or the original bytes of the file content.
The base64-encoded bytes of the file content.
"""
mime_type: str = 'text/plain'
tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py
-122
View File
@@ -1,122 +0,0 @@
# 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 json
from unittest.mock import MagicMock
from unittest.mock import patch
from google.adk.agents.invocation_context import InvocationContext
from google.adk.code_executors.agent_engine_sandbox_code_executor import AgentEngineSandboxCodeExecutor
from google.adk.code_executors.code_execution_utils import CodeExecutionInput
import pytest
@pytest.fixture
def mock_invocation_context() -> InvocationContext:
"""Fixture for a mock InvocationContext."""
mock = MagicMock(spec=InvocationContext)
mock.invocation_id = "test-invocation-123"
return mock
class TestAgentEngineSandboxCodeExecutor:
"""Unit tests for the AgentEngineSandboxCodeExecutor."""
def test_init_with_sandbox_overrides(self):
"""Tests that class attributes can be overridden at instantiation."""
executor = AgentEngineSandboxCodeExecutor(
sandbox_resource_name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789",
)
assert executor.sandbox_resource_name == (
"projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789"
)
def test_init_with_sandbox_overrides_throws_error(self):
"""Tests that class attributes can be overridden at instantiation."""
with pytest.raises(ValueError):
AgentEngineSandboxCodeExecutor(
sandbox_resource_name="projects/123/locations/us-central1/reasoningEgines/456/sandboxes/789",
)
def test_init_with_agent_engine_overrides_throws_error(self):
"""Tests that class attributes can be overridden at instantiation."""
with pytest.raises(ValueError):
AgentEngineSandboxCodeExecutor(
agent_engine_resource_name=(
"projects/123/locations/us-central1/reason/456"
),
)
@patch(
"google.adk.code_executors.agent_engine_sandbox_code_executor.vertexai"
)
def test_execute_code_success(
self,
mock_vertexai,
mock_invocation_context,
):
# Setup Mocks
mock_api_client = MagicMock()
mock_vertexai.Client.return_value = mock_api_client
mock_response = MagicMock()
mock_json_output = MagicMock()
mock_json_output.mime_type = "application/json"
mock_json_output.data = json.dumps(
{"stdout": "hello world", "stderr": ""}
).encode("utf-8")
mock_json_output.metadata = None
mock_file_output = MagicMock()
mock_file_output.mime_type = "text/plain"
mock_file_output.data = b"file content"
mock_file_output.metadata = MagicMock()
mock_file_output.metadata.attributes = {"file_name": b"file.txt"}
mock_png_file_output = MagicMock()
mock_png_file_output.mime_type = "image/png"
sample_png_bytes = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00\x1f\x15\xc4\x89"
mock_png_file_output.data = sample_png_bytes
mock_png_file_output.metadata = MagicMock()
mock_png_file_output.metadata.attributes = {"file_name": b"file.png"}
mock_response.outputs = [
mock_json_output,
mock_file_output,
mock_png_file_output,
]
mock_api_client.agent_engines.sandboxes.execute_code.return_value = (
mock_response
)
# Execute
executor = AgentEngineSandboxCodeExecutor(
sandbox_resource_name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789"
)
code_input = CodeExecutionInput(code='print("hello world")')
result = executor.execute_code(mock_invocation_context, code_input)
# Assert
assert result.stdout == "hello world"
assert not result.stderr
assert result.output_files[0].mime_type == "text/plain"
assert result.output_files[0].content == b"file content"
assert result.output_files[0].name == "file.txt"
assert result.output_files[1].mime_type == "image/png"
assert result.output_files[1].name == "file.png"
assert result.output_files[1].content == sample_png_bytes
mock_api_client.agent_engines.sandboxes.execute_code.assert_called_once_with(
name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789",
input_data={"code": 'print("hello world")'},
)