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 <ekzhu@users.noreply.github.com>

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__)

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):

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__)
 

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__)

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

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__)

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)

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

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__), "../../.."))