From ca3600543d47a52e403bd15b3da63608342f694e Mon Sep 17 00:00:00 2001 From: Malachi Soord Date: Fri, 26 Sep 2025 23:34:13 +0200 Subject: [PATCH] add: prompt rendering with rich --- docs/help.md | 1 + llm/cli.py | 84 +++++++++++++++++++++++++++++++------------------- pyproject.toml | 1 + 3 files changed, 55 insertions(+), 31 deletions(-) diff --git a/docs/help.md b/docs/help.md index a5a63566..028e91d7 100644 --- a/docs/help.md +++ b/docs/help.md @@ -152,6 +152,7 @@ Options: -u, --usage Show token usage -x, --extract Extract first fenced code block --xl, --extract-last Extract last fenced code block + --render Render for terminal -h, --help Show this message and exit. ``` diff --git a/llm/cli.py b/llm/cli.py index 2e11e2c8..8c580501 100644 --- a/llm/cli.py +++ b/llm/cli.py @@ -5,6 +5,11 @@ import io import json import os + +from rich.console import Console +from rich.live import Live +from rich.markdown import Markdown + from llm import ( Attachment, AsyncConversation, @@ -471,6 +476,7 @@ def cli(): is_flag=True, help="Extract last fenced code block", ) +@click.option("--render", is_flag=True, help="Render for terminal") def prompt( prompt, system, @@ -502,6 +508,7 @@ def prompt( usage, extract, extract_last, + render, ): """ Execute a prompt @@ -830,42 +837,45 @@ def read_prompt(): # Merge in options for the .prompt() methods kwargs.update(validated_options) + console = Console() + try: if async_: async def inner(): + response = prompt_method( + prompt, + attachments=resolved_attachments, + system=system, + schema=schema, + fragments=resolved_fragments, + system_fragments=resolved_system_fragments, + **kwargs, + ) + if should_stream: - response = prompt_method( - prompt, - attachments=resolved_attachments, - system=system, - schema=schema, - fragments=resolved_fragments, - system_fragments=resolved_system_fragments, - **kwargs, - ) - async for chunk in response: - print(chunk, end="") - sys.stdout.flush() - print("") + accumulated_text = "" + with Live(accumulated_text, console=console, refresh_per_second=10) as live: + async for chunk in response: + accumulated_text += chunk + + if render: + display_content = Markdown(accumulated_text) + else: + display_content = accumulated_text + + live.update(display_content) else: - response = prompt_method( - prompt, - fragments=resolved_fragments, - attachments=resolved_attachments, - schema=schema, - system=system, - system_fragments=resolved_system_fragments, - **kwargs, - ) text = await response.text() if extract or extract_last: - text = ( - extract_fenced_code_block(text, last=extract_last) or text - ) - print(text) + text = extract_fenced_code_block(text, last=extract_last) or text + if render: + text = Markdown(text) + console.print(text) + return response + response = asyncio.run(inner()) else: response = prompt_method( @@ -877,16 +887,28 @@ async def inner(): system_fragments=resolved_system_fragments, **kwargs, ) + + + if should_stream: - for chunk in response: - print(chunk, end="") - sys.stdout.flush() - print("") + accumulated_text = "" + with Live(accumulated_text, console=console, refresh_per_second=10) as live: + for chunk in response: + accumulated_text += chunk + + if render: + display_content = Markdown(accumulated_text) + else: + display_content = accumulated_text + + live.update(display_content) else: text = response.text() if extract or extract_last: text = extract_fenced_code_block(text, last=extract_last) or text - print(text) + if render: + text = Markdown(text) + console.print(text) # List of exceptions that should never be raised in pytest: except (ValueError, NotImplementedError) as ex: raise click.ClickException(str(ex)) diff --git a/pyproject.toml b/pyproject.toml index a78b7c53..19077edd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ dependencies = [ "pip", "pyreadline3; sys_platform == 'win32'", "puremagic", + "rich" ] [project.urls]