Skip to content

Commit

Permalink
Merge branch 'main' into feat/add-remove-assistant-agent-tools
Browse files Browse the repository at this point in the history
  • Loading branch information
JMLX42 authored Dec 4, 2024
2 parents 8c25461 + e1d5833 commit af527fd
Show file tree
Hide file tree
Showing 4 changed files with 586 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -81,9 +81,23 @@ async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token:
if code_blocks:
# Execute the code blocks.
result = await self._code_executor.execute_code_blocks(code_blocks, cancellation_token=cancellation_token)
return Response(chat_message=TextMessage(content=result.output, source=self.name))

code_output = result.output
if code_output.strip() == "":
# No output
code_output = f"The script ran but produced no output to console. The POSIX exit code was: {result.exit_code}. If you were expecting output, consider revising the script to ensure content is printed to stdout."
elif result.exit_code != 0:
# Error
code_output = f"The script ran, then exited with an error (POSIX exit code: {result.exit_code})\nIts output was:\n{result.output}"

return Response(chat_message=TextMessage(content=code_output, source=self.name))
else:
return Response(chat_message=TextMessage(content="No code blocks found in the thread.", source=self.name))
return Response(
chat_message=TextMessage(
content="No code blocks found in the thread. Please provide at least one markdown-encoded code block to execute (i.e., quoting code in ```python or ```sh code blocks).",
source=self.name,
)
)

async def on_reset(self, cancellation_token: CancellationToken) -> None:
"""It it's a no-op as the code executor agent has no mutable state."""
Expand Down
166 changes: 166 additions & 0 deletions python/packages/autogen-agentchat/tests/test_code_executor_agent.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,166 @@
import pytest
from autogen_agentchat.agents import CodeExecutorAgent
from autogen_agentchat.base import Response
from autogen_agentchat.messages import TextMessage
from autogen_core import CancellationToken
from autogen_core.components.code_executor import LocalCommandLineCodeExecutor


@pytest.mark.asyncio
async def test_basic_code_execution() -> None:
"""Test basic code execution"""

agent = CodeExecutorAgent(name="code_executor", code_executor=LocalCommandLineCodeExecutor())

messages = [
TextMessage(
content="""
```python
import math
number = 42
square_root = math.sqrt(number)
print("%0.3f" % (square_root,))
```
""".strip(),
source="assistant",
)
]
response = await agent.on_messages(messages, CancellationToken())

assert isinstance(response, Response)
assert isinstance(response.chat_message, TextMessage)
assert response.chat_message.content.strip() == "6.481"
assert response.chat_message.source == "code_executor"


@pytest.mark.asyncio
async def test_code_execution_error() -> None:
"""Test basic code execution"""

agent = CodeExecutorAgent(name="code_executor", code_executor=LocalCommandLineCodeExecutor())

messages = [
TextMessage(
content="""
```python
import math
number = -1.0
square_root = math.sqrt(number)
print("%0.3f" % (square_root,))
```
""".strip(),
source="assistant",
)
]
response = await agent.on_messages(messages, CancellationToken())

assert isinstance(response, Response)
assert isinstance(response.chat_message, TextMessage)
assert "The script ran, then exited with an error (POSIX exit code: 1)" in response.chat_message.content
assert "ValueError: math domain error" in response.chat_message.content


@pytest.mark.asyncio
async def test_code_execution_no_output() -> None:
"""Test basic code execution"""

agent = CodeExecutorAgent(name="code_executor", code_executor=LocalCommandLineCodeExecutor())

messages = [
TextMessage(
content="""
```python
import math
number = 42
square_root = math.sqrt(number)
```
""".strip(),
source="assistant",
)
]
response = await agent.on_messages(messages, CancellationToken())

assert isinstance(response, Response)
assert isinstance(response.chat_message, TextMessage)
assert (
"The script ran but produced no output to console. The POSIX exit code was: 0. If you were expecting output, consider revising the script to ensure content is printed to stdout."
in response.chat_message.content
)


@pytest.mark.asyncio
async def test_code_execution_no_block() -> None:
"""Test basic code execution"""

agent = CodeExecutorAgent(name="code_executor", code_executor=LocalCommandLineCodeExecutor())

messages = [
TextMessage(
content="""
import math
number = 42
square_root = math.sqrt(number)
""".strip(),
source="assistant",
)
]
response = await agent.on_messages(messages, CancellationToken())

assert isinstance(response, Response)
assert isinstance(response.chat_message, TextMessage)
assert (
"No code blocks found in the thread. Please provide at least one markdown-encoded code block"
in response.chat_message.content
)


@pytest.mark.asyncio
async def test_code_execution_multiple_blocks() -> None:
"""Test basic code execution"""

agent = CodeExecutorAgent(name="code_executor", code_executor=LocalCommandLineCodeExecutor())

messages = [
TextMessage(
content="""
```python
import math
number = 42
square_root = math.sqrt(number)
print("%0.3f" % (square_root,))
```
And also:
```python
import time
print(f"The current time is: {time.time()}")
```
And this should result in an error:
```python
import math
number = -1.0
square_root = math.sqrt(number)
print("%0.3f" % (square_root,))
```
""".strip(),
source="assistant",
)
]
response = await agent.on_messages(messages, CancellationToken())

assert isinstance(response, Response)
assert isinstance(response.chat_message, TextMessage)
assert "6.481" in response.chat_message.content
assert "The current time is:" in response.chat_message.content
assert "The script ran, then exited with an error (POSIX exit code: 1)" in response.chat_message.content
assert "ValueError: math domain error" in response.chat_message.content
Loading

0 comments on commit af527fd

Please sign in to comment.