google-adk 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

google/adk/code_executors/code_executor_context.py

@@ -0,0 +1,202 @@
+# 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.
+
+"""The persistent context used to configure the code executor."""
+
+import copy
+import dataclasses
+import datetime
+from typing import Any
+from typing import Optional
+
+from ..sessions.state import State
+from .code_execution_utils import File
+
+_CONTEXT_KEY = '_code_execution_context'
+_SESSION_ID_KEY = 'execution_session_id'
+_PROCESSED_FILE_NAMES_KEY = 'processed_input_files'
+_INPUT_FILE_KEY = '_code_executor_input_files'
+_ERROR_COUNT_KEY = '_code_executor_error_counts'
+
+_CODE_EXECUTION_RESULTS_KEY = '_code_execution_results'
+
+
+class CodeExecutorContext:
+  """The persistent context used to configure the code executor."""
+
+  _context: dict[str, Any]
+
+  def __init__(self, session_state: State):
+    """Initializes the code executor context.
+
+    Args:
+      session_state: The session state to get the code executor context from.
+    """
+    self._context = self._get_code_executor_context(session_state)
+    self._session_state = session_state
+
+  def get_state_delta(self) -> dict[str, Any]:
+    """Gets the state delta to update in the persistent session state.
+
+    Returns:
+      The state delta to update in the persistent session state.
+    """
+    context_to_update = copy.deepcopy(self._context)
+    return {_CONTEXT_KEY: context_to_update}
+
+  def get_execution_id(self) -> Optional[str]:
+    """Gets the session ID for the code executor.
+
+    Returns:
+      The session ID for the code executor context.
+    """
+    if _SESSION_ID_KEY not in self._context:
+      return None
+    return self._context[_SESSION_ID_KEY]
+
+  def set_execution_id(self, session_id: str):
+    """Sets the session ID for the code executor.
+
+    Args:
+      session_id: The session ID for the code executor.
+    """
+    self._context[_SESSION_ID_KEY] = session_id
+
+  def get_processed_file_names(self) -> list[str]:
+    """Gets the processed file names from the session state.
+
+    Returns:
+      A list of processed file names in the code executor context.
+    """
+    if _PROCESSED_FILE_NAMES_KEY not in self._context:
+      return []
+    return self._context[_PROCESSED_FILE_NAMES_KEY]
+
+  def add_processed_file_names(self, file_names: [str]):
+    """Adds the processed file name to the session state.
+
+    Args:
+      file_names: The processed file names to add to the session state.
+    """
+    if _PROCESSED_FILE_NAMES_KEY not in self._context:
+      self._context[_PROCESSED_FILE_NAMES_KEY] = []
+    self._context[_PROCESSED_FILE_NAMES_KEY].extend(file_names)
+
+  def get_input_files(self) -> list[File]:
+    """Gets the code executor input file names from the session state.
+
+    Returns:
+      A list of input files in the code executor context.
+    """
+    if _INPUT_FILE_KEY not in self._session_state:
+      return []
+    return [File(**file) for file in self._session_state[_INPUT_FILE_KEY]]
+
+  def add_input_files(
+      self,
+      input_files: list[File],
+  ):
+    """Adds the input files to the code executor context.
+
+    Args:
+      input_files: The input files to add to the code executor context.
+    """
+    if _INPUT_FILE_KEY not in self._session_state:
+      self._session_state[_INPUT_FILE_KEY] = []
+    for input_file in input_files:
+      self._session_state[_INPUT_FILE_KEY].append(
+          dataclasses.asdict(input_file)
+      )
+
+  def clear_input_files(self):
+    """Removes the input files and processed file names to the code executor context."""
+    if _INPUT_FILE_KEY in self._session_state:
+      self._session_state[_INPUT_FILE_KEY] = []
+    if _PROCESSED_FILE_NAMES_KEY in self._context:
+      self._context[_PROCESSED_FILE_NAMES_KEY] = []
+
+  def get_error_count(self, invocation_id: str) -> int:
+    """Gets the error count from the session state.
+
+    Args:
+      invocation_id: The invocation ID to get the error count for.
+
+    Returns:
+      The error count for the given invocation ID.
+    """
+    if _ERROR_COUNT_KEY not in self._session_state:
+      return 0
+    return self._session_state[_ERROR_COUNT_KEY].get(invocation_id, 0)
+
+  def increment_error_count(self, invocation_id: str):
+    """Increments the error count from the session state.
+
+    Args:
+      invocation_id: The invocation ID to increment the error count for.
+    """
+    if _ERROR_COUNT_KEY not in self._session_state:
+      self._session_state[_ERROR_COUNT_KEY] = {}
+    self._session_state[_ERROR_COUNT_KEY][invocation_id] = (
+        self.get_error_count(invocation_id) + 1
+    )
+
+  def reset_error_count(self, invocation_id: str):
+    """Resets the error count from the session state.
+
+    Args:
+      invocation_id: The invocation ID to reset the error count for.
+    """
+    if _ERROR_COUNT_KEY not in self._session_state:
+      return
+    if invocation_id in self._session_state[_ERROR_COUNT_KEY]:
+      del self._session_state[_ERROR_COUNT_KEY][invocation_id]
+
+  def update_code_execution_result(
+      self,
+      invocation_id: str,
+      code: str,
+      result_stdout: str,
+      result_stderr: str,
+  ):
+    """Updates the code execution result.
+
+    Args:
+      invocation_id: The invocation ID to update the code execution result for.
+      code: The code to execute.
+      result_stdout: The standard output of the code execution.
+      result_stderr: The standard error of the code execution.
+    """
+    if _CODE_EXECUTION_RESULTS_KEY not in self._session_state:
+      self._session_state[_CODE_EXECUTION_RESULTS_KEY] = {}
+    if invocation_id not in self._session_state[_CODE_EXECUTION_RESULTS_KEY]:
+      self._session_state[_CODE_EXECUTION_RESULTS_KEY][invocation_id] = []
+    self._session_state[_CODE_EXECUTION_RESULTS_KEY][invocation_id].append({
+        'code': code,
+        'result_stdout': result_stdout,
+        'result_stderr': result_stderr,
+        'timestamp': int(datetime.datetime.now().timestamp()),
+    })
+
+  def _get_code_executor_context(self, session_state: State) -> dict[str, Any]:
+    """Gets the code executor context from the session state.
+
+    Args:
+      session_state: The session state to get the code executor context from.
+
+    Returns:
+      A dict of code executor context.
+    """
+    if _CONTEXT_KEY not in session_state:
+      session_state[_CONTEXT_KEY] = {}
+    return session_state[_CONTEXT_KEY]

