adds a terminal_output function to pretty-print model output in CLI.

PiperOrigin-RevId: 820159448

Author: aelissee

examples/realtime_simple_cli.py

@@ -58,7 +58,6 @@
 from genai_processors.core import text_to_speech
 from genai_processors.examples import models
 import pyaudio
-import termcolor
 
 # You need to define the project id in the environment variables.
 GOOGLE_PROJECT_ID = os.environ['GOOGLE_PROJECT_ID']
@@ -144,23 +143,10 @@ async def run_conversation() -> None:
       + realtime.LiveProcessor(turn_processor=genai_processor + tts)
       + play_output
   )
-
-  async for part in conversation_agent(text.terminal_input()):
-    # Print the transcription and the output of the model (should include status
-    # parts and other metadata parts)
-    match part.role:
-      case 'user':
-        color = 'green'
-      case 'model':
-        color = 'red'
-      case _:
-        color = 'yellow'
-    part_role = part.role or 'default'
-    print(
-        termcolor.colored(
-            f'{part_role}: {part.text}', color, 'on_grey', attrs=['bold']
-        )
-    )
+  prompt = 'USER (ctrl+D to end)> '
+  await text.terminal_output(
+      conversation_agent(text.terminal_input(prompt=prompt)), prompt=prompt
+  )
 
 
 def main(argv: Sequence[str]):

genai_processors/core/text.py

@@ -23,6 +23,7 @@
 import dataclasses_json
 from genai_processors import content_api
 from genai_processors import processor
+import termcolor
 
 
 _MAX_LOOP_COUNT = 1000
@@ -229,7 +230,7 @@ def _extract_part(
         ]
         break
       offset += len(c.text)
-    to_process = left_over + part_buffer[part_idx + 1:]
+    to_process = left_over + part_buffer[part_idx + 1 :]
     return to_yield, to_process
 
   async def call(
@@ -414,3 +415,53 @@ async def terminal_input(
     except EOFError:
       # Exit on ctrl+D.
       return
+
+
+async def terminal_output(
+    content: AsyncIterable[content_api.ProcessorPartTypes],
+    prompt: str = '',
+) -> None:
+  """Prints the part to the terminal.
+
+  Consumes all the content and prints it to the terminal. Prints the prompt
+  when an `end_of_turn` part is encountered.
+
+  The parts are printed with their role in green, red or yellow. The text parts
+  are printed in bold.
+
+  Args:
+    content: The content to print.
+    prompt: The prompt to print when the model is done.
+  """
+  old_part_role = None
+  async for part in content:
+    match part.role:
+      case 'user':
+        color = 'green'
+      case 'model':
+        color = 'red'
+      case _:
+        color = 'yellow'
+    part_role = part.role or 'default'
+    if content_api.is_text(part.mimetype):
+      content_text = part.text
+    else:
+      content_text = f'<{part.mimetype}>'
+    if part_role != old_part_role:
+      old_part_role = part_role
+      print(
+          termcolor.colored(
+              f'\n{part_role}: {content_text}', color, attrs=['bold']
+          ),
+          end='',
+          flush=True,
+      )
+    else:
+      print(
+          termcolor.colored(f'{content_text}', color, attrs=['bold']),
+          end='',
+          flush=True,
+      )
+    # Reprint the prompt when the model is done.
+    if content_api.is_end_of_turn(part):
+      print('\n' + prompt, end='', flush=True)