Lspinheiro/chore/migrate azure executor autogen ext (#3652)

* migrate code executor and tests * update extras * update dependencies and examples * fix imports * fix uv lock * add code_executor to toctree --------- Co-authored-by: Leonardo Pinheiro <[email protected]> Co-authored-by: Eric Zhu <[email protected]>
microsoft · Oct 8, 2024 · 53e5951 · 53e5951
1 parent ffb16d5
commit 53e5951
Show file tree

Hide file tree

Showing 12 changed files with 319 additions and 246 deletions.
diff --git a/python/README.md b/python/README.md
@@ -8,7 +8,7 @@ See [`autogen-core`](./packages/autogen-core/) package for main functionality.
 **TL;DR**, run all checks with:
 
 ```sh
-uv sync
+uv sync  --all-extras
 source .venv/bin/activate
 poe check
 ```
@@ -24,7 +24,7 @@ in the current directory,
 run:
 
 ```sh
-uv sync
+uv sync  --all-extras
 source .venv/bin/activate
 ```
 

diff --git a/...ckages/autogen-core/docs/src/core-user-guide/cookbook/azure-container-code-executor.ipynb b/...ckages/autogen-core/docs/src/core-user-guide/cookbook/azure-container-code-executor.ipynb
@@ -23,7 +23,7 @@
     "\n",
     "## AzureContainerCodeExecutor\n",
     "\n",
-    "The {py:class}`~autogen_core.components.code_executor.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",
+    "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",
@@ -65,7 +65,8 @@
     "\n",
     "from anyio import open_file\n",
     "from autogen_core.base import CancellationToken\n",
-    "from autogen_core.components.code_executor import AzureContainerCodeExecutor, CodeBlock\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"
    ]
   },

diff --git a/python/packages/autogen-core/docs/src/reference/index.md b/python/packages/autogen-core/docs/src/reference/index.md
@@ -40,6 +40,7 @@ python/autogen_core/autogen_core
 
 python/autogen_ext/autogen_ext
 python/autogen_ext/autogen_ext.tools
