diff --git a/contributing/samples/agent_engine_code_execution/README b/contributing/samples/agent_engine_code_execution/README index 8d5a4442..b0443ae2 100644 --- a/contributing/samples/agent_engine_code_execution/README +++ b/contributing/samples/agent_engine_code_execution/README @@ -7,9 +7,9 @@ This sample data science agent uses Agent Engine Code Execution Sandbox to execu ## 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. +* 1. Follow https://docs.cloud.google.com/agent-builder/agent-engine/code-execution/quickstart#create-an-agent-engine-instance to create an agent engine instance. Replace the AGENT_ENGINE_RESOURCE_NAME with the one you just created. A new sandbox environment under this agent engine instance will be created for each session with TTL of 1 year. But sandbox can only main its state for up to 14 days. This is the recommended usage for production environments. -* 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. +* 2. For testing or protyping purposes, create a sandbox environment by following this guide: https://docs.cloud.google.com/agent-builder/agent-engine/code-execution/quickstart#create_a_sandbox. Replace the SANDBOX_RESOURCE_NAME with the one you just created. This will be used as the default sandbox environment for all the code executions throughout the lifetime of the agent. As the sandbox is re-used across sessions, all sessions will share the same Python environment and variable values." ## Sample prompt diff --git a/contributing/samples/agent_engine_code_execution/agent.py b/contributing/samples/agent_engine_code_execution/agent.py index d85989eb..a32e4ca4 100644 --- a/contributing/samples/agent_engine_code_execution/agent.py +++ b/contributing/samples/agent_engine_code_execution/agent.py @@ -85,11 +85,10 @@ When plotting trends, you should make sure to sort and order the data by the x-a """, code_executor=AgentEngineSandboxCodeExecutor( - # Replace with your sandbox resource name if you already have one. - sandbox_resource_name="SANDBOX_RESOURCE_NAME", + # Replace with your sandbox resource name if you already have one. Only use it for testing or prototyping purposes, because this will use the same sandbox for all requests. # "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. + sandbox_resource_name=None, + # Replace with agent engine resource name used for creating sandbox environment. agent_engine_resource_name="AGENT_ENGINE_RESOURCE_NAME", ), ) diff --git a/src/google/adk/code_executors/agent_engine_sandbox_code_executor.py b/src/google/adk/code_executors/agent_engine_sandbox_code_executor.py index 69d1778a..071d59dc 100644 --- a/src/google/adk/code_executors/agent_engine_sandbox_code_executor.py +++ b/src/google/adk/code_executors/agent_engine_sandbox_code_executor.py @@ -38,10 +38,15 @@ class AgentEngineSandboxCodeExecutor(BaseCodeExecutor): 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 + 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 """ sandbox_resource_name: str = None + agent_engine_resource_name: str = None + def __init__( self, sandbox_resource_name: Optional[str] = None, @@ -67,30 +72,19 @@ class AgentEngineSandboxCodeExecutor(BaseCodeExecutor): 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 ) ) + self.sandbox_resource_name = sandbox_resource_name elif agent_engine_resource_name is not None: - from vertexai import types - 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 + self.agent_engine_resource_name = agent_engine_resource_name else: raise ValueError( 'Either sandbox_resource_name or agent_engine_resource_name must be' @@ -103,6 +97,45 @@ class AgentEngineSandboxCodeExecutor(BaseCodeExecutor): invocation_context: InvocationContext, code_execution_input: CodeExecutionInput, ) -> CodeExecutionResult: + # default to the sandbox resource name if set. + sandbox_name = self.sandbox_resource_name + if self.sandbox_resource_name is None: + from google.api_core import exceptions + from vertexai import types + + # use sandbox name stored in session if available. + sandbox_name = invocation_context.session.state.get('sandbox_name', None) + create_new_sandbox = False + if sandbox_name is None: + create_new_sandbox = True + else: + # Check if the sandbox is still running OR already expired due to ttl. + try: + sandbox = self._get_api_client().agent_engines.sandboxes.get( + name=sandbox_name + ) + if sandbox is None or sandbox.state != 'STATE_RUNNING': + create_new_sandbox = True + except exceptions.NotFound: + create_new_sandbox = True + + if create_new_sandbox: + # Create a new sandbox and assign it to sandbox_name. + operation = self._get_api_client().agent_engines.sandboxes.create( + spec={'code_execution_environment': {}}, + name=self.agent_engine_resource_name, + config=types.CreateAgentEngineSandboxConfig( + # VertexAiSessionService has a default TTL of 1 year, so we set + # the sandbox TTL to 1 year as well. For the current code + # execution sandbox, if it hasn't been used for 14 days, the + # state will be lost. + display_name='default_sandbox', + ttl='31536000s', + ), + ) + sandbox_name = operation.response.name + invocation_context.session.state['sandbox_name'] = sandbox_name + # Execute the code. input_data = { 'code': code_execution_input.code, @@ -119,7 +152,7 @@ class AgentEngineSandboxCodeExecutor(BaseCodeExecutor): code_execution_response = ( self._get_api_client().agent_engines.sandboxes.execute_code( - name=self.sandbox_resource_name, + name=sandbox_name, input_data=input_data, ) ) diff --git a/tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py b/tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py index 6022527f..9b27b82c 100644 --- a/tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py +++ b/tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py @@ -19,6 +19,7 @@ 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 +from google.adk.sessions.session import Session import pytest @@ -27,6 +28,10 @@ def mock_invocation_context() -> InvocationContext: """Fixture for a mock InvocationContext.""" mock = MagicMock(spec=InvocationContext) mock.invocation_id = "test-invocation-123" + session = MagicMock(spec=Session) + mock.session = session + session.state = {} + return mock @@ -118,3 +123,129 @@ class TestAgentEngineSandboxCodeExecutor: name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789", input_data={"code": 'print("hello world")'}, ) + + @patch("vertexai.Client") + def test_execute_code_recreates_sandbox_when_get_returns_none( + self, + mock_vertexai_client, + mock_invocation_context, + ): + # Setup Mocks + mock_api_client = MagicMock() + mock_vertexai_client.return_value = mock_api_client + + # Existing sandbox name stored in session, but get() will return None + existing_sandbox_name = "projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/old" + mock_invocation_context.session.state = { + "sandbox_name": existing_sandbox_name + } + + # Mock get to return None (simulating missing/expired sandbox) + mock_api_client.agent_engines.sandboxes.get.return_value = None + + # Mock create operation to return a new sandbox resource name + operation_mock = MagicMock() + created_sandbox_name = "projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + operation_mock.response.name = created_sandbox_name + mock_api_client.agent_engines.sandboxes.create.return_value = operation_mock + + # Mock execute_code response + mock_response = MagicMock() + mock_json_output = MagicMock() + mock_json_output.mime_type = "application/json" + mock_json_output.data = json.dumps( + {"stdout": "recreated sandbox run", "stderr": ""} + ).encode("utf-8") + mock_json_output.metadata = None + mock_response.outputs = [mock_json_output] + mock_api_client.agent_engines.sandboxes.execute_code.return_value = ( + mock_response + ) + + # Execute using agent_engine_resource_name so a sandbox can be created + executor = AgentEngineSandboxCodeExecutor( + agent_engine_resource_name=( + "projects/123/locations/us-central1/reasoningEngines/456" + ) + ) + code_input = CodeExecutionInput(code='print("hello world")') + result = executor.execute_code(mock_invocation_context, code_input) + + # Assert get was called for the existing sandbox + mock_api_client.agent_engines.sandboxes.get.assert_called_once_with( + name=existing_sandbox_name + ) + + # Assert create was called and session updated with new sandbox + mock_api_client.agent_engines.sandboxes.create.assert_called_once() + assert ( + mock_invocation_context.session.state["sandbox_name"] + == created_sandbox_name + ) + + # Assert execute_code used the created sandbox name + mock_api_client.agent_engines.sandboxes.execute_code.assert_called_once_with( + name=created_sandbox_name, + input_data={"code": 'print("hello world")'}, + ) + + @patch("vertexai.Client") + def test_execute_code_creates_sandbox_if_missing( + self, + mock_vertexai_client, + mock_invocation_context, + ): + # Setup Mocks + mock_api_client = MagicMock() + mock_vertexai_client.return_value = mock_api_client + + # Mock create operation to return a sandbox resource name + operation_mock = MagicMock() + created_sandbox_name = "projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + operation_mock.response.name = created_sandbox_name + mock_api_client.agent_engines.sandboxes.create.return_value = operation_mock + + # Mock execute_code response + mock_response = MagicMock() + mock_json_output = MagicMock() + mock_json_output.mime_type = "application/json" + mock_json_output.data = json.dumps( + {"stdout": "created sandbox run", "stderr": ""} + ).encode("utf-8") + mock_json_output.metadata = None + mock_response.outputs = [mock_json_output] + mock_api_client.agent_engines.sandboxes.execute_code.return_value = ( + mock_response + ) + + # Ensure session.state behaves like a dict for storing sandbox_name + mock_invocation_context.session.state = {} + + # Execute using agent_engine_resource_name so a sandbox will be created + executor = AgentEngineSandboxCodeExecutor( + agent_engine_resource_name=( + "projects/123/locations/us-central1/reasoningEngines/456" + ), + sandbox_resource_name=None, + ) + code_input = CodeExecutionInput(code='print("hello world")') + result = executor.execute_code(mock_invocation_context, code_input) + + # Assert sandbox creation was called and session state updated + mock_api_client.agent_engines.sandboxes.create.assert_called_once() + create_call_kwargs = ( + mock_api_client.agent_engines.sandboxes.create.call_args.kwargs + ) + assert create_call_kwargs["name"] == ( + "projects/123/locations/us-central1/reasoningEngines/456" + ) + assert ( + mock_invocation_context.session.state["sandbox_name"] + == created_sandbox_name + ) + + # Assert execute_code used the created sandbox name + mock_api_client.agent_engines.sandboxes.execute_code.assert_called_once_with( + name=created_sandbox_name, + input_data={"code": 'print("hello world")'}, + )