fix

Author: inverse

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,
@@ -41,7 +46,6 @@
     remove_alias,
 )
 from llm.models import _BaseConversation, ChainResponse
-from .cli_utils import prompt_output
 
 from .migrations import migrate
 from .plugins import pm, load_plugins
@@ -833,41 +837,76 @@ def read_prompt():
         # Merge in options for the .prompt() methods
         kwargs.update(validated_options)
 
+    console = Console()
+
     try:
         if async_:
 
             async def inner():
-                return prompt_output(
-                    prompt_method,
+                response = prompt_method(
                     prompt,
-                    should_stream,
-                    render,
-                    extract,
-                    extract_last,
-                    resolved_fragments,
-                    resolved_attachments,
-                    system,
-                    schema,
-                    resolved_system_fragments,
-                    kwargs,
+                    attachments=resolved_attachments,
+                    system=system,
+                    schema=schema,
+                    fragments=resolved_fragments,
+                    system_fragments=resolved_system_fragments,
+                    **kwargs,
                 )
 
+                if should_stream:
+                    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:
+                    text = await 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)
+
+
             response = asyncio.run(inner())
         else:
-            response = prompt_output(
-                prompt_method,
+            response = prompt_method(
                 prompt,
-                should_stream,
-                render,
-                extract,
-                extract_last,
-                resolved_fragments,
-                resolved_attachments,
-                system,
-                schema,
-                resolved_system_fragments,
-                kwargs,
+                fragments=resolved_fragments,
+                attachments=resolved_attachments,
+                system=system,
+                schema=schema,
+                system_fragments=resolved_system_fragments,
+                **kwargs,
             )
+
+
+
+            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)
     # List of exceptions that should never be raised in pytest:
     except (ValueError, NotImplementedError) as ex:
         raise click.ClickException(str(ex))