Terminal colors

guidance/_utils.py

@@ -6,9 +6,18 @@
 import sys
 import textwrap
 import types
+import re
 
 import numpy as np
 
+from html.parser import HTMLParser
+try:
+    from colored import Fore  # type: ignore[import-untyped]
+except ImportError:
+    colored_is_imported = False
+else:
+    colored_is_imported = True
+
 
 class _Rewrite(ast.NodeTransformer):
     def visit_Constant(self, node):
@@ -261,3 +270,38 @@ def softmax(array: np.ndarray, axis: int = -1) -> np.ndarray:
     array_maxs = np.amax(array, axis=axis, keepdims=True)
     exp_x_shifted = np.exp(array - array_maxs)
     return exp_x_shifted / np.sum(exp_x_shifted, axis=axis, keepdims=True)
+
+
+class ModelStateHTMLParser(HTMLParser):
+    """Parse jupyter-flavored HTML that contains colored text to color text for the command-line"""
+    def __init__(self):
+        super().__init__()
+        self.colored_text = ''
+
+    def feed(self, data):
+        self.colored_text = ''
+        # Remove html insertion tags (I suppose it is for jupyter to recognize it as html markdown)
+        data = data.replace('<||_html:', '').replace('_||>', '')
+        super().feed(data)
+        return self.colored_text
+
+    def handle_starttag(self, tag, attrs):
+        # Start ANSI text coloring when span tag opens
+        if tag == 'span':
+            if colored_is_imported:
+                # Use bg color we would have used in jupyter as fg color (get from style attributes)
+                style = dict(attrs)['style']
+                # just capture integer rgb parts of rgba color
+                rgb = re.search(r'background-color:\s*rgba\((\d+)\.?\d*,\s*(\d+)\.?\d*,\s*(\d)+\.?\d*,\s*\d+\.?\d*\)', style).groups()
+                self.colored_text += Fore.rgb(*rgb)
+            else:
+                # Default to ANSI green (32m) if colored is not available
+                self.colored_text += '\033[32m'
+
+    def handle_endtag(self, tag):
+        # ANSI reset color when tag closes
+        if tag == 'span':
+            self.colored_text += '\033[0m'
+
+    def handle_data(self, data):
+        self.colored_text += data

guidance/models/_model.py

@@ -17,11 +17,19 @@
 import numpy as np
 
 try:
+    from IPython import get_ipython
     from IPython.display import clear_output, display, HTML
-
-    ipython_is_imported = True
 except ImportError:
     ipython_is_imported = False
+    notebook_mode = False
+else:
+    ipython_is_imported = True
+    _ipython = get_ipython()
+    notebook_mode = (
+        _ipython is not None
+        and "IPKernelApp" in _ipython.config
+    )
+
 try:
     import torch
 
@@ -39,7 +47,7 @@
     )
     from .. import _cpp as cpp
 from ._guidance_engine_metrics import GuidanceEngineMetrics
-from .._utils import softmax, CaptureEvents
+from .._utils import softmax, CaptureEvents, ModelStateHTMLParser
 from .._parser import EarleyCommitParser, Parser
 from .._grammar import (
     GrammarFunction,
@@ -862,6 +870,8 @@ def __init__(self, engine, echo=True, **kwargs):
         self._event_parent = None
         self._last_display = 0  # used to track the last display call to enable throttling
         self._last_event_stream = 0  # used to track the last event streaming call to enable throttling
+        self._state_html_parser = ModelStateHTMLParser() # used to parse the state for cli display
+        self._last_state_len = 0 # used to track the last state length for appending to cli display
 
     @property
     def active_role_end(self):
@@ -975,11 +985,18 @@ def _update_display(self, throttle=True):
                 else:
                     self._last_display = curr_time
 
-            if ipython_is_imported:
+            if notebook_mode:
                 clear_output(wait=True)
                 display(HTML(self._html()))
             else:
-                pprint(self._state)
+                print(
+                    self._state_html_parser.feed(
+                        self._state[self._last_state_len:]
+                    ),
+                    end='',
+                    flush=True
+                )
+                self._last_state_len = len(self._state)
 
     def reset(self, clear_variables=True):
         """This resets the state of the model object.

setup.py

@@ -37,6 +37,7 @@
     "openai": ["openai>=1.0"],
     "schemas": ["jsonschema"],
     "server": ["fastapi", "uvicorn"],
+    "cli": ["colored"]
 }
 
 # Create the union of all our requirements