Custom roles (TheR1D#183)

* Implementing custom roles.
* Implementing system role messages.
* Small stdin change with new lines at the end.
* Small option name fix --list-chats.
* Refactoring module imports.
* Adding tests for roles.

README.md

@@ -1,11 +1,11 @@
-# Shell GPT
-A command-line productivity tool powered by OpenAI's GPT-3.5 model. As developers, we can leverage AI capabilities to generate shell commands, code snippets, comments, and documentation, among other things. Forget about cheat sheets and notes, with this tool you can get accurate answers right in your terminal, and you'll probably find yourself reducing your daily Google searches, saving you valuable time and effort.
+# ShellGPT
+A command-line productivity tool powered by OpenAI's GPT models. As developers, we can leverage AI capabilities to generate shell commands, code snippets, comments, and documentation, among other things. Forget about cheat sheets and notes, with this tool you can get accurate answers right in your terminal, and you'll probably find yourself reducing your daily Google searches, saving you valuable time and effort. ShellGPT is cross-platform compatible and supports all major operating systems, including Linux, macOS, and Windows with all major shells, such as PowerShell, CMD, Bash, Zsh, Fish, and many others.
 
 https://user-images.githubusercontent.com/16740832/231569156-a3a9f9d4-18b1-4fff-a6e1-6807651aa894.mp4
 
 ## Installation
 ```shell
-pip install shell-gpt==0.8.9
+pip install shell-gpt==0.9.0
 ```
 You'll need an OpenAI API key, you can generate one [here](https://beta.openai.com/account/api-keys).
 
@@ -218,9 +218,9 @@ print(response.text)
 ```
 
 ### Chat sessions
-To list all the current chat sessions, use the `--list-chat` option:
+To list all the current chat sessions, use the `--list-chats` option:
 ```shell
-sgpt --list-chat
+sgpt --list-chats
 # .../shell_gpt/chat_cache/number
 # .../shell_gpt/chat_cache/python_request
 ```
@@ -233,6 +233,26 @@ sgpt --show-chat number
 # assistant: Your favorite number is 4, so if we add 4 to it, the result would be 8.
 ```
 
+### 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 --role json "random: user, password, email, address"
+{
+  "user": "JohnDoe",
+  "password": "p@ssw0rd",
+  "email": "[email protected]",
+  "address": {
+    "street": "123 Main St",
+    "city": "Anytown",
+    "state": "CA",
+    "zip": "12345"
+  }
+}
+```
+
 ### Request cache
 Control cache using `--cache` (default) and `--no-cache` options. This caching applies for all `sgpt` requests to OpenAI API:
 ```shell
@@ -264,32 +284,42 @@ REQUEST_TIMEOUT=60
 DEFAULT_MODEL=gpt-3.5-turbo
 # Default color for OpenAI completions.
 DEFAULT_COLOR=magenta
+# Force use system role messages (not recommended).
+SYSTEM_ROLES=false
 ```
 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.
 
+Switch `SYSTEM_ROLES` to force use [system roles](https://help.openai.com/en/articles/7042661-chatgpt-api-transition-guide) messages, this is not recommended, since it doesn't perform well with current GPT models.
+
 ### Full list of arguments
 ```text
-╭─ Arguments ────────────────────────────────────────────────────────────────────────────────────────────────╮
-│   prompt      [PROMPT]  The prompt to generate completions for.                                            │
-╰────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
-╭─ Options ──────────────────────────────────────────────────────────────────────────────────────────────────╮
-│ --model            [gpt-3.5-turbo|gpt-4|gpt-4-32k]  OpenAI GPT model to use. [default: gpt-3.5-turbo]      │
-│ --temperature      FLOAT RANGE [0.0<=x<=1.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] │
-│ --editor                                            Open $EDITOR to provide a prompt. [default: no-editor] │
-│ --cache                                             Cache completion results. [default: cache]             │
-│ --help                                              Show this message and exit.                            │
-╰────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
-╭─ Assistance Options ───────────────────────────────────────────────────────────────────────────────────────╮
-│ --shell  -s                 Generate and execute shell commands.                                           │
-│ --code       --no-code      Generate only code. [default: no-code]                                         │
-╰────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
-╭─ Chat Options ─────────────────────────────────────────────────────────────────────────────────────────────╮
-│ --chat        TEXT  Follow conversation with id, use "temp" for quick session. [default: None]             │
-│ --repl        TEXT  Start a REPL (Read–eval–print loop) session. [default: None]                           │
-│ --show-chat   TEXT  Show all messages from provided chat id. [default: None]                               │
-│ --list-chat         List all existing chat ids. [default: no-list-chat]                                    │
-╰────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
+╭─ Arguments ─────────────────────────────────────────────────────────────────────────────────────────────────╮
+│   prompt      [PROMPT]  The prompt to generate completions for.                                             │
+╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
+╭─ Options ───────────────────────────────────────────────────────────────────────────────────────────────────╮
+│ --model            [gpt-3.5-turbo|gpt-4|gpt-4-32k]  OpenAI GPT model to use. [default: gpt-3.5-turbo]       │
+│ --temperature      FLOAT RANGE [0.0<=x<=1.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]  │
+│ --editor                                            Open $EDITOR to provide a prompt. [default: no-editor]  │
+│ --cache                                             Cache completion results. [default: cache]              │
+│ --help                                              Show this message and exit.                             │
+╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
+╭─ Assistance Options ────────────────────────────────────────────────────────────────────────────────────────╮
+│ --shell  -s                 Generate and execute shell commands.                                            │
+│ --code       --no-code      Generate only code. [default: no-code]                                          │
+╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
+╭─ Chat Options ──────────────────────────────────────────────────────────────────────────────────────────────╮
+│ --chat        TEXT  Follow conversation with id, use "temp" for quick session. [default: None]              │
+│ --repl        TEXT  Start a REPL (Read–eval–print loop) session. [default: None]                            │
+│ --show-chat   TEXT  Show all messages from provided chat id. [default: None]                                │
+│ --list-chats        List all existing chat ids. [default: no-list-chats]                                    │
+╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
+╭─ Role Options ──────────────────────────────────────────────────────────────────────────────────────────────╮
+│ --role         TEXT  System role for GPT model. [default: None]                                             │
+│ --create-role  TEXT  Create role. [default: None]                                                           │
+│ --show-role    TEXT  Show role. [default: None]                                                             │
+│ --list-roles         List roles. [default: no-list-roles]                                                   │
+╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
 ```
 
 ## Docker

sgpt/__init__.py

@@ -1,12 +1,4 @@
-from .config import cfg as cfg
-from .cache import Cache as Cache
-from .client import OpenAIClient as OpenAIClient
-from .handlers.chat_handler import ChatHandler as ChatHandler
-from .handlers.default_handler import DefaultHandler as DefaultHandler
-from .handlers.repl_handler import ReplHandler as ReplHandler
-from . import utils as utils
 from .app import main as main
 from .app import entry_point as cli  # noqa: F401
-from . import make_prompt as make_prompt
 
-__version__ = "0.8.9"
+__version__ = "0.9.0"

sgpt/app.py

@@ -1,24 +1,23 @@
 """
-shell-gpt: An interface to OpenAI's ChatGPT (GPT-3.5) API
-
-This module provides a simple interface for OpenAI's ChatGPT API using Typer
+This module provides a simple interface for OpenAI API using Typer
 as the command line interface. It supports different modes of output including
 shell commands and code, and allows users to specify the desired OpenAI model
 and length and other options of the output. Additionally, it supports executing
 shell commands directly from the interface.
-
-API Key is stored locally for easy use in future runs.
 """
 # To allow users to use arrow keys in the REPL.
 import readline  # noqa: F401
 import sys
 
 import typer
-
-# Click is part of typer.
 from click import BadArgumentUsage, MissingParameter
 
-from sgpt import ChatHandler, DefaultHandler, OpenAIClient, ReplHandler, cfg
+from sgpt.client import OpenAIClient
+from sgpt.config import cfg
+from sgpt.handlers.chat_handler import ChatHandler
+from sgpt.handlers.default_handler import DefaultHandler
+from sgpt.handlers.repl_handler import ReplHandler
+from sgpt.role import DefaultRoles, SystemRole
 from sgpt.utils import ModelOptions, get_edited_prompt, run_command
 
 
@@ -80,17 +79,40 @@ def main(
         callback=ChatHandler.show_messages_callback,
         rich_help_panel="Chat Options",
     ),
-    list_chat: bool = typer.Option(
+    list_chats: bool = typer.Option(
         False,
         help="List all existing chat ids.",
         callback=ChatHandler.list_ids,
         rich_help_panel="Chat Options",
     ),
+    role: str = typer.Option(
+        None,
+        help="System role for GPT model.",
+        rich_help_panel="Role Options",
+    ),
+    create_role: str = typer.Option(
+        None,
+        help="Create role.",
+        callback=SystemRole.create,
+        rich_help_panel="Role Options",
+    ),
+    show_role: str = typer.Option(
+        None,
+        help="Show role.",
+        callback=SystemRole.show,
+        rich_help_panel="Role Options",
+    ),
+    list_roles: bool = typer.Option(
+        False,
+        help="List roles.",
+        callback=SystemRole.list,
+        rich_help_panel="Role Options",
+    ),
 ) -> None:
     stdin_passed = not sys.stdin.isatty()
 
     if stdin_passed and not repl:
-        prompt = sys.stdin.read() + (prompt or "")
+        prompt = f"{sys.stdin.read()}\n\n{prompt or ''}"
 
     if not prompt and not editor and not repl:
         raise MissingParameter(param_hint="PROMPT", param_type="string")
@@ -111,9 +133,11 @@ def main(
     api_key = cfg.get("OPENAI_API_KEY")
     client = OpenAIClient(api_host, api_key)
 
+    role_class = DefaultRoles.get(shell, code) if not role else SystemRole.get(role)
+
     if repl:
         # Will be in infinite loop here until user exits with Ctrl+C.
-        ReplHandler(client, repl, shell, code).handle(
+        ReplHandler(client, repl, role_class).handle(
             prompt,
             model=model.value,
             temperature=temperature,
@@ -123,7 +147,7 @@ def main(
         )
 
     if chat:
-        full_completion = ChatHandler(client, chat, shell, code).handle(
+        full_completion = ChatHandler(client, chat, role_class).handle(
             prompt,
             model=model.value,
             temperature=temperature,
@@ -132,7 +156,7 @@ def main(
             caching=cache,
         )
     else:
-        full_completion = DefaultHandler(client, shell, code).handle(
+        full_completion = DefaultHandler(client, role_class).handle(
             prompt,
             model=model.value,
             temperature=temperature,

sgpt/client.py

@@ -4,7 +4,8 @@
 
 import requests
 
-from sgpt import Cache, cfg
+from .cache import Cache
+from .config import cfg
 
 CACHE_LENGTH = int(cfg.get("CACHE_LENGTH"))
 CACHE_PATH = Path(cfg.get("CACHE_PATH"))
@@ -27,7 +28,7 @@ def _request(
         top_probability: float = 1,
     ) -> Generator[str, None, None]:
         """
-        Make request to OpenAI ChatGPT API, read more:
+        Make request to OpenAI API, read more:
         https://platform.openai.com/docs/api-reference/chat
 
         :param messages: List of messages {"role": user or assistant, "content": message_string}

sgpt/config.py

@@ -6,26 +6,28 @@
 
 from click import UsageError
 
-from sgpt.utils import ModelOptions
+from .utils import ModelOptions
 
 CONFIG_FOLDER = os.path.expanduser("~/.config")
-CONFIG_PATH = Path(CONFIG_FOLDER) / "shell_gpt" / ".sgptrc"
+SHELL_GPT_CONFIG_FOLDER = Path(CONFIG_FOLDER) / "shell_gpt"
+SHELL_GPT_CONFIG_PATH = SHELL_GPT_CONFIG_FOLDER / ".sgptrc"
+ROLE_STORAGE_PATH = SHELL_GPT_CONFIG_FOLDER / "roles"
+CHAT_CACHE_PATH = Path(gettempdir()) / "chat_cache"
+CACHE_PATH = Path(gettempdir()) / "cache"
 
 # TODO: Refactor ENV variables with SGPT_ prefix.
 DEFAULT_CONFIG = {
     # TODO: Refactor it to CHAT_STORAGE_PATH.
-    "CHAT_CACHE_PATH": os.getenv(
-        "CHAT_CACHE_PATH", str(Path(gettempdir()) / "shell_gpt" / "chat_cache")
-    ),
-    "CACHE_PATH": os.getenv(
-        "CACHE_PATH", str(Path(gettempdir()) / "shell_gpt" / "cache")
-    ),
+    "CHAT_CACHE_PATH": os.getenv("CHAT_CACHE_PATH", str(CHAT_CACHE_PATH)),
+    "CACHE_PATH": os.getenv("CACHE_PATH", str(CACHE_PATH)),
     "CHAT_CACHE_LENGTH": int(os.getenv("CHAT_CACHE_LENGTH", "100")),
     "CACHE_LENGTH": int(os.getenv("CHAT_CACHE_LENGTH", "100")),
     "REQUEST_TIMEOUT": int(os.getenv("REQUEST_TIMEOUT", "60")),
     "DEFAULT_MODEL": os.getenv("DEFAULT_MODEL", ModelOptions.GPT3.value),
     "OPENAI_API_HOST": os.getenv("OPENAI_API_HOST", "https://api.openai.com"),
     "DEFAULT_COLOR": os.getenv("DEFAULT_COLOR", "magenta"),
+    "ROLE_STORAGE_PATH": os.getenv("ROLE_STORAGE_PATH", str(ROLE_STORAGE_PATH)),
+    "SYSTEM_ROLES": os.getenv("SYSTEM_ROLES", "false"),
     # New features might add their own config variables here.
 }
 
@@ -77,4 +79,4 @@ def get(self, key: str) -> str:  # type: ignore
         return value
 
 
-cfg = Config(CONFIG_PATH, **DEFAULT_CONFIG)
+cfg = Config(SHELL_GPT_CONFIG_PATH, **DEFAULT_CONFIG)