Markdown formatting for chat history (#444)

Author: jeanlucthumm

README.md

@@ -365,6 +365,8 @@ sgpt --role json_generator "random: user, password, email, address"
 }
 ```
 
+If the description of the role contains the words "APPLY MARKDOWN" (case sensitive), then chats will be displayed using markdown formatting.
+
 ### Request cache
 Control cache using `--cache` (default) and `--no-cache` options. This caching applies for all `sgpt` requests to OpenAI API:
 ```shell

sgpt/handlers/chat_handler.py

@@ -4,6 +4,8 @@
 
 import typer
 from click import BadArgumentUsage
+from rich.console import Console
+from rich.markdown import Markdown
 
 from ..config import cfg
 from ..role import DefaultRoles, SystemRole
@@ -107,15 +109,15 @@ def __init__(self, chat_id: str, role: SystemRole) -> None:
     def initiated(self) -> bool:
         return self.chat_session.exists(self.chat_id)
 
-    @property
-    def initial_message(self) -> str:
-        chat_history = self.chat_session.get_messages(self.chat_id)
-        return chat_history[0] if chat_history else ""
-
     @property
     def is_same_role(self) -> bool:
         # TODO: Should be optimized for REPL mode.
-        return self.role.same_role(self.initial_message)
+        return self.role.same_role(self.initial_message(self.chat_id))
+
+    @classmethod
+    def initial_message(cls, chat_id: str) -> str:
+        chat_history = cls.chat_session.get_messages(chat_id)
+        return chat_history[0] if chat_history else ""
 
     @classmethod
     @option_callback
@@ -126,7 +128,15 @@ def list_ids(cls, value: str) -> None:
 
     @classmethod
     def show_messages(cls, chat_id: str) -> None:
-        # Prints all messages from a specified chat ID to the console.
+        if "APPLY MARKDOWN" in cls.initial_message(chat_id):
+            for message in cls.chat_session.get_messages(chat_id):
+                if message.startswith("assistant:"):
+                    Console().print(Markdown(message))
+                else:
+                    typer.secho(message, fg=cfg.get("DEFAULT_COLOR"))
+                typer.echo()
+            return
+
         for index, message in enumerate(cls.chat_session.get_messages(chat_id)):
             color = "magenta" if index % 2 == 0 else "green"
             typer.secho(message, fg=color)
@@ -138,7 +148,7 @@ def show_messages_callback(cls, chat_id: str) -> None:
 
     def validate(self) -> None:
         if self.initiated:
-            chat_role_name = self.role.get_role_name(self.initial_message)
+            chat_role_name = self.role.get_role_name(self.initial_message(self.chat_id))
             if not chat_role_name:
                 raise BadArgumentUsage(
                     f'Could not determine chat role of "{self.chat_id}"'

tests/test_default.py

@@ -38,6 +38,43 @@ def test_default_stdin(completion):
     assert "Prague" in result.stdout
 
 
+@patch("rich.console.Console.print")
+@patch("litellm.completion")
+def test_show_chat_use_markdown(completion, console_print):
+    completion.return_value = mock_comp("ok")
+    chat_name = "_test"
+    chat_path = Path(cfg.get("CHAT_CACHE_PATH")) / chat_name
+    chat_path.unlink(missing_ok=True)
+
+    args = {"prompt": "my number is 2", "--chat": chat_name}
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 0
+    assert chat_path.exists()
+
+    result = runner.invoke(app, ["--show-chat", chat_name])
+    assert result.exit_code == 0
+    console_print.assert_called()
+
+
+@patch("rich.console.Console.print")
+@patch("litellm.completion")
+def test_show_chat_no_use_markdown(completion, console_print):
+    completion.return_value = mock_comp("ok")
+    chat_name = "_test"
+    chat_path = Path(cfg.get("CHAT_CACHE_PATH")) / chat_name
+    chat_path.unlink(missing_ok=True)
+
+    # Flag '--code' doesn't use markdown
+    args = {"prompt": "my number is 2", "--chat": chat_name, "--code": True}
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 0
+    assert chat_path.exists()
+
+    result = runner.invoke(app, ["--show-chat", chat_name])
+    assert result.exit_code == 0
+    console_print.assert_not_called()
+
+
 @patch("litellm.completion")
 def test_default_chat(completion):
     completion.side_effect = [mock_comp("ok"), mock_comp("4")]