From b143bc56ac773c54bcbffeab4534a3a0b347cb7f Mon Sep 17 00:00:00 2001
From: Paul Gauthier <aider@paulg.org>
Date: Tue, 23 Jan 2024 09:20:31 -0800
Subject: [PATCH] switch to mdstream

---
 aider/coders/base_coder.py | 23 ++++++++++-------------
 aider/mdstream.py          | 12 ++++++++----
 2 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/aider/coders/base_coder.py b/aider/coders/base_coder.py
index e728e2478..a0c764dfa 100755
--- a/aider/coders/base_coder.py
+++ b/aider/coders/base_coder.py
@@ -13,13 +13,13 @@ from pathlib import Path
 import openai
 from jsonschema import Draft7Validator
 from rich.console import Console, Text
-from rich.live import Live
 from rich.markdown import Markdown
 
 from aider import models, prompts, utils
 from aider.commands import Commands
 from aider.history import ChatSummary
 from aider.io import InputOutput
+from aider.mdstream import MarkdownStream
 from aider.repo import GitRepo
 from aider.repomap import RepoMap
 from aider.sendchat import send_with_retries
@@ -725,14 +725,12 @@ class Coder:
             self.io.tool_output(tokens)
 
     def show_send_output_stream(self, completion):
-        live = None
         if self.show_pretty():
-            live = Live(vertical_overflow="scroll")
+            mdstream = MarkdownStream()
+        else:
+            mdstream = None
 
         try:
-            if live:
-                live.start()
-
             for chunk in completion:
                 if len(chunk.choices) == 0:
                     continue
@@ -762,22 +760,21 @@ class Coder:
                     text = None
 
                 if self.show_pretty():
-                    self.live_incremental_response(live, False)
+                    self.live_incremental_response(mdstream, False)
                 elif text:
                     sys.stdout.write(text)
                     sys.stdout.flush()
         finally:
-            if live:
-                self.live_incremental_response(live, True)
-                live.stop()
+            if mdstream:
+                self.live_incremental_response(mdstream, True)
 
-    def live_incremental_response(self, live, final):
+    def live_incremental_response(self, mdstream, final):
         show_resp = self.render_incremental_response(final)
         if not show_resp:
             return
 
-        md = Markdown(show_resp, style=self.assistant_output_color, code_theme=self.code_theme)
-        live.update(md)
+        mdargs = dict(style=self.assistant_output_color, code_theme=self.code_theme)
+        mdstream.update(show_resp, mdargs=mdargs, final=final)
 
     def render_incremental_response(self, final):
         return self.partial_response_content
diff --git a/aider/mdstream.py b/aider/mdstream.py
index bd3ce1292..3ecb5079f 100755
--- a/aider/mdstream.py
+++ b/aider/mdstream.py
@@ -60,7 +60,7 @@ class MarkdownStream:
         self.printed = []
         self.when = 0
 
-    def update(self, text, final=False, min_delay=0.100):
+    def update(self, text, final=False, min_delay=0.100, mdargs=None):
         now = time.time()
         if not final and now - self.when < min_delay:
             return
@@ -69,7 +69,11 @@ class MarkdownStream:
         string_io = io.StringIO()
         console = Console(file=string_io, force_terminal=True)
 
-        markdown = Markdown(text)
+        if not mdargs:
+            mdargs = dict()
+
+        markdown = Markdown(text, **mdargs)
+
         console.print(markdown)
         output = string_io.getvalue()
 
@@ -78,7 +82,7 @@ class MarkdownStream:
         if not final:
             num_lines -= 4
 
-        if num_lines <= 1:
+        if num_lines <= 1 and not final:
             return
 
         num_printed = len(self.printed)
@@ -111,6 +115,6 @@ if __name__ == "__main__":
     pm = MarkdownStream()
     for i in range(6, len(_text)):
         pm.update(_text[:i])
-        # time.sleep(0.001)
+        time.sleep(0.001)
 
     pm.update(_text, final=True)