google/adk/code_executors/container_code_executor.py

@@ -0,0 +1,196 @@
+# 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 atexit
+import os
+from typing import Optional
+
+import docker
+from docker.client import DockerClient
+from docker.models.containers import Container
+from pydantic import Field
+from typing_extensions import override
+
+from ..agents.invocation_context import InvocationContext
+from .base_code_executor import BaseCodeExecutor
+from .code_execution_utils import CodeExecutionInput
+from .code_execution_utils import CodeExecutionResult
+
+
+DEFAULT_IMAGE_TAG = 'adk-code-executor:latest'
+
+
+class ContainerCodeExecutor(BaseCodeExecutor):
+  """A code executor that uses a custom container to execute code.
+
+  Attributes:
+    base_url: Optional. The base url of the user hosted Docker client.
+    image: The tag of the predefined image or custom image to run on the
+      container. Either docker_path or image must be set.
+    docker_path: The path to the directory containing the Dockerfile. If set,
+      build the image from the dockerfile path instead of using the predefined
+      image. Either docker_path or image must be set.
+  """
+
+  base_url: Optional[str] = None
+  """
+  Optional. The base url of the user hosted Docker client.
+  """
+
+  image: str = None
+  """
+  The tag of the predefined image or custom image to run on the container.
+  Either docker_path or image must be set.
+  """
+
+  docker_path: str = None
+  """
+  The path to the directory containing the Dockerfile.
+  If set, build the image from the dockerfile path instead of using the
+  predefined image. Either docker_path or image must be set.
+  """
+
+  # Overrides the BaseCodeExecutor attribute: this executor cannot be stateful.
+  stateful: bool = Field(default=False, frozen=True, exclude=True)
+
+  # Overrides the BaseCodeExecutor attribute: this executor cannot
+  # optimize_data_file.
+  optimize_data_file: bool = Field(default=False, frozen=True, exclude=True)
+
+  _client: DockerClient = None
+  _container: Container = None
+
+  def __init__(
+      self,
+      base_url: Optional[str] = None,
+      image: Optional[str] = None,
+      docker_path: Optional[str] = None,
+      **data,
+  ):
+    """Initializes the ContainerCodeExecutor.
+
+    Args:
+      base_url: Optional. The base url of the user hosted Docker client.
+      image: The tag of the predefined image or custom image to run on the
+        container. Either docker_path or image must be set.
+      docker_path: The path to the directory containing the Dockerfile. If set,
+        build the image from the dockerfile path instead of using the predefined
+        image. Either docker_path or image must be set.
+      **data: The data to initialize the ContainerCodeExecutor.
+    """
+    if not image and not docker_path:
+      raise ValueError(
+          'Either image or docker_path must be set for ContainerCodeExecutor.'
+      )
+    if 'stateful' in data and data['stateful']:
+      raise ValueError('Cannot set `stateful=True` in ContainerCodeExecutor.')
+    if 'optimize_data_file' in data and data['optimize_data_file']:
+      raise ValueError(
+          'Cannot set `optimize_data_file=True` in ContainerCodeExecutor.'
+      )
+
+    super().__init__(**data)
+    self.base_url = base_url
+    self.image = image if image else DEFAULT_IMAGE_TAG
+    self.docker_path = os.path.abspath(docker_path) if docker_path else None
+
+    self._client = (
+        docker.from_env()
+        if not self.base_url
+        else docker.DockerClient(base_url=self.base_url)
+    )
+    # Initialize the container.
+    self.__init_container()
+
+    # Close the container when the on exit.
+    atexit.register(self.__cleanup_container)
+
+  @override
+  def execute_code(
+      self,
+      invocation_context: InvocationContext,
+      code_execution_input: CodeExecutionInput,
+  ) -> CodeExecutionResult:
+    output = ''
+    error = ''
+    exec_result = self._container.exec_run(
+        ['python3', '-c', code_execution_input.code],
+        demux=True,
+    )
+
+    if exec_result.output and exec_result.output[0]:
+      output = exec_result.output[0].decode('utf-8')
+    if (
+        exec_result.output
+        and len(exec_result.output) > 1
+        and exec_result.output[1]
+    ):
+      error = exec_result.output[1].decode('utf-8')
+
+    # Collect the final result.
+    return CodeExecutionResult(
+        stdout=output,
+        stderr=error,
+        output_files=[],
+    )
+
+  def _build_docker_image(self):
+    """Builds the Docker image."""
+    if not self.docker_path:
+      raise ValueError('Docker path is not set.')
+    if not os.path.exists(self.docker_path):
+      raise FileNotFoundError(f'Invalid Docker path: {self.docker_path}')
+
+    print('Building Docker image...')
+    self._client.images.build(
+        path=self.docker_path,
+        tag=self.image,
+        rm=True,
+    )
+    print(f'Docker image: {self.image} built.')
+
+  def _verify_python_installation(self):
+    """Verifies the container has python3 installed."""
+    exec_result = self._container.exec_run(['which', 'python3'])
+    if exec_result.exit_code != 0:
+      raise ValueError('python3 is not installed in the container.')
+
+  def __init_container(self):
+    """Initializes the container."""
+    if not self._client:
+      raise RuntimeError('Docker client is not initialized.')
+
+    if self.docker_path:
+      self._build_docker_image()
+
+    print('Starting container for ContainerCodeExecutor...')
+    self._container = self._client.containers.run(
+        image=self.image,
+        detach=True,
+        tty=True,
+    )
+    print(f'Container {self._container.id} started.')
+
+    # Verify the container is able to run python3.
+    self._verify_python_installation()
+
+  def __cleanup_container(self):
+    """Closes the container on exit."""
+    if not self._container:
+      return
+
+    print('[Cleanup] Stopping the container...')
+    self._container.stop()
+    self._container.remove()
+    print(f'Container {self._container.id} stopped and removed.')

