From e35db7ef5f5f23ab2a77bca545a9bb0e53b949c7 Mon Sep 17 00:00:00 2001 From: Jack Gerrits Date: Tue, 19 Mar 2024 14:09:43 -0400 Subject: [PATCH] Encapsulate colored into a module (#2057) * Encapsulate colored into a module * lint fix * add missing file * undo change * conform with original colored func * change import strategy --------- Co-authored-by: Eric Zhu --- autogen/agentchat/chat.py | 2 +- .../contrib/capabilities/teachability.py | 2 +- .../agentchat/contrib/compressible_agent.py | 8 +-- autogen/agentchat/contrib/llava_agent.py | 2 +- .../contrib/retrieve_user_proxy_agent.py | 2 +- autogen/agentchat/conversable_agent.py | 9 +-- autogen/formatting_utils.py | 70 +++++++++++++++++++ .../capabilities/chat_with_teachable_agent.py | 2 +- .../capabilities/test_teachable_agent.py | 2 +- 9 files changed, 78 insertions(+), 21 deletions(-) create mode 100644 autogen/formatting_utils.py diff --git a/autogen/agentchat/chat.py b/autogen/agentchat/chat.py index bd4559ff425..d5e127e971f 100644 --- a/autogen/agentchat/chat.py +++ b/autogen/agentchat/chat.py @@ -7,7 +7,7 @@ from .utils import consolidate_chat_info import datetime import warnings -from termcolor import colored +from ..formatting_utils import colored logger = logging.getLogger(__name__) diff --git a/autogen/agentchat/contrib/capabilities/teachability.py b/autogen/agentchat/contrib/capabilities/teachability.py index 9e18f99a345..58ba35ed425 100644 --- a/autogen/agentchat/contrib/capabilities/teachability.py +++ b/autogen/agentchat/contrib/capabilities/teachability.py @@ -6,7 +6,7 @@ from autogen.agentchat.assistant_agent import ConversableAgent from autogen.agentchat.contrib.capabilities.agent_capability import AgentCapability from autogen.agentchat.contrib.text_analyzer_agent import TextAnalyzerAgent -from autogen.agentchat.conversable_agent import colored +from ....formatting_utils import colored class Teachability(AgentCapability): diff --git a/autogen/agentchat/contrib/compressible_agent.py b/autogen/agentchat/contrib/compressible_agent.py index 0264d265654..152cc871a56 100644 --- a/autogen/agentchat/contrib/compressible_agent.py +++ b/autogen/agentchat/contrib/compressible_agent.py @@ -7,13 +7,7 @@ import inspect from autogen.token_count_utils import count_token, get_max_token_limit, num_tokens_from_functions -try: - from termcolor import colored -except ImportError: - - def colored(x, *args, **kwargs): - return x - +from ...formatting_utils import colored logger = logging.getLogger(__name__) diff --git a/autogen/agentchat/contrib/llava_agent.py b/autogen/agentchat/contrib/llava_agent.py index c26f576ab39..182f72837b7 100644 --- a/autogen/agentchat/contrib/llava_agent.py +++ b/autogen/agentchat/contrib/llava_agent.py @@ -8,7 +8,7 @@ from autogen.agentchat.contrib.img_utils import get_image_data, llava_formatter from autogen.agentchat.contrib.multimodal_conversable_agent import MultimodalConversableAgent from autogen.code_utils import content_str -from autogen.agentchat.conversable_agent import colored +from ...formatting_utils import colored logger = logging.getLogger(__name__) diff --git a/autogen/agentchat/contrib/retrieve_user_proxy_agent.py b/autogen/agentchat/contrib/retrieve_user_proxy_agent.py index 5d9657d8143..f252f60e5ec 100644 --- a/autogen/agentchat/contrib/retrieve_user_proxy_agent.py +++ b/autogen/agentchat/contrib/retrieve_user_proxy_agent.py @@ -12,7 +12,7 @@ from autogen.token_count_utils import count_token from autogen.code_utils import extract_code from autogen import logger -from autogen.agentchat.conversable_agent import colored +from ...formatting_utils import colored PROMPT_DEFAULT = """You're a retrieve augmented chatbot. You answer user's questions based on your own knowledge and the diff --git a/autogen/agentchat/conversable_agent.py b/autogen/agentchat/conversable_agent.py index 838af6e7df7..accb2c54d9b 100644 --- a/autogen/agentchat/conversable_agent.py +++ b/autogen/agentchat/conversable_agent.py @@ -15,6 +15,7 @@ from ..coding.base import CodeExecutor from ..coding.factory import CodeExecutorFactory +from ..formatting_utils import colored from ..oai.client import OpenAIWrapper, ModelClient from ..runtime_logging import logging_enabled, log_new_agent @@ -36,14 +37,6 @@ from .agent import Agent, LLMAgent from .._pydantic import model_dump -try: - from termcolor import colored -except ImportError: - - def colored(x, *args, **kwargs): - return x - - __all__ = ("ConversableAgent",) logger = logging.getLogger(__name__) diff --git a/autogen/formatting_utils.py b/autogen/formatting_utils.py new file mode 100644 index 00000000000..be8d8aaa562 --- /dev/null +++ b/autogen/formatting_utils.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from typing import Iterable, Literal + +try: + from termcolor import colored +except ImportError: + # termcolor is an optional dependency - if it cannot be imported then no color is used. + # Alternatively the envvar NO_COLOR can be used to disable color. + # To allow for proper typing and for termcolor to be optional we need to re-define the types used in the lib here. + # This is the direct function definition from termcolor. + Attribute = Literal[ + "bold", + "dark", + "underline", + "blink", + "reverse", + "concealed", + ] + + Highlight = Literal[ + "on_black", + "on_grey", + "on_red", + "on_green", + "on_yellow", + "on_blue", + "on_magenta", + "on_cyan", + "on_light_grey", + "on_dark_grey", + "on_light_red", + "on_light_green", + "on_light_yellow", + "on_light_blue", + "on_light_magenta", + "on_light_cyan", + "on_white", + ] + + Color = Literal[ + "black", + "grey", + "red", + "green", + "yellow", + "blue", + "magenta", + "cyan", + "light_grey", + "dark_grey", + "light_red", + "light_green", + "light_yellow", + "light_blue", + "light_magenta", + "light_cyan", + "white", + ] + + def colored( + text: object, + color: Color | None = None, + on_color: Highlight | None = None, + attrs: Iterable[Attribute] | None = None, + *, + no_color: bool | None = None, + force_color: bool | None = None, + ) -> str: + return str(text) diff --git a/test/agentchat/contrib/capabilities/chat_with_teachable_agent.py b/test/agentchat/contrib/capabilities/chat_with_teachable_agent.py index 66d61386d61..27f030045a7 100755 --- a/test/agentchat/contrib/capabilities/chat_with_teachable_agent.py +++ b/test/agentchat/contrib/capabilities/chat_with_teachable_agent.py @@ -2,10 +2,10 @@ import os import sys -from termcolor import colored from autogen import UserProxyAgent, config_list_from_json from autogen.agentchat.contrib.capabilities.teachability import Teachability from autogen import ConversableAgent +from autogen.formatting_utils import colored sys.path.append(os.path.join(os.path.dirname(__file__), "../..")) from test_assistant_agent import OAI_CONFIG_LIST, KEY_LOC # noqa: E402 diff --git a/test/agentchat/contrib/capabilities/test_teachable_agent.py b/test/agentchat/contrib/capabilities/test_teachable_agent.py index 44904b26d36..e9fdf68c813 100755 --- a/test/agentchat/contrib/capabilities/test_teachable_agent.py +++ b/test/agentchat/contrib/capabilities/test_teachable_agent.py @@ -3,7 +3,7 @@ import pytest import os import sys -from termcolor import colored +from autogen.formatting_utils import colored from autogen import ConversableAgent, config_list_from_json sys.path.append(os.path.join(os.path.dirname(__file__), "../../.."))