diff --git a/mentat/resources/textual/terminal_app.tcss b/mentat/resources/textual/terminal_app.tcss index 330dc80a7..cff1f0499 100644 --- a/mentat/resources/textual/terminal_app.tcss +++ b/mentat/resources/textual/terminal_app.tcss @@ -43,4 +43,16 @@ ContextContainer > Tree { #loading-display { align-horizontal: center; width: 100%; +} + +ContentContainer > RichLog { + height: auto; + overflow: hidden; +} + +.content-piece { + margin: 0; + padding: 0; + height: auto; + width: 100%; } \ No newline at end of file diff --git a/mentat/session.py b/mentat/session.py index 62cc61eac..607c8bb25 100644 --- a/mentat/session.py +++ b/mentat/session.py @@ -9,7 +9,12 @@ import attr import sentry_sdk -from openai import APITimeoutError, BadRequestError, RateLimitError +from openai import ( + APITimeoutError, + BadRequestError, + PermissionDeniedError, + RateLimitError, +) from mentat.agent_handler import AgentHandler from mentat.auto_completer import AutoCompleter @@ -211,7 +216,12 @@ async def _main(self): except ReturnToUser: need_user_request = True continue - except (APITimeoutError, RateLimitError, BadRequestError) as e: + except ( + APITimeoutError, + RateLimitError, + BadRequestError, + PermissionDeniedError, + ) as e: stream.send(f"Error accessing OpenAI API: {e.message}", style="error") break diff --git a/mentat/terminal/patched_autocomplete.py b/mentat/terminal/patched_autocomplete.py index b0f49f5a6..937ad71c0 100644 --- a/mentat/terminal/patched_autocomplete.py +++ b/mentat/terminal/patched_autocomplete.py @@ -166,7 +166,7 @@ def reposition( return if input_cursor_position is None: - input_cursor_position = self.input_widget.cursor_position + input_cursor_position = self.input_widget.cursor_position # pyright: ignore top, right, bottom, left = self.styles.margin # pyright: ignore x, y, width, height = self.input_widget.content_region # pyright: ignore diff --git a/mentat/terminal/terminal_app.py b/mentat/terminal/terminal_app.py index b65687705..4324dd0d8 100644 --- a/mentat/terminal/terminal_app.py +++ b/mentat/terminal/terminal_app.py @@ -9,9 +9,7 @@ from rich.markup import escape from textual import on from textual.app import App, ComposeResult -from textual.message import Message -from textual.reactive import reactive -from textual.widgets import Input, ProgressBar, Static, Tree +from textual.widgets import Input, ProgressBar, RichLog, Static, Tree from textual.widgets._tree import TreeNode from typing_extensions import override @@ -28,29 +26,6 @@ history_file_location = mentat_dir_path / "prompt_history" -class ContentDisplay(Static): - content = reactive("") - - def add_content(self, new_content: str, color: Optional[str] = None): - """ - Using any sort of outside text renderer / colorer (like termcolor.colored or pygments.formatted) - will lead to strange artifacts appearing on the ContentDisplay. Do *NOT* use anything other than - the SessionStream's style/color kwargs! - """ - - new_content = escape(new_content) - if color is not None and color: - new_content = f"[{color}]{new_content}[/{color}]" - self.content += new_content - - def watch_content(self, content: str): - self.update(content) - self.post_message(self.ContentAdded()) - - class ContentAdded(Message): - pass - - class ContentContainer(Static): BINDINGS = [ ("up", "history_up", "Move up in history"), @@ -70,12 +45,17 @@ def __init__( self.last_user_input = "" self.suggester = HistorySuggester(history_file=history_file_location) self.loading_bar = None + self.cur_line = "" super().__init__(renderable, **kwargs) @override def compose(self) -> ComposeResult: - yield ContentDisplay() + self.content = RichLog(wrap=True, markup=True, auto_scroll=True) + yield self.content + # RichLogs can't edit existing lines, only add new ones, so we have a 'buffer' widget until we reach a new line. + self.last_content = Static(self.cur_line, classes="content-piece") + yield self.last_content yield PatchedAutoComplete( Input( classes="user-input", @@ -91,10 +71,6 @@ def on_user_input(self, event: Input.Submitted): self.suggester.append_to_history(event.value) self.input_event.set() - @on(ContentDisplay.ContentAdded) - def on_content_added(self, event: ContentDisplay.ContentAdded): - self.scroll_end(animate=False) - async def collect_user_input(self, default_prompt: str) -> str: self.input_event.clear() @@ -106,10 +82,7 @@ async def collect_user_input(self, default_prompt: str) -> str: user_input.value = "" user_input.disabled = True - content_display = self.query_one(ContentDisplay) - content_display.add_content( - f">>> {self.last_user_input}\n", color=self.theme["prompt"] - ) + self.add_content(f">>> {self.last_user_input}\n", color=self.theme["prompt"]) return self.last_user_input def action_history_up(self): @@ -138,6 +111,20 @@ def end_loading(self): self.loading_bar.remove() self.loading_bar = None + def add_content(self, new_content: str, color: str | None): + new_content = escape(new_content) + lines = [ + f"[{color}]{line}[/{color}]" if color else line + for line in new_content.split("\n") + ] + for line in lines[:-1]: + line = self.cur_line + line + self.cur_line = "" + self.content.write(line) + self.cur_line += lines[-1] + self.last_content.update(self.cur_line) + self.scroll_end(animate=False) + class ContextContainer(Static): def _build_path_tree(self, files: list[str], cwd: Path): @@ -258,9 +245,9 @@ def display_stream_message(self, message: StreamMessage): style = message.extra["style"] color = self.theme[style] - content_display = self.query_one(ContentDisplay) + content_container = self.query_one(ContentContainer) content = str(message.data) + end - content_display.add_content(content, color) + content_container.add_content(content, color) async def get_user_input( self, default_prompt: str, command_autocomplete: bool @@ -299,6 +286,7 @@ def action_on_interrupt(self): def disable_app(self): self.query_one(Input).disabled = True + self.end_loading() def start_loading(self): self.query_one(ContentContainer).start_loading()