add: prompt rendering with rich

Author: inverse

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.
 ```
 

llm/cli.py

@@ -41,6 +41,7 @@
     remove_alias,
 )
 from llm.models import _BaseConversation, ChainResponse
+from .cli_utils import prompt_output
 
 from .migrations import migrate
 from .plugins import pm, load_plugins
@@ -471,6 +472,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 +504,7 @@ def prompt(
     usage,
     extract,
     extract_last,
+    render,
 ):
     """
     Execute a prompt
@@ -834,59 +837,37 @@ def read_prompt():
         if async_:
 
             async def inner():
-                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("")
-                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)
-                return response
+                return prompt_output(
+                    prompt_method,
+                    prompt,
+                    should_stream,
+                    render,
+                    extract,
+                    extract_last,
+                    resolved_fragments,
+                    resolved_attachments,
+                    system,
+                    schema,
+                    resolved_system_fragments,
+                    kwargs,
+                )
 
             response = asyncio.run(inner())
         else:
-            response = prompt_method(
+            response = prompt_output(
+                prompt_method,
                 prompt,
-                fragments=resolved_fragments,
-                attachments=resolved_attachments,
-                system=system,
-                schema=schema,
-                system_fragments=resolved_system_fragments,
-                **kwargs,
+                should_stream,
+                render,
+                extract,
+                extract_last,
+                resolved_fragments,
+                resolved_attachments,
+                system,
+                schema,
+                resolved_system_fragments,
+                kwargs,
             )
-            if should_stream:
-                for chunk in response:
-                    print(chunk, end="")
-                    sys.stdout.flush()
-                print("")
-            else:
-                text = response.text()
-                if extract or extract_last:
-                    text = extract_fenced_code_block(text, last=extract_last) or text
-                print(text)
     # List of exceptions that should never be raised in pytest:
     except (ValueError, NotImplementedError) as ex:
         raise click.ClickException(str(ex))

llm/cli_utils.py

@@ -0,0 +1,56 @@
+from rich.console import Console
+from rich.live import Live
+from rich.markdown import Markdown
+
+from llm import AsyncResponse, Attachment
+
+from .utils import Fragment, extract_fenced_code_block
+
+
+def prompt_output(
+    prompt_method: AsyncResponse,
+    prompt: str,
+    should_stream: bool,
+    render: bool,
+    extract: bool,
+    extract_last: bool,
+    resolved_fragments: list[Fragment],
+    resolved_attachments: list[Attachment],
+    system: str,
+    schema: str,
+    resolved_system_fragments: list[Fragment],
+    kwargs: dict,
+):
+    response = prompt_method(
+        prompt,
+        fragments=resolved_fragments,
+        attachments=resolved_attachments,
+        system=system,
+        schema=schema,
+        system_fragments=resolved_system_fragments,
+        **kwargs,
+    )
+
+    console = Console()
+
+    if should_stream:
+        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
+        if render:
+            text = Markdown(text)
+        console.print(text)
+
+    return response

pyproject.toml

@@ -39,6 +39,7 @@ dependencies = [
     "pip",
     "pyreadline3; sys_platform == 'win32'",
     "puremagic",
+    "rich"
 ]
 
 [project.urls]