{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Azure Container Code Executor\n",
    "\n",
    "This guide will explain the Azure Container Apps dynamic sessions in Azure Container Apps and show you how to use the Azure Container Code Executor class.\n",
    "\n",
    "\n",
    "## Azure Container Apps dynamic sessions\n",
    "\n",
    "The [Azure Container Apps dynamic sessions](https://learn.microsoft.com/en-us/azure/container-apps/sessions) is a component in the Azure Container Apps service. The environment is hosted on remote Azure instances and will not execute any code locally. The interpreter is capable of executing python code in a jupyter environment with a pre-installed base of commonly used packages. [Custom environments](https://learn.microsoft.com/en-us/azure/container-apps/sessions-custom-container) can be created by users for their applications. Files can additionally be [uploaded to, or downloaded from](https://learn.microsoft.com/en-us/azure/container-apps/sessions-code-interpreter#upload-a-file-to-a-session) each session.\n",
    "\n",
    "The code interpreter can run multiple sessions of code, each of which are delineated by a session identifier string.\n",
    "\n",
    "### Create a Container Apps Session Pool\n",
    "\n",
    "In your Azure portal, create a new `Container App Session Pool` resource with the pool type set to `Python code interpreter` and note the `Pool management endpoint`. The format for the endpoint should be something like `https://{region}.dynamicsessions.io/subscriptions/{subscription_id}/resourceGroups/{resource_group_name}/sessionPools/{session_pool_name}`.\n",
    "\n",
    "Alternatively, you can use the [Azure CLI to create a session pool.](https://learn.microsoft.com/en-us/azure/container-apps/sessions-code-interpreter#create-a-session-pool-with-azure-cli)\n",
    "\n",
    "## AzureContainerCodeExecutor\n",
    "\n",
    "The {py:class}`~autogen_ext.code_executor.aca_dynamic_sessions.AzureContainerCodeExecutor` class is a python code executor that creates and executes arbitrary python code on a default Serverless code interpreter session. Its interface is as follows\n",
    "\n",
    "### Initialization\n",
    "\n",
    "First, you will need to find or create a credentialing object that implements the {py:class}`~autogen_core.components.code_executor.TokenProvider` interface. This is any object that implements the following function\n",
    "```python\n",
    "def get_token(\n",
    "    self, *scopes: str, claims: Optional[str] = None, tenant_id: Optional[str] = None, **kwargs: Any\n",
    ") -> azure.core.credentials.AccessToken\n",
    "```\n",
    "An example of such an object is the [azure.identity.DefaultAzureCredential](https://learn.microsoft.com/en-us/python/api/azure-identity/azure.identity.defaultazurecredential?view=azure-python) class.\n",
    "\n",
    "Lets start by installing that"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install azure.identity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, lets import all the necessary modules and classes for our code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tempfile\n",
    "\n",
    "from anyio import open_file\n",
    "from autogen_core.base import CancellationToken\n",
    "from autogen_core.components.code_executor import CodeBlock\n",
    "from autogen_ext.code_executor.aca_dynamic_sessions import AzureContainerCodeExecutor\n",
    "from azure.identity import DefaultAzureCredential"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we create our Azure code executor and run some test code along with verification that it ran correctly. We'll create the executor with a temporary working directory to ensure a clean environment as we show how to use each feature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cancellation_token = CancellationToken()\n",
    "POOL_MANAGEMENT_ENDPOINT = \"...\"\n",
    "\n",
    "with tempfile.TemporaryDirectory() as temp_dir:\n",
    "    executor = AzureContainerCodeExecutor(\n",
    "        pool_management_endpoint=POOL_MANAGEMENT_ENDPOINT, credential=DefaultAzureCredential(), work_dir=temp_dir\n",
    "    )\n",
    "\n",
    "    code_blocks = [CodeBlock(code=\"import sys; print('hello world!')\", language=\"python\")]\n",
    "    code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n",
    "    assert code_result.exit_code == 0 and \"hello world!\" in code_result.output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, lets try uploading some files and verifying their integrity. All files uploaded to the Serverless code interpreter is uploaded into the `/mnt/data` directory. All downloadable files must also be placed in the directory. By default, the current working directory for the code executor is set to `/mnt/data`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tempfile.TemporaryDirectory() as temp_dir:\n",
    "    test_file_1 = \"test_upload_1.txt\"\n",
    "    test_file_1_contents = \"test1 contents\"\n",
    "    test_file_2 = \"test_upload_2.txt\"\n",
    "    test_file_2_contents = \"test2 contents\"\n",
    "\n",
    "    async with await open_file(os.path.join(temp_dir, test_file_1), \"w\") as f:  # type: ignore[syntax]\n",
    "        await f.write(test_file_1_contents)\n",
    "    async with await open_file(os.path.join(temp_dir, test_file_2), \"w\") as f:  # type: ignore[syntax]\n",
    "        await f.write(test_file_2_contents)\n",
    "\n",
    "    assert os.path.isfile(os.path.join(temp_dir, test_file_1))\n",
    "    assert os.path.isfile(os.path.join(temp_dir, test_file_2))\n",
    "\n",
    "    executor = AzureContainerCodeExecutor(\n",
    "        pool_management_endpoint=POOL_MANAGEMENT_ENDPOINT, credential=DefaultAzureCredential(), work_dir=temp_dir\n",
    "    )\n",
    "    await executor.upload_files([test_file_1, test_file_2], cancellation_token)\n",
    "\n",
    "    file_list = await executor.get_file_list(cancellation_token)\n",
    "    assert test_file_1 in file_list\n",
    "    assert test_file_2 in file_list\n",
    "\n",
    "    code_blocks = [\n",
    "        CodeBlock(\n",
    "            code=f\"\"\"\n",
    "with open(\"{test_file_1}\") as f:\n",
    "  print(f.read())\n",
    "with open(\"{test_file_2}\") as f:\n",
    "  print(f.read())\n",
    "\"\"\",\n",
    "            language=\"python\",\n",
    "        )\n",
    "    ]\n",
    "    code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n",
    "    assert code_result.exit_code == 0\n",
    "    assert test_file_1_contents in code_result.output\n",
    "    assert test_file_2_contents in code_result.output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Downloading files works in a similar way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tempfile.TemporaryDirectory() as temp_dir:\n",
    "    test_file_1 = \"test_upload_1.txt\"\n",
    "    test_file_1_contents = \"test1 contents\"\n",
    "    test_file_2 = \"test_upload_2.txt\"\n",
    "    test_file_2_contents = \"test2 contents\"\n",
    "\n",
    "    assert not os.path.isfile(os.path.join(temp_dir, test_file_1))\n",
    "    assert not os.path.isfile(os.path.join(temp_dir, test_file_2))\n",
    "\n",
    "    executor = AzureContainerCodeExecutor(\n",
    "        pool_management_endpoint=POOL_MANAGEMENT_ENDPOINT, credential=DefaultAzureCredential(), work_dir=temp_dir\n",
    "    )\n",
    "\n",
    "    code_blocks = [\n",
    "        CodeBlock(\n",
    "            code=f\"\"\"\n",
    "with open(\"{test_file_1}\", \"w\") as f:\n",
    "  f.write(\"{test_file_1_contents}\")\n",
    "with open(\"{test_file_2}\", \"w\") as f:\n",
    "  f.write(\"{test_file_2_contents}\")\n",
    "\"\"\",\n",
    "            language=\"python\",\n",
    "        ),\n",
    "    ]\n",
    "    code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n",
    "    assert code_result.exit_code == 0\n",
    "\n",
    "    file_list = await executor.get_file_list(cancellation_token)\n",
    "    assert test_file_1 in file_list\n",
    "    assert test_file_2 in file_list\n",
    "\n",
    "    await executor.download_files([test_file_1, test_file_2], cancellation_token)\n",
    "\n",
    "    assert os.path.isfile(os.path.join(temp_dir, test_file_1))\n",
    "    async with await open_file(os.path.join(temp_dir, test_file_1), \"r\") as f:  # type: ignore[syntax]\n",
    "        content = await f.read()\n",
    "        assert test_file_1_contents in content\n",
    "    assert os.path.isfile(os.path.join(temp_dir, test_file_2))\n",
    "    async with await open_file(os.path.join(temp_dir, test_file_2), \"r\") as f:  # type: ignore[syntax]\n",
    "        content = await f.read()\n",
    "        assert test_file_2_contents in content"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### New Sessions\n",
    "\n",
    "Every instance of the {py:class}`~autogen_core.components.code_executor.AzureContainerCodeExecutor` class will have a unique session ID. Every call to a particular code executor will be executed on the same session until the {py:meth}`~autogen_core.components.code_executor.AzureContainerCodeExecutor.restart` function is called on it. Previous sessions cannot be reused.\n",
    "\n",
    "Here we'll run some code on the code session, restart it, then verify that a new session has been opened."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "executor = AzureContainerCodeExecutor(\n",
    "    pool_management_endpoint=POOL_MANAGEMENT_ENDPOINT, credential=DefaultAzureCredential()\n",
    ")\n",
    "\n",
    "code_blocks = [CodeBlock(code=\"x = 'abcdefg'\", language=\"python\")]\n",
    "code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n",
    "assert code_result.exit_code == 0\n",
    "\n",
    "code_blocks = [CodeBlock(code=\"print(x)\", language=\"python\")]\n",
    "code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n",
    "assert code_result.exit_code == 0 and \"abcdefg\" in code_result.output\n",
    "\n",
    "await executor.restart()\n",
    "code_blocks = [CodeBlock(code=\"print(x)\", language=\"python\")]\n",
    "code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n",
    "assert code_result.exit_code != 0 and \"NameError\" in code_result.output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Available Packages\n",
    "\n",
    "Each code execution instance is pre-installed with most of the commonly used packages. However, the list of available packages and versions are not available outside of the execution environment. The packages list on the environment can be retrieved by calling the {py:meth}`~autogen_core.components.code_executor.AzureContainerCodeExecutor.get_available_packages` function on the code executor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(executor.get_available_packages(cancellation_token))"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}