Render markdown for default output (TheR1D#400)

Co-authored-by: Jaycen Horton <jaycen@nori.com>

README.md

@@ -264,38 +264,13 @@ import requests
 response = requests.get('https://localhost:443')
 print(response.text)
 ```
-
-### Picking up on a chat mode conversation with REPL mode
-
-```text
-sgpt --repl number
-───── Chat History──────
-user: ###
-Role name: default
-You are Command Line App ShellGPT, a programming and system administration assistant.
-You are managing Darwin/MacOS 13.3.1 operating system with zsh shell.
-Provide only plain text without Markdown formatting.
-Do not show any warnings or information regarding your capabilities.
-If you need to store any data, assume it will be stored in the chat.
-
-Request: please remember my favorite number: 4
-###
-assistant: Sure, I have stored your favorite number as 4.
-user: what would be my favorite number raised to the power of 4
-assistant: Your favorite number raised to the power of 4 would be 256.
-────────────────────────────────────────────────────────
-Entering REPL mode, press Ctrl+C to exit.
->>> What is the sum of my favorite number and your previous response?
-The sum of your favorite number (4) and my previous response (256) would be 260.
-```
-
+It is also possible to pickup conversations from chat sessions (which were created using `--chat` option) and continue them in REPL mode.
 
 ### Roles
 ShellGPT allows you to create custom roles, which can be utilized to generate code, shell commands, or to fulfill your specific needs. To create a new role, use the `--create-role` option followed by the role name. You will be prompted to provide a description for the role, along with other details. This will create a JSON file in `~/.config/shell_gpt/roles` with the role name. Inside this directory, you can also edit default `sgpt` roles, such as **shell**, **code**, and **default**. Use the `--list-roles` option to list all available roles, and the `--show-role` option to display the details of a specific role. Here's an example of a custom role:
 ```shell
-sgpt --create-role json
-# Enter role description: You are JSON generator, provide only valid json as response.
-# Enter expecting result, e.g. answer, code, shell command, etc.: json
+sgpt --create-role json_generator
+# Enter role description: Provide only valid json as response.
 sgpt --role json "random: user, password, email, address"
 {
   "user": "JohnDoe",
@@ -339,14 +314,17 @@ CACHE_PATH=/tmp/shell_gpt/cache
 REQUEST_TIMEOUT=60
 # Default OpenAI model to use.
 DEFAULT_MODEL=gpt-3.5-turbo
-# Default color for OpenAI completions.
+# Default color for shell and code completions.
 DEFAULT_COLOR=magenta
 # When in --shell mode, default to "Y" for no input.
 DEFAULT_EXECUTE_SHELL_CMD=false
 # Disable streaming of responses
 DISABLE_STREAMING=false
+# The pygment theme to view markdown (default/describe role).
+CODE_THEME=default
 ```
 Possible options for `DEFAULT_COLOR`: black, red, green, yellow, blue, magenta, cyan, white, bright_black, bright_red, bright_green, bright_yellow, bright_blue, bright_magenta, bright_cyan, bright_white.
+Possible options for `CODE_THEME`: https://pygments.org/styles/
 
 ### Full list of arguments
 ```text
@@ -355,8 +333,8 @@ Possible options for `DEFAULT_COLOR`: black, red, green, yellow, blue, magenta,
 ╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
 ╭─ Options ───────────────────────────────────────────────────────────────────────────────────────────────────╮
 │ --model            TEXT                             OpenAI GPT model to use. [default: gpt-3.5-turbo]       │
-│ --temperature      FLOAT RANGE [0.0<=x<=2.0]        Randomness of generated output. [default: 0.1]          │
-│ --top-probability  FLOAT RANGE [0.1<=x<=1.0]        Limits highest probable tokens (words). [default: 1.0]  │
+│ --temperature      FLOAT RANGE [0.0<=x<=2.0]        Randomness of generated output. [default: 0.0]          │
+│ --top-probability  FLOAT RANGE [0.0<=x<=1.0]        Limits highest probable tokens (words). [default: 1.0]  │
 │ --editor                                            Open $EDITOR to provide a prompt. [default: no-editor]  │
 │ --cache                                             Cache completion results. [default: cache]              │
 │ --help                                              Show this message and exit.                             │

sgpt/config.py

@@ -26,7 +26,8 @@
     "DEFAULT_COLOR": os.getenv("DEFAULT_COLOR", "magenta"),
     "ROLE_STORAGE_PATH": os.getenv("ROLE_STORAGE_PATH", str(ROLE_STORAGE_PATH)),
     "DEFAULT_EXECUTE_SHELL_CMD": os.getenv("DEFAULT_EXECUTE_SHELL_CMD", "false"),
-    "DISABLE_STREAMING": os.getenv("DISABLE_STREAMING", "false")
+    "DISABLE_STREAMING": os.getenv("DISABLE_STREAMING", "false"),
+    "CODE_THEME": os.getenv("CODE_THEME", "dracula"),
     # New features might add their own config variables here.
 }
 

sgpt/handlers/handler.py

@@ -1,6 +1,9 @@
 from typing import Any, Dict, Generator, List
 
 import typer
+from rich.console import Console
+from rich.live import Live
+from rich.markdown import Markdown
 
 from ..client import OpenAIClient
 from ..config import cfg
@@ -13,22 +16,50 @@ def __init__(self, role: SystemRole) -> None:
             cfg.get("OPENAI_API_HOST"), cfg.get("OPENAI_API_KEY")
         )
         self.role = role