google/adk/code_executors/unsafe_local_code_executor.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 contextlib import redirect_stdout
+import io
+
+from pydantic import Field
+from typing_extensions import override
+
+from ..agents.invocation_context import InvocationContext
+from .base_code_executor import BaseCodeExecutor
+from .code_execution_utils import CodeExecutionInput
+from .code_execution_utils import CodeExecutionResult
+
+
+class UnsafeLocalCodeExecutor(BaseCodeExecutor):
+  """A code executor that unsafely execute code in the current local context."""
+
+  # Overrides the BaseCodeExecutor attribute: this executor cannot be stateful.
+  stateful: bool = Field(default=False, frozen=True, exclude=True)
+
+  # Overrides the BaseCodeExecutor attribute: this executor cannot
+  # optimize_data_file.
+  optimize_data_file: bool = Field(default=False, frozen=True, exclude=True)
+
+  def __init__(self, **data):
+    """Initializes the UnsafeLocalCodeExecutor."""
+    if 'stateful' in data and data['stateful']:
+      raise ValueError('Cannot set `stateful=True` in UnsafeLocalCodeExecutor.')
+    if 'optimize_data_file' in data and data['optimize_data_file']:
+      raise ValueError(
+          'Cannot set `optimize_data_file=True` in UnsafeLocalCodeExecutor.'
+      )
+    super().__init__(**data)
+
+  @override
+  def execute_code(
+      self,
+      invocation_context: InvocationContext,
+      code_execution_input: CodeExecutionInput,
+  ) -> CodeExecutionResult:
+    # Execute the code.
+    output = ''
+    error = ''
+    try:
+      globals_ = {}
+      locals_ = {}
+      stdout = io.StringIO()
+      with redirect_stdout(stdout):
+        exec(code_execution_input.code, globals_, locals_)
+      output = stdout.getvalue()
+    except Exception as e:
+      error = str(e)
+
+    # Collect the final result.
+    return CodeExecutionResult(
+        stdout=output,
+        stderr=error,
+        output_files=[],
+    )

