From 824003a8471f908bab3de4b274ade6076ca96d39 Mon Sep 17 00:00:00 2001 From: gziz Date: Tue, 19 Nov 2024 19:45:02 -0800 Subject: [PATCH 1/4] Console to return last processed (#4277) --- .../src/autogen_agentchat/task/_console.py | 23 +++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/task/_console.py b/python/packages/autogen-agentchat/src/autogen_agentchat/task/_console.py index 3c6464b1f07..5dbdd2a56e4 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/task/_console.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/task/_console.py @@ -22,19 +22,25 @@ async def Console( stream: AsyncGenerator[AgentMessage | TaskResult, None] | AsyncGenerator[AgentMessage | Response, None], *, no_inline_images: bool = False, -) -> None: - """Consume the stream from :meth:`~autogen_agentchat.base.Team.run_stream` +) -> TaskResult | Response: + """ + Consume the stream from :meth:`~autogen_agentchat.base.Team.run_stream` or :meth:`~autogen_agentchat.base.ChatAgent.on_messages_stream` - and print the messages to the console. + print the messages to the console and return the last processed TaskResult or Response. Args: stream (AsyncGenerator[AgentMessage | TaskResult, None] | AsyncGenerator[AgentMessage | Response, None]): Stream to render no_inline_images (bool, optional): If terminal is iTerm2 will render images inline. Use this to disable this behavior. Defaults to False. - """ + Returns: + last_processed: The last processed TaskResult or Response. + """ render_image_iterm = _is_running_in_iterm() and _is_output_a_tty() and not no_inline_images start_time = time.time() total_usage = RequestUsage(prompt_tokens=0, completion_tokens=0) + + last_processed: TaskResult | Response | None = None + async for message in stream: if isinstance(message, TaskResult): duration = time.time() - start_time @@ -47,6 +53,8 @@ async def Console( f"Duration: {duration:.2f} seconds\n" ) sys.stdout.write(output) + last_processed = message + elif isinstance(message, Response): duration = time.time() - start_time @@ -71,6 +79,8 @@ async def Console( f"Duration: {duration:.2f} seconds\n" ) sys.stdout.write(output) + last_processed = message + else: output = f"{'-' * 10} {message.source} {'-' * 10}\n{_message_to_str(message, render_image_iterm=render_image_iterm)}\n" if message.models_usage: @@ -79,6 +89,11 @@ async def Console( total_usage.prompt_tokens += message.models_usage.prompt_tokens sys.stdout.write(output) + if last_processed is None: + raise ValueError("No TaskResult or Response was processed.") + + return last_processed + # iTerm2 image rendering protocol: https://iterm2.com/documentation-images.html def _image_to_iterm(image: Image) -> str: From b632a9c166c0863b3542c7bc2a550c1ea46d1fea Mon Sep 17 00:00:00 2001 From: Jack Gerrits Date: Wed, 20 Nov 2024 11:24:04 -0500 Subject: [PATCH 2/4] Preserve input generator type --- .../src/autogen_agentchat/task/_console.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/task/_console.py b/python/packages/autogen-agentchat/src/autogen_agentchat/task/_console.py index 5dbdd2a56e4..6b5849e4cf4 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/task/_console.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/task/_console.py @@ -1,7 +1,7 @@ import os import sys import time -from typing import AsyncGenerator, List +from typing import AsyncGenerator, List, Optional, TypeVar, cast from autogen_core.components import Image from autogen_core.components.models import RequestUsage @@ -18,11 +18,14 @@ def _is_output_a_tty() -> bool: return sys.stdout.isatty() +T = TypeVar("T", bound=TaskResult | Response) + + async def Console( - stream: AsyncGenerator[AgentMessage | TaskResult, None] | AsyncGenerator[AgentMessage | Response, None], + stream: AsyncGenerator[AgentMessage | T, None], *, no_inline_images: bool = False, -) -> TaskResult | Response: +) -> T: """ Consume the stream from :meth:`~autogen_agentchat.base.Team.run_stream` or :meth:`~autogen_agentchat.base.ChatAgent.on_messages_stream` @@ -39,7 +42,7 @@ async def Console( start_time = time.time() total_usage = RequestUsage(prompt_tokens=0, completion_tokens=0) - last_processed: TaskResult | Response | None = None + last_processed: Optional[T] = None async for message in stream: if isinstance(message, TaskResult): @@ -53,7 +56,8 @@ async def Console( f"Duration: {duration:.2f} seconds\n" ) sys.stdout.write(output) - last_processed = message + # mypy ignore + last_processed = message # type: ignore elif isinstance(message, Response): duration = time.time() - start_time @@ -79,9 +83,12 @@ async def Console( f"Duration: {duration:.2f} seconds\n" ) sys.stdout.write(output) - last_processed = message + # mypy ignore + last_processed = message # type: ignore else: + # Cast required for mypy to be happy + message = cast(AgentMessage, message) # type: ignore output = f"{'-' * 10} {message.source} {'-' * 10}\n{_message_to_str(message, render_image_iterm=render_image_iterm)}\n" if message.models_usage: output += f"[Prompt tokens: {message.models_usage.prompt_tokens}, Completion tokens: {message.models_usage.completion_tokens}]\n" From 2fcb4d73936d2fd515d4b3a41ff5646cf2527e2d Mon Sep 17 00:00:00 2001 From: Eric Zhu Date: Wed, 20 Nov 2024 11:51:37 -0500 Subject: [PATCH 3/4] Add tests --- .../tests/test_group_chat.py | 41 ++++++++++++++++++- 1 file changed, 40 insertions(+), 1 deletion(-) diff --git a/python/packages/autogen-agentchat/tests/test_group_chat.py b/python/packages/autogen-agentchat/tests/test_group_chat.py index b6c4f4bfdc7..829f9960788 100644 --- a/python/packages/autogen-agentchat/tests/test_group_chat.py +++ b/python/packages/autogen-agentchat/tests/test_group_chat.py @@ -24,7 +24,7 @@ ToolCallMessage, ToolCallResultMessage, ) -from autogen_agentchat.task import HandoffTermination, MaxMessageTermination, TextMentionTermination +from autogen_agentchat.task import HandoffTermination, MaxMessageTermination, TextMentionTermination, Console from autogen_agentchat.teams import ( RoundRobinGroupChat, SelectorGroupChat, @@ -314,6 +314,14 @@ async def test_round_robin_group_chat_with_tools(monkeypatch: pytest.MonkeyPatch else: assert message == result.messages[index] index += 1 + + # Test Console. + tool_use_agent._model_context.clear() # pyright: ignore + mock.reset() + index = 0 + await team.reset() + result2 = await Console(team.run_stream(task="Write a program that prints 'Hello, world!'")) + assert result2 == result @pytest.mark.asyncio @@ -475,6 +483,14 @@ async def test_selector_group_chat(monkeypatch: pytest.MonkeyPatch) -> None: else: assert message == result.messages[index] index += 1 + + # Test Console. + mock.reset() + agent1._count = 0 # pyright: ignore + index = 0 + await team.reset() + result2 = await Console(team.run_stream(task="Write a program that prints 'Hello, world!'")) + assert result2 == result @pytest.mark.asyncio @@ -527,6 +543,14 @@ async def test_selector_group_chat_two_speakers(monkeypatch: pytest.MonkeyPatch) else: assert message == result.messages[index] index += 1 + + # Test Console. + mock.reset() + agent1._count = 0 # pyright: ignore + index = 0 + await team.reset() + result2 = await Console(team.run_stream(task="Write a program that prints 'Hello, world!'")) + assert result2 == result @pytest.mark.asyncio @@ -594,6 +618,13 @@ async def test_selector_group_chat_two_speakers_allow_repeated(monkeypatch: pyte else: assert message == result.messages[index] index += 1 + + # Test Console. + mock.reset() + index = 0 + await team.reset() + result2 = await Console(team.run_stream(task="Write a program that prints 'Hello, world!'")) + assert result2 == result @pytest.mark.asyncio @@ -791,6 +822,14 @@ async def test_swarm_handoff_using_tool_calls(monkeypatch: pytest.MonkeyPatch) - else: assert message == result.messages[index] index += 1 + + # Test Console + agent1._model_context.clear() # pyright: ignore + mock.reset() + index = 0 + await team.reset() + result2 = await Console(team.run_stream(task="task")) + assert result2 == result @pytest.mark.asyncio From 785964202f8ed0b3363bbb584ab13abbfa8becd5 Mon Sep 17 00:00:00 2001 From: Eric Zhu Date: Wed, 20 Nov 2024 11:55:31 -0500 Subject: [PATCH 4/4] format --- .../autogen-agentchat/tests/test_group_chat.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/python/packages/autogen-agentchat/tests/test_group_chat.py b/python/packages/autogen-agentchat/tests/test_group_chat.py index 829f9960788..7df27abcbcd 100644 --- a/python/packages/autogen-agentchat/tests/test_group_chat.py +++ b/python/packages/autogen-agentchat/tests/test_group_chat.py @@ -24,7 +24,7 @@ ToolCallMessage, ToolCallResultMessage, ) -from autogen_agentchat.task import HandoffTermination, MaxMessageTermination, TextMentionTermination, Console +from autogen_agentchat.task import Console, HandoffTermination, MaxMessageTermination, TextMentionTermination from autogen_agentchat.teams import ( RoundRobinGroupChat, SelectorGroupChat, @@ -314,7 +314,7 @@ async def test_round_robin_group_chat_with_tools(monkeypatch: pytest.MonkeyPatch else: assert message == result.messages[index] index += 1 - + # Test Console. tool_use_agent._model_context.clear() # pyright: ignore mock.reset() @@ -483,7 +483,7 @@ async def test_selector_group_chat(monkeypatch: pytest.MonkeyPatch) -> None: else: assert message == result.messages[index] index += 1 - + # Test Console. mock.reset() agent1._count = 0 # pyright: ignore @@ -543,7 +543,7 @@ async def test_selector_group_chat_two_speakers(monkeypatch: pytest.MonkeyPatch) else: assert message == result.messages[index] index += 1 - + # Test Console. mock.reset() agent1._count = 0 # pyright: ignore @@ -618,7 +618,7 @@ async def test_selector_group_chat_two_speakers_allow_repeated(monkeypatch: pyte else: assert message == result.messages[index] index += 1 - + # Test Console. mock.reset() index = 0 @@ -822,7 +822,7 @@ async def test_swarm_handoff_using_tool_calls(monkeypatch: pytest.MonkeyPatch) - else: assert message == result.messages[index] index += 1 - + # Test Console agent1._model_context.clear() # pyright: ignore mock.reset()