+        self.disable_stream = cfg.get("DISABLE_STREAMING") == "false"
         self.color = cfg.get("DEFAULT_COLOR")
+        self.theme_name = cfg.get("CODE_THEME")
 
-    def make_messages(self, prompt: str) -> List[Dict[str, str]]:
-        raise NotImplementedError
-
-    def get_completion(self, **kwargs: Any) -> Generator[str, None, None]:
-        yield from self.client.get_completion(**kwargs)
+    def _handle_with_markdown(self, prompt: str, **kwargs: Any) -> str:
+        messages = self.make_messages(prompt.strip())
+        full_completion = ""
+        with Live(
+            Markdown(markup="", code_theme=self.theme_name),
+            console=Console(),
+            refresh_per_second=8,
+        ) as live:
+            if not self.disable_stream:
+                live.update(
+                    Markdown(markup="Loading...\r", code_theme=self.theme_name),
+                    refresh=True,
+                )
+            for word in self.get_completion(messages=messages, **kwargs):
+                full_completion += word
+                live.update(
+                    Markdown(full_completion, code_theme=self.theme_name),
+                    refresh=not self.disable_stream,
+                )
+        return full_completion
 
-    def handle(self, prompt: str, **kwargs: Any) -> str:
+    def _handle_with_plain_text(self, prompt: str, **kwargs: Any) -> str:
         messages = self.make_messages(prompt.strip())
         full_completion = ""
-        stream = cfg.get("DISABLE_STREAMING") == "false"
-        if not stream:
+        if not self.disable_stream:
             typer.echo("Loading...\r", nl=False)
         for word in self.get_completion(messages=messages, **kwargs):
             typer.secho(word, fg=self.color, bold=True, nl=False)
             full_completion += word
-        typer.echo("\033[K" if not stream else "")  # Overwrite "loading..."
+        # Overwrite "loading..."
+        typer.echo("\033[K" if not self.disable_stream else "")
         return full_completion
+
+    def make_messages(self, prompt: str) -> List[Dict[str, str]]:
+        raise NotImplementedError
+
+    def get_completion(self, **kwargs: Any) -> Generator[str, None, None]:
+        yield from self.client.get_completion(**kwargs)
+
+    def handle(self, prompt: str, **kwargs: Any) -> str:
+        if self.role.name == "ShellGPT" or self.role.name == "Shell Command Descriptor":
+            return self._handle_with_markdown(prompt, **kwargs)
+        return self._handle_with_plain_text(prompt, **kwargs)

sgpt/role.py

@@ -22,9 +22,9 @@
 """
 
 DESCRIBE_SHELL_ROLE = """Provide a terse, single sentence description of the given shell command.
-Describe each argument and option of the command, in plain text (no markdown).
-Provide only plain text without Markdown formatting.
-Do not include symbols such as `."""
+Describe each argument and option of the command.
+Use markdown formatting when possible.
+"""
 
 CODE_ROLE = """Provide only code as output without any description.
 Provide only code in plain text format without Markdown formatting.
@@ -35,7 +35,7 @@
 
 DEFAULT_ROLE = """You are programming and system administration assistant.
 You are managing {os} operating system with {shell} shell.
-Provide only plain text without Markdown formatting.
+Use Markdown formatting when possible.
 If you need to store any data, assume it will be stored in the conversation."""
 
 ROLE_TEMPLATE = "You are {name}\n{role}"

tests/test_integration.py

@@ -418,14 +418,14 @@ def test_shell_stdin_with_prompt(self):
         assert "sort" in result.stdout
 
     def test_role(self):
-        test_role = Path(cfg.get("ROLE_STORAGE_PATH")) / "test_json.json"
+        test_role = Path(cfg.get("ROLE_STORAGE_PATH")) / "json_generator.json"
         test_role.unlink(missing_ok=True)
         dict_arguments = {
             "prompt": "test",
-            "--create-role": "test_json",
+            "--create-role": "json_generator",
         }
         input = (
-            "You are a JSON generator, return only valid plain JSON as response. "
+            "Provide only valid plain JSON as response with valid field values. "
             + "Do not include any markdown formatting such as ```.\n"
         )
         result = runner.invoke(app, self.get_arguments(**dict_arguments), input=input)
@@ -437,20 +437,20 @@ def test_role(self):
         }
         result = runner.invoke(app, self.get_arguments(**dict_arguments))
         assert result.exit_code == 0
-        assert "test_json" in result.stdout
+        assert "json_generator" in result.stdout
 
         dict_arguments = {
             "prompt": "test",
-            "--show-role": "test_json",
+            "--show-role": "json_generator",
         }
         result = runner.invoke(app, self.get_arguments(**dict_arguments))
         assert result.exit_code == 0
-        assert "You are a JSON generator" in result.stdout
+        assert "You are json_generator" in result.stdout
 
         # Test with command line argument prompt.
         dict_arguments = {
             "prompt": "random username, password, email",
-            "--role": "test_json",
+            "--role": "json_generator",
         }
         result = runner.invoke(app, self.get_arguments(**dict_arguments))
         assert result.exit_code == 0
@@ -462,7 +462,7 @@ def test_role(self):
         # Test with stdin prompt.
         dict_arguments = {
             "prompt": "",
-            "--role": "test_json",
+            "--role": "json_generator",
         }
         stdin = "random username, password, email"
         result = runner.invoke(app, self.get_arguments(**dict_arguments), input=stdin)