Feat/better log: Add colorize function and TermColor enum for text coloring (#5410)

Co-authored-by: Engel Nyst <enyst@users.noreply.github.com>

openhands/core/logger.py

@@ -5,7 +5,8 @@ import re
 import sys
 import traceback
 from datetime import datetime
-from typing import Literal, Mapping
+from types import TracebackType
+from typing import Any, Literal, Mapping
 
 from termcolor import colored
 
@@ -61,7 +62,8 @@ class NoColorFormatter(logging.Formatter):
 
 
 def strip_ansi(s: str) -> str:
-    """
+    """Remove ANSI escape sequences (terminal color/formatting codes) from string.
+
     Removes ANSI escape sequences from str, as defined by ECMA-048 in
     http://www.ecma-international.org/publications/files/ECMA-ST/Ecma-048.pdf
     # https://github.com/ewen-lbh/python-strip-ansi/blob/master/strip_ansi/__init__.py
@@ -136,6 +138,7 @@ class RollingLogger:
 
     def print_lines(self):
         """Display the last n log_lines in the console (not for file logging).
+
         This will create the effect of a rolling display in the console.
         """
         self.move_back()
@@ -143,18 +146,14 @@ class RollingLogger:
             self.replace_current_line(line)
 
     def move_back(self, amount=-1):
-        """
-        '\033[F'    moves the cursor up one line.
-        """
+        r"""'\033[F' moves the cursor up one line."""
         if amount == -1:
             amount = self.max_lines
         self._write('\033[F' * (self.max_lines))
         self._flush()
 
     def replace_current_line(self, line=''):
-        """
-        '\033[2K\r' clears the line and moves the cursor to the beginning of the line.
-        """
+        r"""'\033[2K\r' clears the line and moves the cursor to the beginning of the line."""
         self._write('\033[2K' + line + '\n')
         self._flush()
 
@@ -232,18 +231,21 @@ def get_file_handler(log_dir: str, log_level: int = logging.INFO):
 logging.basicConfig(level=logging.ERROR)
 
 
-def log_uncaught_exceptions(ex_cls, ex, tb):
+def log_uncaught_exceptions(
+    ex_cls: type[BaseException], ex: BaseException, tb: TracebackType | None
+) -> Any:
     """Logs uncaught exceptions along with the traceback.
 
     Args:
-        ex_cls (type): The type of the exception.
-        ex (Exception): The exception instance.
-        tb (traceback): The traceback object.
+        ex_cls: The type of the exception.
+        ex: The exception instance.
+        tb: The traceback object.
 
     Returns:
         None
     """
-    logging.error(''.join(traceback.format_tb(tb)))
+    if tb:  # Add check since tb can be None
+        logging.error(''.join(traceback.format_tb(tb)))
     logging.error('{0}: {1}'.format(ex_cls, ex))
 
 
@@ -283,7 +285,7 @@ logging.getLogger('LiteLLM Proxy').disabled = True
 
 
 class LlmFileHandler(logging.FileHandler):
-    """# LLM prompt and response logging"""
+    """LLM prompt and response logging."""
 
     def __init__(self, filename, mode='a', encoding='utf-8', delay=False):
         """Initializes an instance of LlmFileHandler.

openhands/runtime/builder/base.py

@@ -9,13 +9,13 @@ class RuntimeBuilder(abc.ABC):
         tags: list[str],
         platform: str | None = None,
     ) -> str:
-        """
-        Build the runtime image.
+        """Build the runtime image.
 
         Args:
             path (str): The path to the runtime image's build directory.
             tags (list[str]): The tags to apply to the runtime image (e.g., ["repo:my-repo", "sha:my-sha"]).
             platform (str, optional): The target platform for the build. Defaults to None.
+
         Returns:
             str: The name:tag of the runtime image after build (e.g., "repo:sha").
                 This can be different from the tags input if the builder chooses to mutate the tags (e.g., adding a
@@ -28,8 +28,7 @@ class RuntimeBuilder(abc.ABC):
 
     @abc.abstractmethod
     def image_exists(self, image_name: str, pull_from_repo: bool = True) -> bool:
-        """
-        Check if the runtime image exists.
+        """Check if the runtime image exists.
 
         Args:
             image_name (str): The name of the runtime image (e.g., "repo:sha").

openhands/runtime/builder/docker.py

@@ -9,6 +9,7 @@ from openhands import __version__ as oh_version
 from openhands.core.logger import RollingLogger
 from openhands.core.logger import openhands_logger as logger
 from openhands.runtime.builder.base import RuntimeBuilder
+from openhands.utils.term_color import TermColor, colorize
 
 
 class DockerRuntimeBuilder(RuntimeBuilder):
@@ -187,7 +188,9 @@ class DockerRuntimeBuilder(RuntimeBuilder):
             return True
         except docker.errors.ImageNotFound:
             if not pull_from_repo:
-                logger.debug(f'Image {image_name} not found locally')
+                logger.debug(
+                    f'Image {image_name} {colorize("not found", TermColor.WARNING)} locally'
+                )
                 return False
             try:
                 logger.debug(
@@ -214,7 +217,7 @@ class DockerRuntimeBuilder(RuntimeBuilder):
                 logger.debug('Could not find image locally or in registry.')
                 return False
             except Exception as e:
-                msg = 'Image could not be pulled: '
+                msg = f'Image {colorize("could not be pulled", TermColor.ERROR)}: '
                 ex_msg = str(e)
                 if 'Not Found' in ex_msg:
                     msg += 'image not found in registry.'
@@ -286,8 +289,7 @@ class DockerRuntimeBuilder(RuntimeBuilder):
             logger.debug(current_line['status'])
 
     def _prune_old_cache_files(self, cache_dir: str, max_age_days: int = 7) -> None:
-        """
-        Prune cache files older than the specified number of days.
+        """Prune cache files older than the specified number of days.
 
         Args:
             cache_dir (str): The path to the cache directory.
@@ -311,8 +313,7 @@ class DockerRuntimeBuilder(RuntimeBuilder):
             logger.warning(f'Error during build cache pruning: {e}')
 
     def _is_cache_usable(self, cache_dir: str) -> bool:
-        """
-        Check if the cache directory is usable (exists and is writable).
+        """Check if the cache directory is usable (exists and is writable).
 
         Args:
             cache_dir (str): The path to the cache directory.

openhands/utils/term_color.py

@@ -0,0 +1,25 @@
+from enum import Enum
+
+from termcolor import colored
+
+
+class TermColor(Enum):
+    """Terminal color codes."""
+
+    WARNING = 'yellow'
+    SUCCESS = 'green'
+    ERROR = 'red'
+    INFO = 'blue'
+
+
+def colorize(text: str, color: TermColor = TermColor.WARNING) -> str:
+    """Colorize text with specified color.
+
+    Args:
+        text (str): Text to be colored
+        color (TermColor, optional): Color to use. Defaults to TermColor.WARNING
+
+    Returns:
+        str: Colored text
+    """
+    return colored(text, color.value)