Remove text_formatting.py.

PiperOrigin-RevId: 713324457

langfun/core/__init__.py

@@ -77,11 +77,6 @@
 from langfun.core.concurrent import with_context_access
 from langfun.core.concurrent import with_retry
 
-# Utility libraries for text formatting.
-from langfun.core.text_formatting import colored
-from langfun.core.text_formatting import colored_print as print  # pylint: disable=redefined-builtin
-from langfun.core.text_formatting import colored_template
-
 # Interface for natural language formattable.
 from langfun.core.natural_language import NaturalLanguageFormattable
 

langfun/core/coding/python/correction.py

@@ -192,9 +192,7 @@ def correct(
 def _error_feedback_str(error: Exception) -> str:
   """Returns the error str for feedback."""
   if isinstance(error, errors.CodeError):
-    return lf.text_formatting.decolored(
-        error.format(include_complete_code=False)
-    )
+    return pg.decolor(error.format(include_complete_code=False))
   else:
     return f"Encountered {error.__class__.__name__}: {error}"
 

langfun/core/coding/python/errors.py

@@ -17,7 +17,8 @@
 import sys
 import textwrap
 import traceback
-import langfun.core as lf
+
+import pyglove as pg
 
 
 class CodeError(RuntimeError):
@@ -62,28 +63,28 @@ def format(self, include_complete_code: bool = True):
     if 'line' not in error_message and self.lineno is not None:
       error_message += f' (<unknown>, line {self.lineno})'
     r.write(
-        lf.colored(
+        pg.colored(
             f'{self.cause.__class__.__name__}: {error_message}', 'magenta'))
 
     if self.lineno is not None:
       r.write('\n\n')
       r.write(textwrap.indent(
-          lf.colored(
+          pg.colored(
               self.code_lines(self.lineno - 1, self.end_lineno), 'magenta'),
           ' ' * 2
       ))
       r.write('\n')
 
     if include_complete_code:
       r.write('\n')
-      r.write(lf.colored('[Generated Code]', 'green', styles=['bold']))
+      r.write(pg.colored('[Generated Code]', 'green', styles=['bold']))
       r.write('\n\n')
-      r.write(lf.colored('  ```python\n', 'green'))
+      r.write(pg.colored('  ```python\n', 'green'))
       r.write(textwrap.indent(
-          lf.colored(self.code, 'green'),
+          pg.colored(self.code, 'green'),
           ' ' * 2
       ))
-      r.write(lf.colored('\n  ```\n', 'green'))
+      r.write(pg.colored('\n  ```\n', 'green'))
     return r.getvalue()
 
 
@@ -98,10 +99,10 @@ def __str__(self):
     r = io.StringIO()
     cause_message = str(self.cause).rstrip()
     if self.message:
-      r.write(lf.colored(self.message, 'magenta'))
+      r.write(pg.colored(self.message, 'magenta'))
       r.write('\n\n')
     r.write(
-        lf.colored(
+        pg.colored(
             f'{self.cause.__class__.__name__}: {cause_message}', 'magenta'
         )
     )

langfun/core/concurrent.py

@@ -25,7 +25,6 @@
 from typing import Any, Callable, Iterable, Iterator, Literal, Sequence, Tuple, Type, Union
 
 from langfun.core import component
-from langfun.core import text_formatting
 import pyglove as pg
 
 
@@ -844,10 +843,10 @@ def update(self, delta: int) -> None:
   def refresh(self):
     s = io.StringIO()
     if self.label is not None:
-      s.write(text_formatting.colored(self.label, 'red', styles=['bold']))
+      s.write(pg.colored(self.label, 'red', styles=['bold']))
       s.write(': ')
     s.write(
-        text_formatting.colored(
+        pg.colored(
             '%d%% (%d/%d)' %
             (
                 self._progress * 100 // self.total,

langfun/core/console.py

@@ -15,7 +15,7 @@
 
 import sys
 from typing import Any
-from langfun.core.text_formatting import colored
+import pyglove as pg
 
 
 def write(
@@ -42,10 +42,15 @@ def write(
   """
   # Print title if present.
   if title is not None:
-    print(colored(title, styles=['bold']))
+    print(pg.colored(title, styles=['bold']))
 
   # Print body.
-  print(colored(str(value), color=color, background=background, styles=styles))
+  print(dir(pg.utils))
+  print(
+      pg.colored(
+          str(value), color=color, background=background, styles=styles
+      )
+  )
 
 
 try:

langfun/core/eval/base.py

@@ -1298,7 +1298,7 @@ def finalize(self) -> pg.Dict:
             id=self.id,
             dir=self.dir,
             model=self.lm.model_id,
-            prompt_template=lf.text_formatting.decolored(str(self.prompt)),
+            prompt_template=pg.decolor(str(self.prompt)),
             method=self.method,
             schema_fn=str(self.schema_fn),
         ),
@@ -2110,8 +2110,7 @@ def monitor(
 
 def _format_error(error: Exception):
   """Formats an error into a string."""
-  return (f'({error.__class__.__name__}) '
-          + lf.text_formatting.decolored(str(error)))
+  return (f'({error.__class__.__name__}) ' + pg.decolor(str(error)))
 
 
 def _error_key(error: Exception) -> str:

langfun/core/language_model.py

@@ -653,7 +653,7 @@ def _debug_model_info(
     """Outputs debugging information about the model."""
     title_suffix = ''
     if usage.total_tokens != 0:
-      title_suffix = console.colored(
+      title_suffix = pg.colored(
           f' (total {usage.total_tokens} tokens)', 'red'
       )
 
@@ -672,7 +672,7 @@ def _debug_prompt(
     """Outputs debugging information about the prompt."""
     title_suffix = ''
     if usage.prompt_tokens != 0:
-      title_suffix = console.colored(f' ({usage.prompt_tokens} tokens)', 'red')
+      title_suffix = pg.colored(f' ({usage.prompt_tokens} tokens)', 'red')
 
     console.write(
         # We use metadata 'formatted_text' for scenarios where the prompt text
@@ -703,7 +703,7 @@ def _debug_response(
     if usage.completion_tokens != 0:
       title_suffix += f'{usage.completion_tokens} tokens '
     title_suffix += f'in {elapse:.2f} seconds)'
-    title_suffix = console.colored(title_suffix, 'red')
+    title_suffix = pg.colored(title_suffix, 'red')
 
     console.write(
         str(response) + '\n',

langfun/core/structured/mapping.py

@@ -46,15 +46,15 @@ def format(self, include_lm_response: bool = True) -> str:
     r = io.StringIO()
     error_message = str(self.cause).rstrip()
     r.write(
-        lf.colored(
+        pg.colored(
             f'{self.cause.__class__.__name__}: {error_message}', 'magenta'
         )
     )
     if include_lm_response:
       r.write('\n\n')
-      r.write(lf.colored('[LM Response]', 'blue', styles=['bold']))
+      r.write(pg.colored('[LM Response]', 'blue', styles=['bold']))
       r.write('\n')
-      r.write(lf.colored(self.lm_response.text, 'blue'))
+      r.write(pg.colored(self.lm_response.text, 'blue'))
     return r.getvalue()
 
 
@@ -163,27 +163,27 @@ def output_repr(
   def natural_language_format(self) -> str:
     result = io.StringIO()
     if self.context:
-      result.write(lf.colored('[CONTEXT]\n', styles=['bold']))
-      result.write(lf.colored(self.context, color='magenta'))
+      result.write(pg.colored('[CONTEXT]\n', styles=['bold']))
+      result.write(pg.colored(self.context, color='magenta'))
       result.write('\n\n')
 
-    result.write(lf.colored('[INPUT]\n', styles=['bold']))
-    result.write(lf.colored(self.input_repr(), color='green'))
+    result.write(pg.colored('[INPUT]\n', styles=['bold']))
+    result.write(pg.colored(self.input_repr(), color='green'))
 
     if self.schema is not None:
       result.write('\n\n')
-      result.write(lf.colored('[SCHEMA]\n', styles=['bold']))
-      result.write(lf.colored(self.schema_repr(), color='red'))
+      result.write(pg.colored('[SCHEMA]\n', styles=['bold']))
+      result.write(pg.colored(self.schema_repr(), color='red'))
 
     if schema_lib.MISSING != self.output:
       result.write('\n\n')
-      result.write(lf.colored('[OUTPUT]\n', styles=['bold']))
-      result.write(lf.colored(self.output_repr(), color='blue'))
+      result.write(pg.colored('[OUTPUT]\n', styles=['bold']))
+      result.write(pg.colored(self.output_repr(), color='blue'))
 
     if self.metadata:
       result.write('\n\n')
-      result.write(lf.colored('[METADATA]\n', styles=['bold']))
-      result.write(lf.colored(str(self.metadata), color='cyan'))
+      result.write(pg.colored('[METADATA]\n', styles=['bold']))
+      result.write(pg.colored(str(self.metadata), color='cyan'))
     return result.getvalue().strip()
 
   @classmethod

langfun/core/structured/mapping_test.py

@@ -29,11 +29,11 @@ def test_format(self):
         lf.AIMessage('hi'), ValueError('Cannot parse message.')
     )
     self.assertEqual(
-        lf.text_formatting.decolored(str(error)),
+        pg.decolor(str(error)),
         'ValueError: Cannot parse message.\n\n[LM Response]\nhi',
     )
     self.assertEqual(
-        lf.text_formatting.decolored(error.format(include_lm_response=False)),
+        pg.decolor(error.format(include_lm_response=False)),
         'ValueError: Cannot parse message.',
     )
 

langfun/core/structured/schema.py

@@ -91,20 +91,25 @@ def __init__(self,
   def __str__(self):
     r = io.StringIO()
     r.write(
-        lf.colored(f'{self.cause.__class__.__name__}: {self.cause}', 'magenta'))
+        pg.colored(
+            f'{self.cause.__class__.__name__}: {self.cause}', 'magenta'
+        )
+    )
 
     r.write('\n')
-    r.write(lf.colored('Schema:', 'red'))
+    r.write(pg.colored('Schema:', 'red'))
     r.write('\n\n')
     r.write(textwrap.indent(
-        lf.colored(schema_repr(self.protocol).repr(self.schema), 'magenta'),
+        pg.colored(
+            schema_repr(self.protocol).repr(self.schema), 'magenta'
+        ),
         ' ' * 2
     ))
     r.write('\n\n')
-    r.write(lf.colored('Generated value:', 'red'))
+    r.write(pg.colored('Generated value:', 'red'))
     r.write('\n\n')
     r.write(textwrap.indent(
-        lf.colored(value_repr(self.protocol).repr(self.value), 'magenta'),
+        pg.colored(value_repr(self.protocol).repr(self.value), 'magenta'),
         ' ' * 2
     ))
     return r.getvalue()
@@ -759,12 +764,15 @@ def __init__(self, json: str, cause: Exception):
   def __str__(self) -> str:
     r = io.StringIO()
     r.write(
-        lf.colored(f'{self.cause.__class__.__name__}: {self.cause}', 'magenta'))
+        pg.colored(
+            f'{self.cause.__class__.__name__}: {self.cause}', 'magenta'
+        )
+    )
 
     r.write('\n\n')
-    r.write(lf.colored('JSON text:', 'red'))
+    r.write(pg.colored('JSON text:', 'red'))
     r.write('\n\n')
-    r.write(textwrap.indent(lf.colored(self.json, 'magenta'), ' ' * 2))
+    r.write(textwrap.indent(pg.colored(self.json, 'magenta'), ' ' * 2))
     return r.getvalue()