Skip to content

Commit

Permalink
Added a -r, --raw option to use the LLM as-is without a specific role…
Browse files Browse the repository at this point in the history
… or system message.

Raw mode increases flexibility by enabling responses similar to those obtained directly through the LLM's API.
  • Loading branch information
Gilad Barnea committed May 2, 2024
1 parent 20405b2 commit faba3d6
Show file tree
Hide file tree
Showing 4 changed files with 185 additions and 4 deletions.
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 ───────────────────────────────────────────────────────────────────────────────────────────╮
Expand Down
13 changes: 10 additions & 3 deletions sgpt/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.",
Expand Down Expand Up @@ -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:
Expand All @@ -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)
)
Expand Down
10 changes: 9 additions & 1 deletion sgpt/role.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}"
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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:
Expand Down
165 changes: 165 additions & 0 deletions tests/test_raw.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit faba3d6

Please sign in to comment.