google/adk/code_executors/vertex_ai_code_executor.py

@@ -0,0 +1,234 @@
+# 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 datetime
+import mimetypes
+import os
+from typing import Any, Optional
+
+from typing_extensions import override
+from vertexai.preview.extensions import Extension
+
+from ..agents.invocation_context import InvocationContext
+from .base_code_executor import BaseCodeExecutor
+from .code_execution_utils import CodeExecutionInput
+from .code_execution_utils import CodeExecutionResult
+from .code_execution_utils import File
+
+_SUPPORTED_IMAGE_TYPES = ['png', 'jpg', 'jpeg']
+_SUPPORTED_DATA_FILE_TYPES = ['csv']
+
+_IMPORTED_LIBRARIES = '''
+import io
+import math
+import re
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import scipy
+
+def crop(s: str, max_chars: int = 64) -> str:
+  """Crops a string to max_chars characters."""
+  return s[: max_chars - 3] + '...' if len(s) > max_chars else s
+
+
+def explore_df(df: pd.DataFrame) -> None:
+  """Prints some information about a pandas DataFrame."""
+
+  with pd.option_context(
+      'display.max_columns', None, 'display.expand_frame_repr', False
+  ):
+    # Print the column names to never encounter KeyError when selecting one.
+    df_dtypes = df.dtypes
+
+    # Obtain information about data types and missing values.
+    df_nulls = (len(df) - df.isnull().sum()).apply(
+        lambda x: f'{x} / {df.shape[0]} non-null'
+    )
+
+    # Explore unique total values in columns using `.unique()`.
+    df_unique_count = df.apply(lambda x: len(x.unique()))
+
+    # Explore unique values in columns using `.unique()`.
+    df_unique = df.apply(lambda x: crop(str(list(x.unique()))))
+
+    df_info = pd.concat(
+        (
+            df_dtypes.rename('Dtype'),
+            df_nulls.rename('Non-Null Count'),
+            df_unique_count.rename('Unique Values Count'),
+            df_unique.rename('Unique Values'),
+        ),
+        axis=1,
+    )
+    df_info.index.name = 'Columns'
+    print(f"""Total rows: {df.shape[0]}
+Total columns: {df.shape[1]}
+
+{df_info}""")
+'''
+
+
+def _get_code_interpreter_extension(resource_name: str = None):
+  """Returns: Load or create the code interpreter extension."""
+  if not resource_name:
+    resource_name = os.environ.get('CODE_INTERPRETER_EXTENSION_NAME')
+  if resource_name:
+    new_code_interpreter = Extension(resource_name)
+  else:
+    print('No CODE_INTERPRETER_ID found in the environment. Create a new one.')
+    new_code_interpreter = Extension.from_hub('code_interpreter')
+    os.environ['CODE_INTERPRETER_EXTENSION_NAME'] = (
+        new_code_interpreter.gca_resource.name
+    )
+  return new_code_interpreter
+
+
+class VertexAiCodeExecutor(BaseCodeExecutor):
+  """A code executor that uses Vertex Code Interpreter Extension to execute code.
+
+  Attributes:
+    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/extensions/456
+  """
+
+  resource_name: str = None
+  """
+  If set, load the existing resource name of the code interpreter extension
+  instead of creating a new one.
+  Format: projects/123/locations/us-central1/extensions/456
+  """
+
+  _code_interpreter_extension: Extension
+
+  def __init__(
+      self,
+      resource_name: str = None,
+      **data,
+  ):
+    """Initializes the VertexAiCodeExecutor.
+
+    Args:
+      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/extensions/456
+      **data: Additional keyword arguments to be passed to the base class.
+    """
+    super().__init__(**data)
+    self.resource_name = resource_name
+    self._code_interpreter_extension = _get_code_interpreter_extension(
+        self.resource_name
+    )
+
+  @override
+  def execute_code(
+      self,
+      invocation_context: InvocationContext,
+      code_execution_input: CodeExecutionInput,
+  ) -> CodeExecutionResult:
+    # Execute the code.
+    code_execution_result = self._execute_code_interpreter(
+        self._get_code_with_imports(code_execution_input.code),
+        code_execution_input.input_files,
+        code_execution_input.execution_id,
+    )
+
+    # Save output file as artifacts.
+    current_timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
+    file_name_prefix = '%s_' % str(current_timestamp)
+    saved_files = []
+    file_count = 0
+    for output_file in code_execution_result['output_files']:
+      file_type = output_file['name'].split('.')[-1]
+      file_name = file_name_prefix + '%d.%s' % (file_count, file_type)
+      if file_type in _SUPPORTED_IMAGE_TYPES:
+        file_count += 1
+        saved_files.append(
+            File(
+                name='plot_' + file_name,
+                content=output_file['contents'],
+                mime_type=f'image/{file_type}',
+            )
+        )
+      elif file_type in _SUPPORTED_DATA_FILE_TYPES:
+        file_count += 1
+        saved_files.append(
+            File(
+                name='data_' + file_name,
+                content=output_file['contents'],
+                mime_type=f'text/{file_type}',
+            )
+        )
+      else:
+        mime_type, _ = mimetypes.guess_type(file_name)
+        saved_files.append(
+            File(
+                name=file_name,
+                content=output_file['contents'],
+                mime_type=mime_type,
+            )
+        )
+
+    # Collect the final result.
+    return CodeExecutionResult(
+        stdout=code_execution_result.get('execution_result', ''),
+        stderr=code_execution_result.get('execution_error', ''),
+        output_files=saved_files,
+    )
+
+  def _execute_code_interpreter(
+      self,
+      code: str,
+      input_files: Optional[list[File]] = None,
+      session_id: Optional[str] = None,
+  ) -> dict[str, Any]:
+    """Executes the code interpreter extension.
+
+    Args:
+      code: The code to execute.
+      input_files: The input files to execute the code with.
+      session_id: The session ID to execute the code with.
+
+    Returns:
+      The response from the code interpreter extension.
+    """
+    operation_params = {'code': code}
+    if input_files:
+      operation_params['files'] = [
+          {'name': f.name, 'contents': f.content} for f in input_files
+      ]
+    if session_id:
+      operation_params['session_id'] = session_id
+    response = self._code_interpreter_extension.execute(
+        operation_id='execute',
+        operation_params=operation_params,
+    )
+    return response
+
+  def _get_code_with_imports(self, code: str) -> str:
+    """Builds the code string with built-in imports.
+
+    Args:
+      code: The code to execute.
+
+    Returns:
+      The code string with built-in imports.
+    """
+    return f"""
+{_IMPORTED_LIBRARIES}
+
+{code}
+"""