mirror of
https://github.com/microsoft/autogen.git
synced 2025-12-03 18:40:04 +00:00
Azure container code execution (#346)
* Azure container code execution * fix check errors. Remove unnecessary file * add TokenProvider class * update deps * addressing PR comments * update docstring * switch to aiohttp * fix client timeout
This commit is contained in:
peterychang
GitHub
parent
698a8f374b
commit
bc26ec3de4
@ -22,7 +22,8 @@ dependencies = [
|
||||
"types-aiofiles",
|
||||
"grpcio",
|
||||
"protobuf",
|
||||
"tiktoken"
|
||||
"tiktoken",
|
||||
"azure-core"
|
||||
]
|
||||
|
||||
[tool.hatch.envs.default]
|
||||
|
||||
@ -1,9 +1,11 @@
|
||||
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 ._impl.command_line_code_result import CommandLineCodeResult
|
||||
from ._impl.local_commandline_code_executor import LocalCommandLineCodeExecutor
|
||||
|
||||
__all__ = [
|
||||
"AzureContainerCodeExecutor",
|
||||
"LocalCommandLineCodeExecutor",
|
||||
"CommandLineCodeResult",
|
||||
"CodeBlock",
|
||||
|
||||
@ -41,6 +41,11 @@ class CodeExecutor(Protocol):
|
||||
|
||||
Returns:
|
||||
CodeResult: The result of the code execution.
|
||||
|
||||
Raises:
|
||||
ValueError: Errors in user inputs
|
||||
asyncio.TimeoutError: Code execution timeouts
|
||||
asyncio.CancelledError: CancellationToken evoked during execution
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@ -0,0 +1,290 @@
|
||||
# File based from: https://github.com/microsoft/autogen/blob/main/autogen/coding/local_commandline_code_executor.py
|
||||
# Credit to original authors
|
||||
|
||||
import asyncio
|
||||
from string import Template
|
||||
from typing import Any, Callable, ClassVar, List, Optional, Protocol, Sequence, Union
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
import aiohttp
|
||||
from azure.core.credentials import AccessToken
|
||||
|
||||
# from azure.mgmt.appcontainers import ContainerAppsAPIClient
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
from ....core import CancellationToken
|
||||
from .._base import CodeBlock, CodeExecutor, CodeResult
|
||||
from .._func_with_reqs import (
|
||||
FunctionWithRequirements,
|
||||
FunctionWithRequirementsStr,
|
||||
build_python_functions_file,
|
||||
to_stub,
|
||||
)
|
||||
from .utils import PYTHON_VARIANTS, get_required_packages, lang_to_cmd # type: ignore
|
||||
|
||||
__all__ = ("AzureContainerCodeExecutor", "TokenProvider")
|
||||
|
||||
A = ParamSpec("A")
|
||||
|
||||
|
||||
class TokenProvider(Protocol):
|
||||
def get_token(
|
||||
self, *scopes: str, claims: Optional[str] = None, tenant_id: Optional[str] = None, **kwargs: Any
|
||||
) -> AccessToken: ...
|
||||
|
||||
|
||||
class AzureContainerCodeExecutor(CodeExecutor):
|
||||
SUPPORTED_LANGUAGES: ClassVar[List[str]] = [
|
||||
"python",
|
||||
]
|
||||
FUNCTION_PROMPT_TEMPLATE: ClassVar[str] = """You have access to the following user defined functions.
|
||||
|
||||
$functions"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pool_management_endpoint: str,
|
||||
credential: TokenProvider,
|
||||
timeout: int = 60,
|
||||
functions: Sequence[
|
||||
Union[
|
||||
FunctionWithRequirements[Any, A],
|
||||
Callable[..., Any],
|
||||
FunctionWithRequirementsStr,
|
||||
]
|
||||
] = [],
|
||||
functions_module: str = "functions",
|
||||
persist_session: bool = False,
|
||||
):
|
||||
"""(Experimental) A code executor class that executes code through a an Azure
|
||||
Container Apps instance.
|
||||
|
||||
**This will execute LLM generated code on an Azure dynamic code container.**
|
||||
|
||||
The execution environment is similar to that of a jupyter notebook which allows for incremental code execution. The parameter functions are executed in order once at the beginning of each session. Each code block is then executed serially and in the order they are received. Each environment has a statically defined set of available packages which cannot be changed.
|
||||
Currently, attempting to use packages beyond what is available on the environment will result in an error. To get the list of supported packages, call the `get_available_packages` function.
|
||||
Currently the only supported language is Python.
|
||||
For Python code, use the language "python" for the code block.
|
||||
|
||||
Args:
|
||||
pool_management_endpoint (str): The azure container apps dynamic sessions endpoint.
|
||||
credential (TokenProvider): An object that implements the get_token function.
|
||||
timeout (int): The timeout for the execution of any single code block. Default is 60.
|
||||
functions (List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]]): A list of functions that are available to the code executor. Default is an empty list.
|
||||
persist_session (bool): True - reuse the same azure session ID until restart() is called. False - Refresh the azure session ID for every call to execute_code(). Default is False.
|
||||
"""
|
||||
|
||||
if timeout < 1:
|
||||
raise ValueError("Timeout must be greater than or equal to 1.")
|
||||
|
||||
if not functions_module.isidentifier():
|
||||
raise ValueError("Module name must be a valid Python identifier")
|
||||
|
||||
self._functions_module = functions_module
|
||||
|
||||
self._timeout = timeout
|
||||
|
||||
self._functions = functions
|
||||
self._func_code: str | None = None
|
||||
# Setup could take some time so we intentionally wait for the first code block to do it.
|
||||
if len(functions) > 0:
|
||||
self._setup_functions_complete = False
|
||||
else:
|
||||
self._setup_functions_complete = True
|
||||
|
||||
self._pool_management_endpoint = pool_management_endpoint
|
||||
self._access_token: str | None = None
|
||||
self._persist_session = persist_session
|
||||
self._uuid: UUID = uuid4()
|
||||
self._available_packages: set[str] | None = None
|
||||
self._credential: TokenProvider = credential
|
||||
|
||||
# TODO: expiration?
|
||||
def _ensure_access_token(self) -> None:
|
||||
if not self._access_token:
|
||||
scope = "https://dynamicsessions.io"
|
||||
self._access_token = self._credential.get_token(scope).token
|
||||
|
||||
def format_functions_for_prompt(self, prompt_template: str = FUNCTION_PROMPT_TEMPLATE) -> str:
|
||||
"""(Experimental) Format the functions for a prompt.
|
||||
|
||||
The template includes one variable:
|
||||
- `$functions`: The functions formatted as stubs with two newlines between each function.
|
||||
|
||||
Args:
|
||||
prompt_template (str): The prompt template. Default is the class default.
|
||||
|
||||
Returns:
|
||||
str: The formatted prompt.
|
||||
"""
|
||||
|
||||
template = Template(prompt_template)
|
||||
return template.substitute(
|
||||
functions="\n\n".join([to_stub(func) for func in self._functions]),
|
||||
)
|
||||
|
||||
@property
|
||||
def functions_module(self) -> str:
|
||||
"""(Experimental) The module name for the functions."""
|
||||
return self._functions_module
|
||||
|
||||
@property
|
||||
def functions(self) -> List[str]:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def timeout(self) -> int:
|
||||
"""(Experimental) The timeout for code execution."""
|
||||
return self._timeout
|
||||
|
||||
async def get_available_packages(self, cancellation_token: CancellationToken) -> set[str]:
|
||||
if self._available_packages is not None:
|
||||
return self._available_packages
|
||||
avail_pkgs = """
|
||||
import pkg_resources\n[d.project_name for d in pkg_resources.working_set]
|
||||
"""
|
||||
ret = await self._execute_code_dont_check_setup(
|
||||
[CodeBlock(code=avail_pkgs, language="python")], cancellation_token
|
||||
)
|
||||
if ret.exit_code != 0:
|
||||
raise ValueError(f"Failed to get list of available packages: {ret.output.strip()}")
|
||||
pkgs = ret.output.strip("[]")
|
||||
pkglist = pkgs.split(",\n")
|
||||
return {pkg.strip(" '") for pkg in pkglist}
|
||||
|
||||
async def _populate_available_packages(self, cancellation_token: CancellationToken) -> None:
|
||||
self._available_packages = await self.get_available_packages(cancellation_token)
|
||||
|
||||
async def _setup_functions(self, cancellation_token: CancellationToken) -> None:
|
||||
if not self._func_code:
|
||||
self._func_code = build_python_functions_file(self._functions)
|
||||
|
||||
# Check required function imports and packages
|
||||
lists_of_packages = [x.python_packages for x in self._functions if isinstance(x, FunctionWithRequirements)]
|
||||
# Should we also be checking the imports?
|
||||
|
||||
flattened_packages = [item for sublist in lists_of_packages for item in sublist]
|
||||
required_packages = set(flattened_packages)
|
||||
if self._available_packages is not None:
|
||||
missing_pkgs = set(required_packages - self._available_packages)
|
||||
if len(missing_pkgs) > 0:
|
||||
raise ValueError(f"Packages unavailable in environment: {missing_pkgs}")
|
||||
|
||||
# Attempt to load the function file to check for syntax errors, imports etc.
|
||||
exec_result = await self._execute_code_dont_check_setup(
|
||||
[CodeBlock(code=self._func_code, language="python")], cancellation_token
|
||||
)
|
||||
|
||||
if exec_result.exit_code != 0:
|
||||
raise ValueError(f"Functions failed to load: {exec_result.output.strip()}")
|
||||
|
||||
self._setup_functions_complete = True
|
||||
|
||||
async def execute_code_blocks(
|
||||
self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken
|
||||
) -> CodeResult:
|
||||
"""(Experimental) Execute the code blocks and return the result.
|
||||
|
||||
Args:
|
||||
code_blocks (List[CodeBlock]): The code blocks to execute.
|
||||
cancellation_token (CancellationToken): a token to cancel the operation
|
||||
|
||||
Returns:
|
||||
CodeResult: The result of the code execution."""
|
||||
if not self._persist_session:
|
||||
self.restart()
|
||||
if self._available_packages is None:
|
||||
await self._populate_available_packages(cancellation_token)
|
||||
if not self._setup_functions_complete:
|
||||
await self._setup_functions(cancellation_token)
|
||||
|
||||
return await self._execute_code_dont_check_setup(code_blocks, cancellation_token)
|
||||
|
||||
# The http call here should be replaced by an actual Azure client call once its available
|
||||
async def _execute_code_dont_check_setup(
|
||||
self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken
|
||||
) -> CodeResult:
|
||||
logs_all = ""
|
||||
exitcode = 0
|
||||
self._ensure_access_token()
|
||||
|
||||
# TODO: Better to use the client auth system rather than headers
|
||||
headers = {"Authorization": f"Bearer {self._access_token}"}
|
||||
properties = {
|
||||
"identifier": str(self._uuid),
|
||||
"codeInputType": "inline",
|
||||
"executionType": "synchronous",
|
||||
"pythonCode": "",
|
||||
"timeoutInSeconds": self._timeout,
|
||||
}
|
||||
timeout = aiohttp.ClientTimeout(total=float(self._timeout))
|
||||
async with aiohttp.ClientSession(timeout=timeout) as client:
|
||||
for code_block in code_blocks:
|
||||
lang, code = code_block.language, code_block.code
|
||||
lang = lang.lower()
|
||||
|
||||
if lang in PYTHON_VARIANTS:
|
||||
lang = "python"
|
||||
|
||||
if lang not in self.SUPPORTED_LANGUAGES:
|
||||
# In case the language is not supported, we return an error message.
|
||||
exitcode = 1
|
||||
logs_all += "\n" + f"unknown language {lang}"
|
||||
break
|
||||
|
||||
if self._available_packages is not None:
|
||||
req_pkgs = get_required_packages(code, lang)
|
||||
missing_pkgs = set(req_pkgs - self._available_packages)
|
||||
if len(missing_pkgs) > 0:
|
||||
# In case the code requires packages that are not available in the environment
|
||||
exitcode = 1
|
||||
logs_all += "\n" + f"Python packages unavailable in environment: {missing_pkgs}"
|
||||
break
|
||||
|
||||
properties["pythonCode"] = code_block.code
|
||||
|
||||
task = asyncio.create_task(
|
||||
client.post(
|
||||
self._pool_management_endpoint + "/python/execute",
|
||||
headers=headers,
|
||||
json={"properties": properties},
|
||||
)
|
||||
)
|
||||
|
||||
cancellation_token.link_future(task)
|
||||
try:
|
||||
response = await asyncio.wait_for(task, self._timeout)
|
||||
response.raise_for_status()
|
||||
data = await response.json()
|
||||
logs_all += data.get("stderr", "") + data.get("stdout", "")
|
||||
|
||||
if "Success" in data["status"]:
|
||||
logs_all += str(data["result"])
|
||||
elif "Failure" in data["status"]:
|
||||
exitcode = 1
|
||||
# This case is in the official code example https://github.com/Azure-Samples/container-apps-dynamic-sessions-samples/blob/dd2b3827bc8ea489b8f088654847239e2d51743f/autogen-python-webapi/aca_sessions_executor.py
|
||||
# I have not seen this case actually occur before
|
||||
if "error" in data:
|
||||
logs_all += f"\n{data['error']}"
|
||||
exitcode = 1
|
||||
|
||||
except asyncio.TimeoutError as e:
|
||||
logs_all += "\n Timeout"
|
||||
# e.add_note is only in py 3.11+
|
||||
raise asyncio.TimeoutError(logs_all) from e
|
||||
except asyncio.CancelledError as e:
|
||||
logs_all += "\n Cancelled"
|
||||
# e.add_note is only in py 3.11+
|
||||
raise asyncio.CancelledError(logs_all) from e
|
||||
except aiohttp.ClientResponseError as e:
|
||||
logs_all += "\nError while sending code block to endpoint"
|
||||
raise ConnectionError(logs_all) from e
|
||||
|
||||
return CodeResult(exit_code=exitcode, output=logs_all)
|
||||
|
||||
def restart(self) -> None:
|
||||
"""(Experimental) Restart the code executor."""
|
||||
self._uuid = uuid4()
|
||||
self._setup_functions_complete = False
|
||||
self._access_token = None
|
||||
self._available_packages = None
|
||||
@ -46,6 +46,24 @@ def silence_pip(code: str, lang: str) -> str:
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def get_required_packages(code: str, lang: str) -> set[str]:
|
||||
ret: set[str] = set()
|
||||
if lang == "python":
|
||||
regex = r"^! ?pip install(.*)$"
|
||||
else:
|
||||
return ret
|
||||
|
||||
# Find lines that start with pip install and make sure "-qqq" flag is added.
|
||||
lines = code.split("\n")
|
||||
for _, line in enumerate(lines):
|
||||
# use regex to find lines that start with pip install.
|
||||
match = re.search(regex, line)
|
||||
if match is not None:
|
||||
reqs = match.group(1).split(",")
|
||||
ret = {req.strip(" ") for req in reqs}
|
||||
return ret
|
||||
|
||||
|
||||
PYTHON_VARIANTS = ["python", "Python", "py"]
|
||||
|
||||
|
||||
|
||||
96
python/tests/execution/test_azure_container_code_executor.py
Normal file
96
python/tests/execution/test_azure_container_code_executor.py
Normal file
@ -0,0 +1,96 @@
|
||||
# File based from: https://github.com/microsoft/autogen/blob/main/test/coding/test_commandline_code_executor.py
|
||||
# Credit to original authors
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
from azure.identity import DefaultAzureCredential
|
||||
from agnext.components.code_executor import CodeBlock, AzureContainerCodeExecutor
|
||||
from agnext.core import CancellationToken
|
||||
|
||||
UNIX_SHELLS = ["bash", "sh", "shell"]
|
||||
WINDOWS_SHELLS = ["ps1", "pwsh", "powershell"]
|
||||
PYTHON_VARIANTS = ["python", "Python", "py"]
|
||||
|
||||
ENVIRON_KEY_AZURE_POOL_ENDPOINT = "AZURE_POOL_ENDPOINT"
|
||||
|
||||
POOL_ENDPOINT = os.getenv(ENVIRON_KEY_AZURE_POOL_ENDPOINT)
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not POOL_ENDPOINT,
|
||||
reason="do not run if pool endpoint is not defined",
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_code() -> None:
|
||||
assert POOL_ENDPOINT is not None
|
||||
cancellation_token = CancellationToken()
|
||||
executor = AzureContainerCodeExecutor(pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential())
|
||||
|
||||
# Test single code block.
|
||||
code_blocks = [CodeBlock(code="import sys; print('hello world!')", language="python")]
|
||||
code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)
|
||||
assert code_result.exit_code == 0 and "hello world!" in code_result.output
|
||||
|
||||
# Test multiple code blocks.
|
||||
code_blocks = [
|
||||
CodeBlock(code="import sys; print('hello world!')", language="python"),
|
||||
CodeBlock(code="a = 100 + 100; print(a)", language="python"),
|
||||
]
|
||||
code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)
|
||||
assert (
|
||||
code_result.exit_code == 0
|
||||
and "hello world!" in code_result.output
|
||||
and "200" in code_result.output
|
||||
)
|
||||
|
||||
# Test bash script.
|
||||
if sys.platform not in ["win32"]:
|
||||
code_blocks = [CodeBlock(code="echo 'hello world!'", language="bash")]
|
||||
code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)
|
||||
assert "unknown language" in code_result.output
|
||||
assert code_result.exit_code == 1
|
||||
|
||||
# Test running code.
|
||||
file_lines = ["import sys", "print('hello world!')", "a = 100 + 100", "print(a)"]
|
||||
code_blocks = [CodeBlock(code="\n".join(file_lines), language="python")]
|
||||
code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)
|
||||
assert (
|
||||
code_result.exit_code == 0
|
||||
and "hello world!" in code_result.output
|
||||
and "200" in code_result.output
|
||||
)
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not POOL_ENDPOINT,
|
||||
reason="do not run if pool endpoint is not defined",
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_container_code_executor_timeout() -> None:
|
||||
assert POOL_ENDPOINT is not None
|
||||
cancellation_token = CancellationToken()
|
||||
executor = AzureContainerCodeExecutor(pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential(), timeout=1)
|
||||
code_blocks = [CodeBlock(code="import time; time.sleep(10); print('hello world!')", language="python")]
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
await executor.execute_code_blocks(code_blocks, cancellation_token)
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not POOL_ENDPOINT,
|
||||
reason="do not run if pool endpoint is not defined",
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_container_code_executor_cancellation() -> None:
|
||||
assert POOL_ENDPOINT is not None
|
||||
cancellation_token = CancellationToken()
|
||||
executor = AzureContainerCodeExecutor(pool_management_endpoint=POOL_ENDPOINT, credential=DefaultAzureCredential())
|
||||
code_blocks = [CodeBlock(code="import time; time.sleep(10); print('hello world!')", language="python")]
|
||||
|
||||
coro = executor.execute_code_blocks(code_blocks, cancellation_token)
|
||||
|
||||
await asyncio.sleep(1)
|
||||
cancellation_token.cancel()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await coro
|
||||
|
||||
@ -1,18 +1,27 @@
|
||||
# File based from: https://github.com/microsoft/autogen/blob/main/test/coding/test_user_defined_functions.py
|
||||
# Credit to original authors
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
|
||||
import polars
|
||||
import pytest
|
||||
from azure.identity import DefaultAzureCredential
|
||||
from agnext.components.code_executor import (
|
||||
CodeBlock,
|
||||
FunctionWithRequirements,
|
||||
LocalCommandLineCodeExecutor,
|
||||
AzureContainerCodeExecutor,
|
||||
with_requirements,
|
||||
)
|
||||
from agnext.core import CancellationToken
|
||||
|
||||
ENVIRON_KEY_AZURE_POOL_ENDPOINT = "AZURE_POOL_ENDPOINT"
|
||||
|
||||
DUMMY_POOL_ENDPOINT = "DUMMY_POOL_ENDPOINT"
|
||||
POOL_ENDPOINT = os.getenv(ENVIRON_KEY_AZURE_POOL_ENDPOINT)
|
||||
|
||||
def add_two_numbers(a: int, b: int) -> int:
|
||||
"""Add two numbers together."""
|
||||
return a + b
|
||||
@ -70,6 +79,32 @@ print(data['name'][0])"""
|
||||
assert result.output == "John\n"
|
||||
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 = f"""import polars
|
||||
|
||||
# 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
|
||||
)
|
||||
assert azure_result.output == "John\n"
|
||||
assert azure_result.exit_code == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_can_load_function() -> None:
|
||||
@ -91,6 +126,30 @@ print(add_two_numbers(1, 2))"""
|
||||
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 = f"""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:
|
||||
@ -109,6 +168,26 @@ function_incorrect_import()"""
|
||||
cancellation_token=cancellation_token
|
||||
)
|
||||
|
||||
@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 = f"""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:
|
||||
@ -128,53 +207,84 @@ function_incorrect_dep()"""
|
||||
cancellation_token=cancellation_token
|
||||
)
|
||||
|
||||
@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 = f"""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 (
|
||||
'''def add_two_numbers(a: int, b: int) -> int:
|
||||
"""Add two numbers together."""
|
||||
'''
|
||||
in result
|
||||
)
|
||||
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:
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
func = FunctionWithRequirements.from_str(
|
||||
'''
|
||||
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 (
|
||||
'''def add_two_numbers(a: int, b: int) -> int:
|
||||
"""Add two numbers together."""
|
||||
'''
|
||||
in result
|
||||
)
|
||||
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:
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
cancellation_token = CancellationToken()
|
||||
func = FunctionWithRequirements.from_str(
|
||||
'''
|
||||
func = FunctionWithRequirements.from_str(
|
||||
'''
|
||||
def add_two_numbers(a: int, b: int) -> int:
|
||||
"""Add two numbers together."""
|
||||
return a + b
|
||||
'''
|
||||
)
|
||||
)
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
cancellation_token = CancellationToken()
|
||||
|
||||
executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, functions=[func])
|
||||
code = f"""from {executor.functions_module} import add_two_numbers
|
||||
@ -189,6 +299,34 @@ print(add_two_numbers(1, 2))"""
|
||||
assert result.output == "3\n"
|
||||
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 = f"""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:
|
||||
|
||||
@ -204,15 +342,15 @@ invaliddef add_two_numbers(a: int, b: int) -> int:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cant_run_broken_str_function_with_reqs() -> None:
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
cancellation_token = CancellationToken()
|
||||
func = FunctionWithRequirements.from_str(
|
||||
'''
|
||||
func = FunctionWithRequirements.from_str(
|
||||
'''
|
||||
def add_two_numbers(a: int, b: int) -> int:
|
||||
"""Add two numbers together."""
|
||||
return a + b
|
||||
'''
|
||||
)
|
||||
)
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
cancellation_token = CancellationToken()
|
||||
|
||||
executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, functions=[func])
|
||||
code = f"""from {executor.functions_module} import add_two_numbers
|
||||
@ -226,3 +364,35 @@ print(add_two_numbers(object(), False))"""
|
||||
)
|
||||
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 = f"""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')
|
||||
print(azure_result.output)
|
||||
assert "TypeError: unsupported operand type(s) for +:" in azure_result.output
|
||||
assert azure_result.exit_code == 1
|
||||
Loading…
x
Reference in New Issue
Block a user