Added a -r, --raw option to use the LLM as-is without a specific role or system message.

Raw mode increases flexibility by enabling responses similar to those obtained directly through the LLM's API.

Author: Gilad Barnea

README.md

@@ -436,6 +436,7 @@ Possible options for `CODE_THEME`: https://pygments.org/styles/
 │ --interaction         --no-interaction    Interactive mode for --shell option. [default: interaction]    │
 │ --describe-shell  -d                      Describe a shell command.                                      │
 │ --code            -c                      Generate only code.                                            │
+│ --raw             -r                      Use the LLM as-is without a specific role or system message.   │
 │ --functions           --no-functions      Allow function calls. [default: functions]                     │
 ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────╯
 ╭─ Chat Options ───────────────────────────────────────────────────────────────────────────────────────────╮

sgpt/app.py

@@ -75,6 +75,13 @@ def main(
         help="Generate only code.",
         rich_help_panel="Assistance Options",
     ),
+    raw: bool = typer.Option(
+        False,
+        "-r",
+        "--raw",
+        help="Use the LLM as-is without a specific role or system message.",
+        rich_help_panel="Assistance Options",
+    ),
     functions: bool = typer.Option(
         cfg.get("OPENAI_USE_FUNCTIONS") == "true",
         help="Allow function calls.",
@@ -183,9 +190,9 @@ def main(
             # Non-interactive shell.
             pass
 
-    if sum((shell, describe_shell, code)) > 1:
+    if sum((shell, describe_shell, code, raw)) > 1:
         raise BadArgumentUsage(
-            "Only one of --shell, --describe-shell, and --code options can be used at a time."
+            "Only one of --shell, --describe-shell, --code and --raw options can be used at a time."
         )
 
     if chat and repl:
@@ -198,7 +205,7 @@ def main(
         prompt = get_edited_prompt()
 
     role_class = (
-        DefaultRoles.check_get(shell, describe_shell, code)
+        DefaultRoles.check_get(shell, describe_shell, code, raw)
         if not role
         else SystemRole.get(role)
     )

sgpt/role.py

@@ -39,6 +39,8 @@
 Provide short responses in about 100 words, unless you are specifically asked for more details.
 If you need to store any data, assume it will be stored in the conversation.
 APPLY MARKDOWN formatting when possible."""
+
+RAW_ROLE = """APPLY MARKDOWN formatting when possible."""
 # Note that output for all roles containing "APPLY MARKDOWN" will be formatted as Markdown.
 
 ROLE_TEMPLATE = "You are {name}\n{role}"
@@ -68,6 +70,7 @@ def create_defaults(cls) -> None:
             SystemRole("Shell Command Generator", SHELL_ROLE, variables),
             SystemRole("Shell Command Descriptor", DESCRIBE_SHELL_ROLE, variables),
             SystemRole("Code Generator", CODE_ROLE),
+            SystemRole("GPT", RAW_ROLE),
         ):
             if not default_role._exists:
                 default_role._save()
@@ -167,15 +170,20 @@ class DefaultRoles(Enum):
     SHELL = "Shell Command Generator"
     DESCRIBE_SHELL = "Shell Command Descriptor"
     CODE = "Code Generator"
+    RAW = "GPT"
 
     @classmethod
-    def check_get(cls, shell: bool, describe_shell: bool, code: bool) -> SystemRole:
+    def check_get(
+        cls, shell: bool, describe_shell: bool, code: bool, raw: bool
+    ) -> SystemRole:
         if shell:
             return SystemRole.get(DefaultRoles.SHELL.value)
         if describe_shell:
             return SystemRole.get(DefaultRoles.DESCRIBE_SHELL.value)
         if code:
             return SystemRole.get(DefaultRoles.CODE.value)
+        if raw:
+            return SystemRole.get(DefaultRoles.RAW.value)
         return SystemRole.get(DefaultRoles.DEFAULT.value)
 
     def get_role(self) -> SystemRole:

tests/test_raw.py

@@ -0,0 +1,165 @@
+from pathlib import Path
+from unittest.mock import patch
+
+import typer
+from typer.testing import CliRunner
+
+from sgpt import config, main
+from sgpt.role import DefaultRoles, SystemRole
+
+from .utils import app, cmd_args, comp_args, mock_comp, runner
+
+role = SystemRole.get(DefaultRoles.RAW.value)
+cfg = config.cfg
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_raw_long_option(completion):
+    completion.return_value = mock_comp("Prague")
+
+    args = {"prompt": "capital of the Czech Republic?", "--raw": True}
+    result = runner.invoke(app, cmd_args(**args))
+
+    completion.assert_called_once_with(**comp_args(role, **args))
+    assert result.exit_code == 0
+    assert "Prague" in result.stdout
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_raw_short_option(completion):
+    completion.return_value = mock_comp("Prague")
+
+    args = {"prompt": "capital of the Czech Republic?", "-r": True}
+    result = runner.invoke(app, cmd_args(**args))
+
+    completion.assert_called_once_with(**comp_args(role, **args))
+    assert result.exit_code == 0
+    assert "Prague" in result.stdout
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_raw_stdin(completion):
+    completion.return_value = mock_comp("Prague")
+
+    args = {"--raw": True}
+    stdin = "capital of the Czech Republic?"
+    result = runner.invoke(app, cmd_args(**args), input=stdin)
+
+    completion.assert_called_once_with(**comp_args(role, stdin))
+    assert result.exit_code == 0
+    assert "Prague" in result.stdout
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_raw_chat(completion):
+    completion.side_effect = [mock_comp("ok"), mock_comp("4")]
+    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", "--raw": True, "--chat": chat_name}
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 0
+    assert "ok" in result.stdout
+    assert chat_path.exists()
+
+    args["prompt"] = "my number + 2?"
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 0
+    assert "4" in result.stdout
+
+    expected_messages = [
+        {"role": "system", "content": role.role},
+        {"role": "user", "content": "my number is 2"},
+        {"role": "assistant", "content": "ok"},
+        {"role": "user", "content": "my number + 2?"},
+        {"role": "assistant", "content": "4"},
+    ]
+    expected_args = comp_args(role, "", messages=expected_messages)
+    completion.assert_called_with(**expected_args)
+    assert completion.call_count == 2
+
+    result = runner.invoke(app, ["--list-chats"])
+    assert result.exit_code == 0
+    assert "_test" in result.stdout
+
+    result = runner.invoke(app, ["--show-chat", chat_name])
+    assert result.exit_code == 0
+    assert "my number is 2" in result.stdout
+    assert "ok" in result.stdout
+    assert "my number + 2?" in result.stdout
+    assert "4" in result.stdout
+
+    args["--shell"] = True
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 2
+    assert "Error" in result.stdout
+
+    args["--code"] = True
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 2
+    assert "Error" in result.stdout
+    chat_path.unlink()
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_raw_repl(completion):
+    completion.side_effect = [mock_comp("ok"), mock_comp("8")]
+    chat_name = "_test"
+    chat_path = Path(cfg.get("CHAT_CACHE_PATH")) / chat_name
+    chat_path.unlink(missing_ok=True)
+
+    args = {"--raw": True, "--repl": chat_name}
+    inputs = ["__sgpt__eof__", "my number is 6", "my number + 2?", "exit()"]
+    result = runner.invoke(app, cmd_args(**args), input="\n".join(inputs))
+
+    expected_messages = [
+        {"role": "system", "content": role.role},
+        {"role": "user", "content": "my number is 6"},
+        {"role": "assistant", "content": "ok"},
+        {"role": "user", "content": "my number + 2?"},
+        {"role": "assistant", "content": "8"},
+    ]
+    expected_args = comp_args(role, "", messages=expected_messages)
+    completion.assert_called_with(**expected_args)
+    assert completion.call_count == 2
+
+    assert result.exit_code == 0
+    assert ">>> my number is 6" in result.stdout
+    assert "ok" in result.stdout
+    assert ">>> my number + 2?" in result.stdout
+    assert "8" in result.stdout
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_raw_repl_stdin(completion):
+    completion.side_effect = [mock_comp("ok init"), mock_comp("ok another")]
+    chat_name = "_test"
+    chat_path = Path(cfg.get("CHAT_CACHE_PATH")) / chat_name
+    chat_path.unlink(missing_ok=True)
+
+    my_runner = CliRunner()
+    my_app = typer.Typer()
+    my_app.command()(main)
+
+    args = {"--raw": True, "--repl": chat_name}
+    inputs = ["this is stdin", "__sgpt__eof__", "prompt", "another", "exit()"]
+    result = my_runner.invoke(my_app, cmd_args(**args), input="\n".join(inputs))
+
+    expected_messages = [
+        {"role": "system", "content": role.role},
+        {"role": "user", "content": "this is stdin\n\n\n\nprompt"},
+        {"role": "assistant", "content": "ok init"},
+        {"role": "user", "content": "another"},
+        {"role": "assistant", "content": "ok another"},
+    ]
+    expected_args = comp_args(role, "", messages=expected_messages)
+    completion.assert_called_with(**expected_args)
+    assert completion.call_count == 2
+
+    assert result.exit_code == 0
+    assert "this is stdin" in result.stdout
+    assert ">>> prompt" in result.stdout
+    assert "ok init" in result.stdout
+    assert ">>> another" in result.stdout
+    assert "ok another" in result.stdout