+python/autogen_ext/autogen_ext.code_executor
 ```
 
 ::::{grid} 1 2 2 3

diff --git a/python/packages/autogen-core/pyproject.toml b/python/packages/autogen-core/pyproject.toml
@@ -22,7 +22,6 @@ dependencies = [
     "grpcio~=1.62.0",
     "protobuf~=4.25.1",
     "tiktoken",
-    "azure-core",
     "docker~=7.0",
     "opentelemetry-api~=1.27.0",
     "asyncio_atexit"

diff --git a/python/packages/autogen-core/src/autogen_core/components/code_executor/__init__.py b/python/packages/autogen-core/src/autogen_core/components/code_executor/__init__.py
@@ -1,23 +1,37 @@
 from ._base import CodeBlock, CodeExecutor, CodeResult
-from ._func_with_reqs import Alias, FunctionWithRequirements, Import, ImportFromModule, with_requirements
-from ._impl.azure_container_code_executor import AzureContainerCodeExecutor
+from ._func_with_reqs import (
+    Alias,
+    FunctionWithRequirements,
+    FunctionWithRequirementsStr,
+    Import,
+    ImportFromModule,
+    build_python_functions_file,
+    to_stub,
+    with_requirements,
+)
 from ._impl.command_line_code_result import CommandLineCodeResult
 from ._impl.docker_command_line_code_executor import DockerCommandLineCodeExecutor
 from ._impl.local_commandline_code_executor import LocalCommandLineCodeExecutor
+from ._impl.utils import get_required_packages, lang_to_cmd
 from ._utils import extract_markdown_code_blocks
 
 __all__ = [
-    "AzureContainerCodeExecutor",
     "LocalCommandLineCodeExecutor",
     "CommandLineCodeResult",
     "CodeBlock",
-    "CodeResult",
     "CodeExecutor",
+    "CodeResult",
     "Alias",
     "ImportFromModule",
     "Import",
     "FunctionWithRequirements",
+    "FunctionWithRequirementsStr",
     "with_requirements",
+    "to_stub",
     "extract_markdown_code_blocks",
+    "get_required_packages",
+    "build_python_functions_file",
     "DockerCommandLineCodeExecutor",
+    "get_required_packages",
+    "lang_to_cmd",
 ]
diff --git a/python/packages/autogen-core/tests/execution/test_user_defined_functions.py b/python/packages/autogen-core/tests/execution/test_user_defined_functions.py
@@ -8,13 +8,11 @@
 import pytest
 from autogen_core.base import CancellationToken
 from autogen_core.components.code_executor import (
-    AzureContainerCodeExecutor,
     CodeBlock,
     FunctionWithRequirements,
     LocalCommandLineCodeExecutor,
     with_requirements,
 )
-from azure.identity import DefaultAzureCredential
 
 ENVIRON_KEY_AZURE_POOL_ENDPOINT = "AZURE_POOL_ENDPOINT"
 
@@ -78,32 +76,35 @@ async def test_can_load_function_with_reqs() -> None:
         assert result.exit_code == 0
 
 
-@pytest.mark.skipif(
-    not POOL_ENDPOINT,
-    reason="do not run if pool endpoint is not defined",
-)
-@pytest.mark.asyncio
-async def test_azure_can_load_function_with_reqs() -> None:
-    assert POOL_ENDPOINT is not None
-    cancellation_token = CancellationToken()
-    azure_executor = AzureContainerCodeExecutor(
-        pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential(), functions=[load_data]
-    )
-    # AzureContainerCodeExecutor doesn't use the functions module import
-    code = """import polars
+def test_local_formatted_prompt() -> None:
+    assert_str = '''def add_two_numbers(a: int, b: int) -> int:
+    """Add two numbers together."""
+'''
+    with tempfile.TemporaryDirectory() as temp_dir:
+        executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, functions=[add_two_numbers])
+
+        result = executor.format_functions_for_prompt()
+        assert assert_str in result
 
-# Get first row's name
-data = load_data()
-print(data['name'][0])"""
 
-    azure_result = await azure_executor.execute_code_blocks(
-        code_blocks=[
-            CodeBlock(language="python", code=code),
-        ],
-        cancellation_token=cancellation_token,
+def test_local_formatted_prompt_str_func() -> None:
+    func = FunctionWithRequirements.from_str(
+        '''
+def add_two_numbers(a: int, b: int) -> int:
+    """Add two numbers together."""
+    return a + b
+'''
     )
-    assert azure_result.output == "John\n"
-    assert azure_result.exit_code == 0
+
+    assert_str = '''def add_two_numbers(a: int, b: int) -> int:
+    """Add two numbers together."""
+'''
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, functions=[func])
+
+        result = executor.format_functions_for_prompt()
+        assert assert_str in result
 
 
 @pytest.mark.asyncio
@@ -124,31 +125,6 @@ async def test_can_load_function() -> None:
         assert result.exit_code == 0
 
 
-@pytest.mark.skipif(
-    not POOL_ENDPOINT,
-    reason="do not run if pool endpoint is not defined",
-)
-@pytest.mark.asyncio
-async def test_azure_can_load_function() -> None:
-    assert POOL_ENDPOINT is not None
-
-    cancellation_token = CancellationToken()
-    azure_executor = AzureContainerCodeExecutor(
-        pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential(), functions=[add_two_numbers]
-    )
-    # AzureContainerCodeExecutor doesn't use the functions module import
-    code = """print(add_two_numbers(1, 2))"""
-
-    azure_result = await azure_executor.execute_code_blocks(
-        code_blocks=[
-            CodeBlock(language="python", code=code),
-        ],
-        cancellation_token=cancellation_token,
-    )
-    assert azure_result.output == "3\n"
-    assert azure_result.exit_code == 0
-
-
 @pytest.mark.asyncio
 async def test_fails_for_function_incorrect_import() -> None:
     with tempfile.TemporaryDirectory() as temp_dir:
@@ -166,30 +142,6 @@ async def test_fails_for_function_incorrect_import() -> None:
             )
 
 
-@pytest.mark.skipif(
-    not POOL_ENDPOINT,
-    reason="do not run if pool endpoint is not defined",
-)
-@pytest.mark.asyncio
-async def test_azure_fails_for_function_incorrect_import() -> None:
-    assert POOL_ENDPOINT is not None
-    cancellation_token = CancellationToken()
-    azure_executor = AzureContainerCodeExecutor(
-        pool_management_endpoint=POOL_ENDPOINT,
-        credential=DefaultAzureCredential(),
-        functions=[function_incorrect_import],
-    )
-    code = """function_incorrect_import()"""
-
-    with pytest.raises(ValueError):
-        await azure_executor.execute_code_blocks(
-            code_blocks=[
-                CodeBlock(language="python", code=code),
-            ],
-            cancellation_token=cancellation_token,
-        )
-
-
 @pytest.mark.asyncio
 async def test_fails_for_function_incorrect_dep() -> None:
     with tempfile.TemporaryDirectory() as temp_dir:
@@ -207,73 +159,6 @@ async def test_fails_for_function_incorrect_dep() -> None:
             )
 
 
-@pytest.mark.skipif(
-    not POOL_ENDPOINT,
-    reason="do not run if pool endpoint is not defined",
-)
-@pytest.mark.asyncio
-async def test_azure_fails_for_function_incorrect_dep() -> None:
-    assert POOL_ENDPOINT is not None
-    cancellation_token = CancellationToken()
-    azure_executor = AzureContainerCodeExecutor(
-        pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential(), functions=[function_incorrect_dep]
-    )
-    code = """function_incorrect_dep()"""
-
-    with pytest.raises(ValueError):
-        await azure_executor.execute_code_blocks(
-            code_blocks=[
-                CodeBlock(language="python", code=code),
-            ],
-            cancellation_token=cancellation_token,
-        )
-
-
-def test_formatted_prompt() -> None:
-    assert_str = '''def add_two_numbers(a: int, b: int) -> int:
-    """Add two numbers together."""
-'''
-    with tempfile.TemporaryDirectory() as temp_dir:
-        executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, functions=[add_two_numbers])
-
-        result = executor.format_functions_for_prompt()
-        assert assert_str in result
-
-    azure_executor = AzureContainerCodeExecutor(
-        pool_management_endpoint=DUMMY_POOL_ENDPOINT, credential=DefaultAzureCredential(), functions=[add_two_numbers]
-    )
-
-    azure_result = azure_executor.format_functions_for_prompt()
-    assert assert_str in azure_result
-
-
-def test_formatted_prompt_str_func() -> None:
-    func = FunctionWithRequirements.from_str(
-        '''
-def add_two_numbers(a: int, b: int) -> int:
-    """Add two numbers together."""
-    return a + b
-'''
-    )
-
-    assert_str = '''def add_two_numbers(a: int, b: int) -> int:
-    """Add two numbers together."""
-'''
-
-    with tempfile.TemporaryDirectory() as temp_dir:
-        executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, functions=[func])
-
-        result = executor.format_functions_for_prompt()
-        assert assert_str in result
-
-    azure_executor = AzureContainerCodeExecutor(
-        pool_management_endpoint=DUMMY_POOL_ENDPOINT, credential=DefaultAzureCredential(), functions=[func]
-    )
-
-    azure_result = azure_executor.format_functions_for_prompt()
-    assert assert_str in azure_result
-
-
 @pytest.mark.asyncio
 async def test_can_load_str_function_with_reqs() -> None:
     func = FunctionWithRequirements.from_str(
@@ -300,36 +185,6 @@ def add_two_numbers(a: int, b: int) -> int:
         assert result.exit_code == 0
 
 
-@pytest.mark.skipif(
-    not POOL_ENDPOINT,
-    reason="do not run if pool endpoint is not defined",
-)
-@pytest.mark.asyncio
-async def test_azure_can_load_str_function_with_reqs() -> None:
-    assert POOL_ENDPOINT is not None
-    cancellation_token = CancellationToken()
-    func = FunctionWithRequirements.from_str(
-        '''
-def add_two_numbers(a: int, b: int) -> int:
-    """Add two numbers together."""
-    return a + b
-'''
-    )
-    azure_executor = AzureContainerCodeExecutor(
-        pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential(), functions=[func]
-    )
-    code = """print(add_two_numbers(1, 2))"""
-
-    azure_result = await azure_executor.execute_code_blocks(
-        code_blocks=[
-            CodeBlock(language="python", code=code),
-        ],
-        cancellation_token=cancellation_token,
-    )
-    assert azure_result.output == "3\n"
-    assert azure_result.exit_code == 0
-
-
 def test_cant_load_broken_str_function_with_reqs() -> None:
     with pytest.raises(ValueError):
         _ = FunctionWithRequirements.from_str(
@@ -365,35 +220,3 @@ def add_two_numbers(a: int, b: int) -> int:
         )
         assert "TypeError: unsupported operand type(s) for +:" in result.output
         assert result.exit_code == 1
-
-
-@pytest.mark.skipif(
-    not POOL_ENDPOINT,
-    reason="do not run if pool endpoint is not defined",
-)
-@pytest.mark.asyncio
-async def test_azure_cant_run_broken_str_function_with_reqs() -> None:
-    assert POOL_ENDPOINT is not None
-    cancellation_token = CancellationToken()
-    func = FunctionWithRequirements.from_str(
-        '''
-def add_two_numbers(a: int, b: int) -> int:
-    """Add two numbers together."""
-    return a + b
-'''
-    )
-
-    azure_executor = AzureContainerCodeExecutor(
-        pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential(), functions=[func]
-    )
-    code = """print(add_two_numbers(object(), False))"""
-
-    azure_result = await azure_executor.execute_code_blocks(
-        code_blocks=[
-            CodeBlock(language="python", code=code),
-        ],
-        cancellation_token=cancellation_token,
-    )
-    # result.output = result.output.encode().decode('unicode_escape')
-    assert "TypeError: unsupported operand type(s) for +:" in azure_result.output
-    assert azure_result.exit_code == 1
diff --git a/python/packages/autogen-ext/pyproject.toml b/python/packages/autogen-ext/pyproject.toml
@@ -19,7 +19,8 @@ dependencies = [
 
 
 [project.optional-dependencies]
-langchain = ["langchain >= 0.3.1"]
+langchain-tools = ["langchain >= 0.3.1"]
+azure-code-executor = ["azure-core"]
 
 [tool.hatch.build.targets.wheel]
 packages = ["src/autogen_ext"]

diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executor/aca_dynamic_sessions/__init__.py b/python/packages/autogen-ext/src/autogen_ext/code_executor/aca_dynamic_sessions/__init__.py
@@ -0,0 +1,5 @@
+from ._azure_container_code_executor import AzureContainerCodeExecutor
+
+__all__ = [
+    "AzureContainerCodeExecutor",
+]