From f9ac86e86ca73b4c27135c30f165f865baaf152f Mon Sep 17 00:00:00 2001 From: gagb Date: Wed, 25 Dec 2024 15:58:08 -0800 Subject: [PATCH 1/8] Implement custom console for MagenticOne CLI with enhanced message rendering --- .../src/autogen_ext/teams/magentic_one_cli.py | 131 +++++++++++++++++- 1 file changed, 128 insertions(+), 3 deletions(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py index 97abac5933b7..2b0fec244d35 100644 --- a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py +++ b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py @@ -1,12 +1,137 @@ import argparse import asyncio +import os +import sys +import time +from typing import Any, AsyncGenerator, List, Optional, TypeVar, cast -from autogen_agentchat.ui import Console +from autogen_agentchat.base import Response, TaskResult +from autogen_agentchat.messages import AgentEvent, ChatMessage, MultiModalMessage +from autogen_core import Image +from autogen_core.models import RequestUsage from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.teams.magentic_one import MagenticOne +def _is_running_in_iterm() -> bool: + return os.getenv("TERM_PROGRAM") == "iTerm.app" + + +def _is_output_a_tty() -> bool: + return sys.stdout.isatty() + + +T = TypeVar("T", bound=TaskResult | Response) + + +# iTerm2 image rendering protocol: https://iterm2.com/documentation-images.html +def _image_to_iterm(image: Image) -> str: + image_data = image.to_base64() + return f"\033]1337;File=inline=1:{image_data}\a\n" + + +def _message_to_str(message: AgentEvent | ChatMessage, *, render_image_iterm: bool = False) -> str: + if isinstance(message, MultiModalMessage): + result: List[str] = [] + for c in message.content: + if isinstance(c, str): + result.append(c) + else: + if render_image_iterm: + result.append(_image_to_iterm(c)) + else: + result.append("") + return "\n".join(result) + else: + return f"{message.content}" + + +def _print(*args: Any, **kwargs: Any) -> None: + print(*args, **kwargs) # noqa: T201 + + +async def CustomConsole( + stream: AsyncGenerator[AgentEvent | ChatMessage | T, None], + *, + no_inline_images: bool = False, +) -> T: + """ + Consumes the message stream from :meth:`~autogen_agentchat.base.TaskRunner.run_stream` + or :meth:`~autogen_agentchat.base.ChatAgent.on_messages_stream` and renders the messages to the console using print(). + Returns the last processed TaskResult or Response. + + Args: + stream (AsyncGenerator[AgentEvent | ChatMessage | TaskResult, None] | AsyncGenerator[AgentEvent | ChatMessage | Response, None]): Message stream to render. + This can be from :meth:`~autogen_agentchat.base.TaskRunner.run_stream` or :meth:`~autogen_agentchat.base.ChatAgent.on_messages_stream`. + 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: A :class:`~autogen_agentchat.base.TaskResult` if the stream is from :meth:`~autogen_agentchat.base.TaskRunner.run_stream` + or a :class:`~autogen_agentchat.base.Response` if the stream is from :meth:`~autogen_agentchat.base.ChatAgent.on_messages_stream`. + """ + 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: Optional[T] = None + + async for message in stream: + if isinstance(message, TaskResult): + duration = time.time() - start_time + output = ( + f"{'-' * 10} Summary {'-' * 10}\n" + f"Number of messages: {len(message.messages)}\n" + f"Finish reason: {message.stop_reason}\n" + f"Total prompt tokens: {total_usage.prompt_tokens}\n" + f"Total completion tokens: {total_usage.completion_tokens}\n" + f"Duration: {duration:.2f} seconds\n" + ) + _print(output) + last_processed = message # type: ignore + + elif isinstance(message, Response): + duration = time.time() - start_time + + # Print final response. + output = f"{'-' * 10} {message.chat_message.source} {'-' * 10}\n{_message_to_str(message.chat_message, render_image_iterm=render_image_iterm)}\n" + if message.chat_message.models_usage: + output += f"[Prompt tokens: {message.chat_message.models_usage.prompt_tokens}, Completion tokens: {message.chat_message.models_usage.completion_tokens}]\n" + total_usage.completion_tokens += message.chat_message.models_usage.completion_tokens + total_usage.prompt_tokens += message.chat_message.models_usage.prompt_tokens + _print(output) + + # Print summary. + if message.inner_messages is not None: + num_inner_messages = len(message.inner_messages) + else: + num_inner_messages = 0 + output = ( + f"{'-' * 10} Summary {'-' * 10}\n" + f"Number of inner messages: {num_inner_messages}\n" + f"Total prompt tokens: {total_usage.prompt_tokens}\n" + f"Total completion tokens: {total_usage.completion_tokens}\n" + f"Duration: {duration:.2f} seconds\n" + ) + _print(output) + last_processed = message # type: ignore + + else: + # Cast required for mypy to be happy + message = cast(AgentEvent | ChatMessage, 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" + total_usage.completion_tokens += message.models_usage.completion_tokens + total_usage.prompt_tokens += message.models_usage.prompt_tokens + _print(output) + + if last_processed is None: + raise ValueError("No TaskResult or Response was processed.") + + return last_processed + + def main() -> None: """ Command-line interface for running a complex task using MagenticOne. @@ -36,10 +161,10 @@ def main() -> None: async def run_task(task: str, hil_mode: bool) -> None: client = OpenAIChatCompletionClient(model="gpt-4o") m1 = MagenticOne(client=client, hil_mode=hil_mode) - await Console(m1.run_stream(task=task)) + await CustomConsole(m1.run_stream(task=task)) task = args.task[0] - asyncio.run(run_task(task, not args.no_hil)) + asyncio.run(asyncio.wait_for(run_task(task, not args.no_hil), timeout=300)) if __name__ == "__main__": From 51489661b78641cfc532490d0c553bc1f04d48e6 Mon Sep 17 00:00:00 2001 From: gagb Date: Wed, 25 Dec 2024 16:09:33 -0800 Subject: [PATCH 2/8] Enhance CustomConsole output with Rich library for improved formatting --- .../src/autogen_ext/teams/magentic_one_cli.py | 30 ++++++++++++------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py index 2b0fec244d35..6dfbea01f905 100644 --- a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py +++ b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py @@ -12,6 +12,10 @@ from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.teams.magentic_one import MagenticOne +from rich.console import Console +from rich.panel import Panel +from rich.text import Text + def _is_running_in_iterm() -> bool: @@ -75,31 +79,31 @@ async def CustomConsole( total_usage = RequestUsage(prompt_tokens=0, completion_tokens=0) last_processed: Optional[T] = None + console = Console() async for message in stream: if isinstance(message, TaskResult): duration = time.time() - start_time output = ( - f"{'-' * 10} Summary {'-' * 10}\n" f"Number of messages: {len(message.messages)}\n" f"Finish reason: {message.stop_reason}\n" f"Total prompt tokens: {total_usage.prompt_tokens}\n" f"Total completion tokens: {total_usage.completion_tokens}\n" f"Duration: {duration:.2f} seconds\n" ) - _print(output) + console.print(Panel(output, title="Summary")) last_processed = message # type: ignore elif isinstance(message, Response): duration = time.time() - start_time # Print final response. - output = f"{'-' * 10} {message.chat_message.source} {'-' * 10}\n{_message_to_str(message.chat_message, render_image_iterm=render_image_iterm)}\n" + output = Text.from_markup(f"[bold]{message.chat_message.source}[/bold]\n{_message_to_str(message.chat_message, render_image_iterm=render_image_iterm)}") if message.chat_message.models_usage: - output += f"[Prompt tokens: {message.chat_message.models_usage.prompt_tokens}, Completion tokens: {message.chat_message.models_usage.completion_tokens}]\n" + output.append(f"\n[Prompt tokens: {message.chat_message.models_usage.prompt_tokens}, Completion tokens: {message.chat_message.models_usage.completion_tokens}]") total_usage.completion_tokens += message.chat_message.models_usage.completion_tokens total_usage.prompt_tokens += message.chat_message.models_usage.prompt_tokens - _print(output) + console.print(Panel(output)) # Print summary. if message.inner_messages is not None: @@ -107,24 +111,23 @@ async def CustomConsole( else: num_inner_messages = 0 output = ( - f"{'-' * 10} Summary {'-' * 10}\n" f"Number of inner messages: {num_inner_messages}\n" f"Total prompt tokens: {total_usage.prompt_tokens}\n" f"Total completion tokens: {total_usage.completion_tokens}\n" f"Duration: {duration:.2f} seconds\n" ) - _print(output) + console.print(Panel(output, title="Summary")) last_processed = message # type: ignore else: # Cast required for mypy to be happy message = cast(AgentEvent | ChatMessage, message) # type: ignore - output = f"{'-' * 10} {message.source} {'-' * 10}\n{_message_to_str(message, render_image_iterm=render_image_iterm)}\n" + output = Text.from_markup(f"[bold]{message.source}[/bold]\n{_message_to_str(message, render_image_iterm=render_image_iterm)}") if message.models_usage: - output += f"[Prompt tokens: {message.models_usage.prompt_tokens}, Completion tokens: {message.models_usage.completion_tokens}]\n" + output.append(f"\n[Prompt tokens: {message.models_usage.prompt_tokens}, Completion tokens: {message.models_usage.completion_tokens}]") total_usage.completion_tokens += message.models_usage.completion_tokens total_usage.prompt_tokens += message.models_usage.prompt_tokens - _print(output) + console.print(Panel(output)) if last_processed is None: raise ValueError("No TaskResult or Response was processed.") @@ -168,4 +171,11 @@ async def run_task(task: str, hil_mode: bool) -> None: if __name__ == "__main__": + import os + import sys + + fd = sys.stdout.fileno() + flags = os.fcntl(fd, os.F_GETFL) + os.fcntl(fd, os.F_SETFL, flags & ~os.O_NONBLOCK) + main() From 6600408b15df7c06533bb5cd14e7bc0759c327fb Mon Sep 17 00:00:00 2001 From: gagb Date: Wed, 25 Dec 2024 16:27:17 -0800 Subject: [PATCH 3/8] Refactor CustomConsole to use AutoGenConsole for MagenticOne CLI --- .../src/autogen_ext/teams/magentic_one_cli.py | 130 +----------------- 1 file changed, 2 insertions(+), 128 deletions(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py index 6dfbea01f905..a0c35c8bf9cd 100644 --- a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py +++ b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py @@ -2,137 +2,11 @@ import asyncio import os import sys -import time -from typing import Any, AsyncGenerator, List, Optional, TypeVar, cast -from autogen_agentchat.base import Response, TaskResult -from autogen_agentchat.messages import AgentEvent, ChatMessage, MultiModalMessage -from autogen_core import Image -from autogen_core.models import RequestUsage +from autogen_agentchat.ui import Console as AutoGenConsole from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.teams.magentic_one import MagenticOne -from rich.console import Console -from rich.panel import Panel -from rich.text import Text - - - -def _is_running_in_iterm() -> bool: - return os.getenv("TERM_PROGRAM") == "iTerm.app" - - -def _is_output_a_tty() -> bool: - return sys.stdout.isatty() - - -T = TypeVar("T", bound=TaskResult | Response) - - -# iTerm2 image rendering protocol: https://iterm2.com/documentation-images.html -def _image_to_iterm(image: Image) -> str: - image_data = image.to_base64() - return f"\033]1337;File=inline=1:{image_data}\a\n" - - -def _message_to_str(message: AgentEvent | ChatMessage, *, render_image_iterm: bool = False) -> str: - if isinstance(message, MultiModalMessage): - result: List[str] = [] - for c in message.content: - if isinstance(c, str): - result.append(c) - else: - if render_image_iterm: - result.append(_image_to_iterm(c)) - else: - result.append("") - return "\n".join(result) - else: - return f"{message.content}" - - -def _print(*args: Any, **kwargs: Any) -> None: - print(*args, **kwargs) # noqa: T201 - - -async def CustomConsole( - stream: AsyncGenerator[AgentEvent | ChatMessage | T, None], - *, - no_inline_images: bool = False, -) -> T: - """ - Consumes the message stream from :meth:`~autogen_agentchat.base.TaskRunner.run_stream` - or :meth:`~autogen_agentchat.base.ChatAgent.on_messages_stream` and renders the messages to the console using print(). - Returns the last processed TaskResult or Response. - - Args: - stream (AsyncGenerator[AgentEvent | ChatMessage | TaskResult, None] | AsyncGenerator[AgentEvent | ChatMessage | Response, None]): Message stream to render. - This can be from :meth:`~autogen_agentchat.base.TaskRunner.run_stream` or :meth:`~autogen_agentchat.base.ChatAgent.on_messages_stream`. - 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: A :class:`~autogen_agentchat.base.TaskResult` if the stream is from :meth:`~autogen_agentchat.base.TaskRunner.run_stream` - or a :class:`~autogen_agentchat.base.Response` if the stream is from :meth:`~autogen_agentchat.base.ChatAgent.on_messages_stream`. - """ - 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: Optional[T] = None - console = Console() - - async for message in stream: - if isinstance(message, TaskResult): - duration = time.time() - start_time - output = ( - f"Number of messages: {len(message.messages)}\n" - f"Finish reason: {message.stop_reason}\n" - f"Total prompt tokens: {total_usage.prompt_tokens}\n" - f"Total completion tokens: {total_usage.completion_tokens}\n" - f"Duration: {duration:.2f} seconds\n" - ) - console.print(Panel(output, title="Summary")) - last_processed = message # type: ignore - - elif isinstance(message, Response): - duration = time.time() - start_time - - # Print final response. - output = Text.from_markup(f"[bold]{message.chat_message.source}[/bold]\n{_message_to_str(message.chat_message, render_image_iterm=render_image_iterm)}") - if message.chat_message.models_usage: - output.append(f"\n[Prompt tokens: {message.chat_message.models_usage.prompt_tokens}, Completion tokens: {message.chat_message.models_usage.completion_tokens}]") - total_usage.completion_tokens += message.chat_message.models_usage.completion_tokens - total_usage.prompt_tokens += message.chat_message.models_usage.prompt_tokens - console.print(Panel(output)) - - # Print summary. - if message.inner_messages is not None: - num_inner_messages = len(message.inner_messages) - else: - num_inner_messages = 0 - output = ( - f"Number of inner messages: {num_inner_messages}\n" - f"Total prompt tokens: {total_usage.prompt_tokens}\n" - f"Total completion tokens: {total_usage.completion_tokens}\n" - f"Duration: {duration:.2f} seconds\n" - ) - console.print(Panel(output, title="Summary")) - last_processed = message # type: ignore - - else: - # Cast required for mypy to be happy - message = cast(AgentEvent | ChatMessage, message) # type: ignore - output = Text.from_markup(f"[bold]{message.source}[/bold]\n{_message_to_str(message, render_image_iterm=render_image_iterm)}") - if message.models_usage: - output.append(f"\n[Prompt tokens: {message.models_usage.prompt_tokens}, Completion tokens: {message.models_usage.completion_tokens}]") - total_usage.completion_tokens += message.models_usage.completion_tokens - total_usage.prompt_tokens += message.models_usage.prompt_tokens - console.print(Panel(output)) - - if last_processed is None: - raise ValueError("No TaskResult or Response was processed.") - - return last_processed def main() -> None: @@ -164,7 +38,7 @@ def main() -> None: async def run_task(task: str, hil_mode: bool) -> None: client = OpenAIChatCompletionClient(model="gpt-4o") m1 = MagenticOne(client=client, hil_mode=hil_mode) - await CustomConsole(m1.run_stream(task=task)) + await AutoGenConsole(m1.run_stream(task=task)) task = args.task[0] asyncio.run(asyncio.wait_for(run_task(task, not args.no_hil), timeout=300)) From 59f00e674519a39a75c75cd8a4f96ae44fd4d357 Mon Sep 17 00:00:00 2001 From: gagb Date: Wed, 25 Dec 2024 16:27:31 -0800 Subject: [PATCH 4/8] Remove unused imports from MagenticOne CLI main execution block --- .../autogen-ext/src/autogen_ext/teams/magentic_one_cli.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py index a0c35c8bf9cd..ef17ac8b5e90 100644 --- a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py +++ b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py @@ -45,9 +45,6 @@ async def run_task(task: str, hil_mode: bool) -> None: if __name__ == "__main__": - import os - import sys - fd = sys.stdout.fileno() flags = os.fcntl(fd, os.F_GETFL) os.fcntl(fd, os.F_SETFL, flags & ~os.O_NONBLOCK) From b3db06be6af9408de813f755e6c7606694535f03 Mon Sep 17 00:00:00 2001 From: gagb Date: Wed, 25 Dec 2024 16:41:00 -0800 Subject: [PATCH 5/8] Replace AutoGenConsole with RichConsole for enhanced output formatting in MagenticOne CLI --- .../src/autogen_ext/teams/magentic_one_cli.py | 6 +- .../src/autogen_ext/teams/rich_console.py | 105 ++++++++++++++++++ 2 files changed, 107 insertions(+), 4 deletions(-) create mode 100644 python/packages/autogen-ext/src/autogen_ext/teams/rich_console.py diff --git a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py index ef17ac8b5e90..486e11fc156f 100644 --- a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py +++ b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py @@ -2,11 +2,9 @@ import asyncio import os import sys - -from autogen_agentchat.ui import Console as AutoGenConsole - from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.teams.magentic_one import MagenticOne +from autogen_ext.teams.rich_console import RichConsole def main() -> None: @@ -38,7 +36,7 @@ def main() -> None: async def run_task(task: str, hil_mode: bool) -> None: client = OpenAIChatCompletionClient(model="gpt-4o") m1 = MagenticOne(client=client, hil_mode=hil_mode) - await AutoGenConsole(m1.run_stream(task=task)) + await RichConsole(m1.run_stream(task=task)) task = args.task[0] asyncio.run(asyncio.wait_for(run_task(task, not args.no_hil), timeout=300)) diff --git a/python/packages/autogen-ext/src/autogen_ext/teams/rich_console.py b/python/packages/autogen-ext/src/autogen_ext/teams/rich_console.py new file mode 100644 index 000000000000..ede61667e2e7 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/teams/rich_console.py @@ -0,0 +1,105 @@ +import os +import sys +import time +from typing import AsyncGenerator, List, Optional, TypeVar, cast + +from autogen_agentchat.base import Response, TaskResult +from autogen_agentchat.messages import AgentEvent, ChatMessage, MultiModalMessage +from autogen_core import Image +from autogen_core.models import RequestUsage +from rich.console import Console +from rich.panel import Panel +from rich.text import Text + +T = TypeVar("T", bound=TaskResult | Response) + +def _is_running_in_iterm() -> bool: + return os.getenv("TERM_PROGRAM") == "iTerm.app" + +def _is_output_a_tty() -> bool: + return sys.stdout.isatty() + +def _image_to_iterm(image: Image) -> str: + image_data = image.to_base64() + return f"\033]1337;File=inline=1:{image_data}\a\n" + +def _message_to_str(message: AgentEvent | ChatMessage, *, render_image_iterm: bool = False) -> str: + if isinstance(message, MultiModalMessage): + result: List[str] = [] + for c in message.content: + if isinstance(c, str): + result.append(c) + else: + if render_image_iterm: + result.append(_image_to_iterm(c)) + else: + result.append("") + return "\n".join(result) + else: + return f"{message.content}" + +async def RichConsole( + stream: AsyncGenerator[AgentEvent | ChatMessage | T, None], + *, + no_inline_images: bool = False, +) -> T: + 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: Optional[T] = None + console = Console() + + async for message in stream: + if isinstance(message, TaskResult): + duration = time.time() - start_time + output = ( + f"Number of messages: {len(message.messages)}\n" + f"Finish reason: {message.stop_reason}\n" + f"Total prompt tokens: {total_usage.prompt_tokens}\n" + f"Total completion tokens: {total_usage.completion_tokens}\n" + f"Duration: {duration:.2f} seconds\n" + ) + console.print(Panel(output, title="Summary")) + last_processed = message # type: ignore + + elif isinstance(message, Response): + duration = time.time() - start_time + + output = Text.from_markup(f"[bold magenta]{message.chat_message.source}[/bold magenta]\n{_message_to_str(message.chat_message, render_image_iterm=render_image_iterm)}") + if message.chat_message.models_usage: + output.append(f"\n[Prompt tokens: {message.chat_message.models_usage.prompt_tokens}, Completion tokens: {message.chat_message.models_usage.completion_tokens}]") + total_usage.completion_tokens += message.chat_message.models_usage.completion_tokens + total_usage.prompt_tokens += message.chat_message.models_usage.prompt_tokens + console.print(Panel(output)) + + if message.inner_messages is not None: + num_inner_messages = len(message.inner_messages) + else: + num_inner_messages = 0 + output = ( + f"Number of inner messages: {num_inner_messages}\n" + f"Total prompt tokens: {total_usage.prompt_tokens}\n" + f"Total completion tokens: {total_usage.completion_tokens}\n" + f"Duration: {duration:.2f} seconds\n" + ) + console.print(Panel(output, title="Summary")) + last_processed = message # type: ignore + + else: + message = cast(AgentEvent | ChatMessage, message) # type: ignore + output = Text.from_markup(f"[bold magenta]{message.source}[/bold magenta]\n{_message_to_str(message, render_image_iterm=render_image_iterm)}") + if message.models_usage: + output.append(f"\n[Prompt tokens: {message.models_usage.prompt_tokens}, Completion tokens: {message.models_usage.completion_tokens}]") + total_usage.completion_tokens += message.models_usage.completion_tokens + total_usage.prompt_tokens += message.models_usage.prompt_tokens + console.print(Panel(output)) + if render_image_iterm and isinstance(message, MultiModalMessage): + for c in message.content: + if isinstance(c, Image): + console.print(_image_to_iterm(c)) + + if last_processed is None: + raise ValueError("No TaskResult or Response was processed.") + + return last_processed From 5cec7bb10a6a0b78fdffb514ce22f13a889654fe Mon Sep 17 00:00:00 2001 From: gagb Date: Wed, 25 Dec 2024 16:57:24 -0800 Subject: [PATCH 6/8] Add primary color option to RichConsole and adjust output formatting --- .../src/autogen_ext/teams/magentic_one_cli.py | 2 +- .../autogen-ext/src/autogen_ext/teams/rich_console.py | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py index 486e11fc156f..871c352b7729 100644 --- a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py +++ b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py @@ -36,7 +36,7 @@ def main() -> None: async def run_task(task: str, hil_mode: bool) -> None: client = OpenAIChatCompletionClient(model="gpt-4o") m1 = MagenticOne(client=client, hil_mode=hil_mode) - await RichConsole(m1.run_stream(task=task)) + await RichConsole(m1.run_stream(task=task), no_inline_images=True) task = args.task[0] asyncio.run(asyncio.wait_for(run_task(task, not args.no_hil), timeout=300)) diff --git a/python/packages/autogen-ext/src/autogen_ext/teams/rich_console.py b/python/packages/autogen-ext/src/autogen_ext/teams/rich_console.py index ede61667e2e7..2a610b71df02 100644 --- a/python/packages/autogen-ext/src/autogen_ext/teams/rich_console.py +++ b/python/packages/autogen-ext/src/autogen_ext/teams/rich_console.py @@ -42,6 +42,7 @@ async def RichConsole( stream: AsyncGenerator[AgentEvent | ChatMessage | T, None], *, no_inline_images: bool = False, + primary_color: str = "magenta", ) -> T: render_image_iterm = _is_running_in_iterm() and _is_output_a_tty() and not no_inline_images start_time = time.time() @@ -66,12 +67,12 @@ async def RichConsole( elif isinstance(message, Response): duration = time.time() - start_time - output = Text.from_markup(f"[bold magenta]{message.chat_message.source}[/bold magenta]\n{_message_to_str(message.chat_message, render_image_iterm=render_image_iterm)}") + output = Text.from_markup(f"{_message_to_str(message.chat_message, render_image_iterm=render_image_iterm)}") if message.chat_message.models_usage: output.append(f"\n[Prompt tokens: {message.chat_message.models_usage.prompt_tokens}, Completion tokens: {message.chat_message.models_usage.completion_tokens}]") total_usage.completion_tokens += message.chat_message.models_usage.completion_tokens total_usage.prompt_tokens += message.chat_message.models_usage.prompt_tokens - console.print(Panel(output)) + console.print(Panel(output, title=f"[bold {primary_color}]{message.chat_message.source}[/bold {primary_color}]")) if message.inner_messages is not None: num_inner_messages = len(message.inner_messages) @@ -88,12 +89,12 @@ async def RichConsole( else: message = cast(AgentEvent | ChatMessage, message) # type: ignore - output = Text.from_markup(f"[bold magenta]{message.source}[/bold magenta]\n{_message_to_str(message, render_image_iterm=render_image_iterm)}") + output = Text.from_markup(f"{_message_to_str(message, render_image_iterm=render_image_iterm)}") if message.models_usage: output.append(f"\n[Prompt tokens: {message.models_usage.prompt_tokens}, Completion tokens: {message.models_usage.completion_tokens}]") total_usage.completion_tokens += message.models_usage.completion_tokens total_usage.prompt_tokens += message.models_usage.prompt_tokens - console.print(Panel(output)) + console.print(Panel(output, title=f"[bold {primary_color}]{message.source}[/bold {primary_color}]")) if render_image_iterm and isinstance(message, MultiModalMessage): for c in message.content: if isinstance(c, Image): From 7284ae80a73c8e8a9f4cc4a3a641312e2767833e Mon Sep 17 00:00:00 2001 From: gagb Date: Wed, 25 Dec 2024 17:01:01 -0800 Subject: [PATCH 7/8] Refactor RichConsole to simplify image handling and improve output formatting --- .../src/autogen_ext/teams/magentic_one_cli.py | 2 +- .../src/autogen_ext/teams/rich_console.py | 19 +++++++++++++++---- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py index 871c352b7729..486e11fc156f 100644 --- a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py +++ b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one_cli.py @@ -36,7 +36,7 @@ def main() -> None: async def run_task(task: str, hil_mode: bool) -> None: client = OpenAIChatCompletionClient(model="gpt-4o") m1 = MagenticOne(client=client, hil_mode=hil_mode) - await RichConsole(m1.run_stream(task=task), no_inline_images=True) + await RichConsole(m1.run_stream(task=task)) task = args.task[0] asyncio.run(asyncio.wait_for(run_task(task, not args.no_hil), timeout=300)) diff --git a/python/packages/autogen-ext/src/autogen_ext/teams/rich_console.py b/python/packages/autogen-ext/src/autogen_ext/teams/rich_console.py index 2a610b71df02..e7583fff5b3f 100644 --- a/python/packages/autogen-ext/src/autogen_ext/teams/rich_console.py +++ b/python/packages/autogen-ext/src/autogen_ext/teams/rich_console.py @@ -13,16 +13,20 @@ T = TypeVar("T", bound=TaskResult | Response) + def _is_running_in_iterm() -> bool: return os.getenv("TERM_PROGRAM") == "iTerm.app" + def _is_output_a_tty() -> bool: return sys.stdout.isatty() + def _image_to_iterm(image: Image) -> str: image_data = image.to_base64() return f"\033]1337;File=inline=1:{image_data}\a\n" + def _message_to_str(message: AgentEvent | ChatMessage, *, render_image_iterm: bool = False) -> str: if isinstance(message, MultiModalMessage): result: List[str] = [] @@ -38,6 +42,7 @@ def _message_to_str(message: AgentEvent | ChatMessage, *, render_image_iterm: bo else: return f"{message.content}" + async def RichConsole( stream: AsyncGenerator[AgentEvent | ChatMessage | T, None], *, @@ -69,10 +74,14 @@ async def RichConsole( output = Text.from_markup(f"{_message_to_str(message.chat_message, render_image_iterm=render_image_iterm)}") if message.chat_message.models_usage: - output.append(f"\n[Prompt tokens: {message.chat_message.models_usage.prompt_tokens}, Completion tokens: {message.chat_message.models_usage.completion_tokens}]") + output.append( + f"\n[Prompt tokens: {message.chat_message.models_usage.prompt_tokens}, Completion tokens: {message.chat_message.models_usage.completion_tokens}]" + ) total_usage.completion_tokens += message.chat_message.models_usage.completion_tokens total_usage.prompt_tokens += message.chat_message.models_usage.prompt_tokens - console.print(Panel(output, title=f"[bold {primary_color}]{message.chat_message.source}[/bold {primary_color}]")) + console.print( + Panel(output, title=f"[bold {primary_color}]{message.chat_message.source}[/bold {primary_color}]") + ) if message.inner_messages is not None: num_inner_messages = len(message.inner_messages) @@ -91,14 +100,16 @@ async def RichConsole( message = cast(AgentEvent | ChatMessage, message) # type: ignore output = Text.from_markup(f"{_message_to_str(message, render_image_iterm=render_image_iterm)}") if message.models_usage: - output.append(f"\n[Prompt tokens: {message.models_usage.prompt_tokens}, Completion tokens: {message.models_usage.completion_tokens}]") + output.append( + f"\n[Prompt tokens: {message.models_usage.prompt_tokens}, Completion tokens: {message.models_usage.completion_tokens}]" + ) total_usage.completion_tokens += message.models_usage.completion_tokens total_usage.prompt_tokens += message.models_usage.prompt_tokens console.print(Panel(output, title=f"[bold {primary_color}]{message.source}[/bold {primary_color}]")) if render_image_iterm and isinstance(message, MultiModalMessage): for c in message.content: if isinstance(c, Image): - console.print(_image_to_iterm(c)) + print(_image_to_iterm(c)) if last_processed is None: raise ValueError("No TaskResult or Response was processed.") From 43a3808332a81b2c15a859cab404fb0de9846504 Mon Sep 17 00:00:00 2001 From: gagb Date: Wed, 25 Dec 2024 17:51:11 -0800 Subject: [PATCH 8/8] Refactor RichConsole to streamline image handling and remove redundant output logic --- .../src/autogen_ext/teams/rich_console.py | 33 +------------------ 1 file changed, 1 insertion(+), 32 deletions(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/teams/rich_console.py b/python/packages/autogen-ext/src/autogen_ext/teams/rich_console.py index e7583fff5b3f..d366e4cff8fd 100644 --- a/python/packages/autogen-ext/src/autogen_ext/teams/rich_console.py +++ b/python/packages/autogen-ext/src/autogen_ext/teams/rich_console.py @@ -34,10 +34,7 @@ def _message_to_str(message: AgentEvent | ChatMessage, *, render_image_iterm: bo if isinstance(c, str): result.append(c) else: - if render_image_iterm: - result.append(_image_to_iterm(c)) - else: - result.append("") + result.append("") return "\n".join(result) else: return f"{message.content}" @@ -68,34 +65,6 @@ async def RichConsole( ) console.print(Panel(output, title="Summary")) last_processed = message # type: ignore - - elif isinstance(message, Response): - duration = time.time() - start_time - - output = Text.from_markup(f"{_message_to_str(message.chat_message, render_image_iterm=render_image_iterm)}") - if message.chat_message.models_usage: - output.append( - f"\n[Prompt tokens: {message.chat_message.models_usage.prompt_tokens}, Completion tokens: {message.chat_message.models_usage.completion_tokens}]" - ) - total_usage.completion_tokens += message.chat_message.models_usage.completion_tokens - total_usage.prompt_tokens += message.chat_message.models_usage.prompt_tokens - console.print( - Panel(output, title=f"[bold {primary_color}]{message.chat_message.source}[/bold {primary_color}]") - ) - - if message.inner_messages is not None: - num_inner_messages = len(message.inner_messages) - else: - num_inner_messages = 0 - output = ( - f"Number of inner messages: {num_inner_messages}\n" - f"Total prompt tokens: {total_usage.prompt_tokens}\n" - f"Total completion tokens: {total_usage.completion_tokens}\n" - f"Duration: {duration:.2f} seconds\n" - ) - console.print(Panel(output, title="Summary")) - last_processed = message # type: ignore - else: message = cast(AgentEvent | ChatMessage, message) # type: ignore output = Text.from_markup(f"{_message_to_str(message, render_image_iterm=render_image_iterm)}")