From a07355db3d13fb8880b4f00d40d6016bf9eeb508 Mon Sep 17 00:00:00 2001 From: Farkhod Sadykov Date: Sun, 19 Feb 2023 02:59:34 +0100 Subject: [PATCH] New arguments, more accurate responses, improvements, tests --- .github/workflows/lint_test.yml | 9 ++-- README.md | 29 +++++----- integrational_tests.py | 96 +++++++++++++++++++++++++++++++++ pre-push.sh | 15 +++--- setup.py | 3 +- sgpt.py | 72 ++++++++++++++++++++++--- tests.py => unittests.py | 24 +++++++-- 7 files changed, 215 insertions(+), 33 deletions(-) create mode 100644 integrational_tests.py rename tests.py => unittests.py (63%) diff --git a/.github/workflows/lint_test.yml b/.github/workflows/lint_test.yml index fbd6c1d..209fbfc 100644 --- a/.github/workflows/lint_test.yml +++ b/.github/workflows/lint_test.yml @@ -24,18 +24,19 @@ jobs: - name: Analysing the code with pylint run: | pylint \ - sgpt.py \ - tests.py \ + *.py \ --disable=missing-function-docstring \ --disable=too-many-arguments \ --disable=missing-module-docstring \ --disable=import-error \ --disable=missing-class-docstring \ --disable=too-many-instance-attributes \ + --disable=too-many-function-args \ + --disable=unspecified-encoding \ --max-line-length=120 - name: Analysing the code with black run: | - black --target-version py310 -l 120 sgpt.py tests.py + black --target-version py310 -l 120 *.py - name: Run unittests run: | - python -m unittest tests.py + python -m unittest unittests.py diff --git a/README.md b/README.md index ce9864c..571e3dc 100644 --- a/README.md +++ b/README.md @@ -1,12 +1,12 @@ # Shell GPT -A command-line interface (CLI) productivity tool powered by OpenAI's Davinci model, will help you accomplish your tasks faster and more efficiently. +A command-line interface (CLI) productivity tool powered by OpenAI's Davinci model, that will help you accomplish your tasks faster and more efficiently.
- +
## Description -`text-davinci-003` is a powerful language model developed by OpenAI that can generate human-like text. It can be used by us, coders, to generate code snippets, comments, documentation and more, helping us increase our productivity and efficiency while coding. +`text-davinci-003` is a powerful language model developed by OpenAI that can generate human-like text. It can be used by us developers, to generate code snippets, comments, documentation and more, helping us increase our productivity and efficiency while coding. 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. @@ -16,7 +16,8 @@ pip install shell-gpt --user ``` On first start you would need to generate and provide your API key, get one [here](https://beta.openai.com/account/api-keys). -## Use cases +## Usage +`sgpt` has a variety of use cases, including simple queries, shell queries, and code queries. ### Simple queries We can use it pretty much as normal search engine, asking about anything, for example: ```shell @@ -28,9 +29,8 @@ sgpt "mass of sun" # -> = 1.99 × 10^30 kg ``` ### Shell queries -Usually we are forgetting commands like `chmod 444` and we want quickly find the answer in google, clicking pages, scrolling, copy&pasting, usually takes some time, but now we "google" and execute it right in the terminal using `--shell` flag `sgpt` will provide only shell commands: +Usually we are forgetting commands like `chmod 444` and we want quickly find the answer in google, but now we "google" and execute it right in the terminal using `--shell` flag `sgpt` will provide only shell commands: ```shell -# Here we are using special flag --shell, which will output only shell commands. sgpt --shell "make all files in current directory read only" # -> chmod 444 * ``` @@ -41,30 +41,35 @@ sgpt --shell --execute "make all files in current directory read only" # -> Execute shell command? [y/N]: y # ... ``` -At this point it is already can solve half of most frequent Google searches, but how far we can push the limits of `text-davinci-003` model? Let's try some docker containers: +Let's try some docker containers: ```shell sgpt -se "start nginx using docker, forward 443 and 80 port, mount current folder with index.html" # -> docker run -d -p 443:443 -p 80:80 -v $(pwd):/usr/share/nginx/html nginx # -> Execute shell command? [y/N]: y # ... ``` -Also, we can provide some parameters name in our prompt, for example, I want to pass input and output file names to ffmpeg: +Also, we can provide some parameters name in our prompt, for example, passing output file names to ffmpeg: ```shell sgpt -se "slow down video twice using ffmpeg, input video name \"input.mp4\" output video name \"output.mp4\"" # -> ffmpeg -i input.mp4 -filter:v "setpts=2.0*PTS" output.mp4 # -> Execute shell command? [y/N]: y # ... ``` -And remember we are in shell, this means we can use outputs of any commands in our prompt, this brings it to another level, here is simple examples with ffmpeg and list of videos in current folder: +We can apply additional shell magic in our prompt, here is simple examples with ffmpeg and list of videos in current folder: ```shell ls # -> 1.mp4 2.mp4 3.mp4 -sgpt -se "using ffmpeg combine multiple videos into one without audio. Video file names $(ls)" +sgpt -se "using ffmpeg combine multiple videos into one without audio. Video file names: $(ls -m)" # -> ffmpeg -i 1.mp4 -i 2.mp4 -i 3.mp4 -filter_complex "[0:v] [1:v] [2:v] concat=n=3:v=1 [v]" -map "[v]" out.mp4 # -> Execute shell command? [y/N]: y # ... ``` -Since GPT-3 models can also do summarization and analyzing of input text, we can ask `text-davinci-003` to find error in logs and provide some details: +Since GPT-3 models can also do summarization and analyzing of input text, we can ask `text-davinci-003` generate for example, commit message: +```shell +sgpt "Generate git commit message with details, my changes: $(git diff)" +# -> Commit message: Implement Model enum and get_edited_prompt() func, add temperature, top_p and editor args for OpenAI request. +``` +Or ask it to find error in logs and provide more details: ```shell sgpt "check these logs, find errors, and explain what the error is about: ${docker logs -n 20 container_name}" # ... @@ -85,7 +90,7 @@ for i in range(1, 101): else: print(i) ``` -Since it is valid python code without any other text, we can redirect the output to file: +Since it is valid python code, we can redirect the output to file: ```shell sgpt --code "solve classic fizz buzz problem using Python" > fizz_buzz.py python fizz_buzz.py diff --git a/integrational_tests.py b/integrational_tests.py new file mode 100644 index 0000000..e0fc2ea --- /dev/null +++ b/integrational_tests.py @@ -0,0 +1,96 @@ +""" +This test module will execute real commands using shell. +This means it will call sgpt.py with command line arguments. +Make sure you have your API key in place ~/.config/shell-gpt/api_key.txt +It is useful for quick tests, saves a bit time. +""" + +import subprocess +import os +from time import sleep +from unittest import TestCase +from tempfile import NamedTemporaryFile + +import typer +from typer.testing import CliRunner +from sgpt import main + +runner = CliRunner() +app = typer.Typer() +app.command()(main) + + +class TestCliApp(TestCase): + def setUp(self) -> None: + # Just to not spam the API. + sleep(2) + + @staticmethod + def get_arguments(prompt, **kwargs): + arguments = [prompt] + for key, value in kwargs.items(): + arguments.append(key) + if isinstance(value, bool): + continue + arguments.append(value) + # arguments.extend([kv for kv in kwargs.items() for kv in kv]) + return arguments + + def test_simple_queries(self): + dict_arguments = { + "prompt": "What is the capital of the Czech Republic?", + "--max-tokens": 32, + "--model": "curie", + } + result = runner.invoke(app, self.get_arguments(**dict_arguments)) + assert result.exit_code == 0 + assert "Prague" in result.stdout + + def test_shell_queries(self): + dict_arguments = { + "prompt": "make a commit using git", + "--max-tokens": 32, + "--model": "davinci", + "--shell": True, + } + result = runner.invoke(app, self.get_arguments(**dict_arguments)) + assert result.exit_code == 0 + assert "git commit" in result.stdout + + def test_code_queries(self): + """ + This test will request from GPT-3 a python code to make CLI app, + which will be written to a temp file, and then it will be executed + in shell with two positional int arguments. As the output we are + expecting the result of multiplying them. + """ + dict_arguments = { + "prompt": ( + "Create a command line application using Python that " + "accepts two integer positional command line arguments " + "and prints the result of multiplying them." + ), + "--max-tokens": 64, + "--model": "davinci", + "--code": True, + } + result = runner.invoke(app, self.get_arguments(**dict_arguments)) + assert result.exit_code == 0 + # Since output will be slightly different, there is no way how to test it precisely. + assert "print" in result.stdout + assert "*" in result.stdout + with NamedTemporaryFile("w+", delete=False) as file: + try: + compile(result.output, file.name, "exec") + except SyntaxError: + assert False, "The output is not valid Python code." + file.seek(0) + file.truncate() + file.write(result.output) + file_path = file.name + number_a = number_b = 2 + # Execute output code in the shell with arguments. + arguments = ["python", file.name, str(number_a), str(number_b)] + script_output = subprocess.run(arguments, stdout=subprocess.PIPE, check=True) + os.remove(file_path) + assert script_output.stdout.decode().strip(), number_a * number_b diff --git a/pre-push.sh b/pre-push.sh index 9a02980..88f00fd 100755 --- a/pre-push.sh +++ b/pre-push.sh @@ -21,24 +21,27 @@ if [ -f "$PWD/venv/bin/activate" ]; then source $PWD/venv/bin/activate fi -if black --check --target-version py310 -l 120 sgpt.py tests.py +if black --exclude venv --check --target-version py310 -l 120 *.py then echo 'Black passed ✅' else echo 'Black failed ❌' - echo 'RUN: black --target-version py310 -l 120 sgpt.py tests.py' + echo 'RUN: black --exclude venv --target-version py310 -l 120 *.py' exit 1 fi -if pylint sgpt.py \ - tests.py \ +if pylint \ + *.py \ --disable=missing-function-docstring \ --disable=too-many-arguments \ --disable=missing-module-docstring \ --disable=import-error \ --disable=missing-class-docstring \ --disable=too-many-instance-attributes \ - --max-line-length=120 + --disable=too-many-function-args \ + --disable=unspecified-encoding \ + --max-line-length=120 \ + --ignore=venv then echo 'Pylint passed ✅' else @@ -46,7 +49,7 @@ else exit 1 fi -if python -m unittest tests.py +if python -m unittest unittests.py then echo 'Unittests passed ✅' else diff --git a/setup.py b/setup.py index 5ee4109..3338e2a 100644 --- a/setup.py +++ b/setup.py @@ -1,8 +1,9 @@ from setuptools import setup, find_packages +# pylint: disable=consider-using-with setup( name="shell_gpt", - version="0.4.1", + version="0.5.0", packages=find_packages(), py_modules=[ "sgpt", diff --git a/sgpt.py b/sgpt.py index 2df4ec9..29d6ec9 100644 --- a/sgpt.py +++ b/sgpt.py @@ -4,20 +4,26 @@ shell-gpt: An interface to OpenAI's GPT-3 API This module provides a simple interface for OpenAI's GPT-3 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 of the output. Additionally, it supports executing shell -commands directly from the interface. +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. """ import os +from enum import Enum from time import sleep from pathlib import Path from getpass import getpass +from types import DynamicClassAttribute +from tempfile import NamedTemporaryFile import typer import requests + +# Click is part of typer. +from click import MissingParameter, BadParameter from rich.progress import Progress, SpinnerColumn, TextColumn @@ -26,6 +32,23 @@ DATA_FOLDER = os.path.expanduser("~/.config") KEY_FILE = Path(DATA_FOLDER) / "shell-gpt" / "api_key.txt" +# pylint: disable=invalid-name +class Model(str, Enum): + davinci = "text-davinci-003" + curie = "text-curie-001" + codex = "code-davinci-002" + + def __str__(self): + return self.name + + @DynamicClassAttribute + def value(self): + return self.name + + +# pylint: enable=invalid-name + + def get_api_key(): if not KEY_FILE.exists(): api_key = getpass(prompt="Please enter your API secret key") @@ -48,13 +71,31 @@ def loading_spinner(func): return wrapper +def get_edited_prompt(): + with NamedTemporaryFile(suffix=".txt", delete=False) as file: + # Create file and store path. + file_path = file.name + editor = os.environ.get("EDITOR", "vim") + # This will write text to file using $EDITOR. + os.system(f"{editor} {file_path}") + # Read file when editor is closed. + with open(file_path, "r") as file: + output = file.read() + os.remove(file_path) + if not output: + raise BadParameter("Couldn't get valid PROMPT from $EDITOR") + return output + + @loading_spinner -def openai_request(prompt, model, max_tokens, api_key): +def openai_request(prompt, model, max_tokens, api_key, temperature, top_p): headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"} data = { "prompt": prompt, "model": model, "max_tokens": max_tokens, + "temperature": temperature, + "top_p": top_p, } response = requests.post(API_URL, headers=headers, json=data, timeout=180) response.raise_for_status() @@ -74,22 +115,39 @@ def typer_writer(text, code, shell, animate): typer.secho(text, fg=color, bold=shell_or_code) +# Using lambda to pass a function to default value, which make it apper as "dynamic" in help. def main( - prompt: str, - model: str = typer.Option("text-davinci-003", help="OpenAI model name."), - max_tokens: int = typer.Option(2048, help="Strict length of output (words)."), + prompt: str = typer.Argument(None, show_default=False, help="The prompt to generate completions for."), + model: Model = typer.Option("davinci", help="GPT-3 model name.", show_choices=True), + max_tokens: int = typer.Option(lambda: 2048, help="Strict length of output (words)."), + temperature: float = typer.Option(lambda: 1.0, min=0.0, max=1.0, help="Randomness of generated output."), + top_probability: float = typer.Option(lambda: 1.0, min=0.1, max=1.0, help="Limits highest probable tokens."), shell: bool = typer.Option(False, "--shell", "-s", help="Provide shell command as output."), execute: bool = typer.Option(False, "--execute", "-e", help="Will execute --shell command."), code: bool = typer.Option(False, help="Provide code as output."), + editor: bool = typer.Option(False, help="Open $EDITOR to provide a prompt."), animation: bool = typer.Option(True, help="Typewriter animation."), spinner: bool = typer.Option(True, help="Show loading spinner during API request."), ): api_key = get_api_key() + if not prompt and not editor: + raise MissingParameter(param_hint="PROMPT", param_type="string") if shell: + # If default values where not changed, make it more accurate. + if temperature == 1.0 == top_probability: + temperature, top_probability = 0.2, 0.9 prompt = f"{prompt}. Provide only shell command as output." elif code: + # If default values where not changed, make it more creative (diverse). + if temperature == 1.0 == top_probability: + temperature, top_probability = 0.8, 0.2 prompt = f"{prompt}. Provide only code as output." - response_text = openai_request(prompt, model, max_tokens, api_key, spinner=spinner) + # Curie has hard cap 2048 + prompt. + if model == "text-curie-001" and max_tokens == 2048: + max_tokens = 1024 + if editor: + prompt = get_edited_prompt() + response_text = openai_request(prompt, model, max_tokens, api_key, temperature, top_probability, spinner=spinner) # For some reason OpenAI returns several leading/trailing white spaces. response_text = response_text.strip() typer_writer(response_text, code, shell, animation) diff --git a/tests.py b/unittests.py similarity index 63% rename from tests.py rename to unittests.py index ad51412..9ed2a49 100644 --- a/tests.py +++ b/unittests.py @@ -15,14 +15,24 @@ class TestMain(unittest.TestCase): self.animation = True self.spinner = True self.api_key = "test-api-key" + self.temperature = 1.0 + self.top_p = 1.0 self.response_text = "Paris" @requests_mock.Mocker() def test_openai_request(self, mock): mock.post(sgpt.API_URL, json={"choices": [{"text": self.response_text}]}, status_code=200) - result = sgpt.openai_request(self.prompt, self.model, self.max_tokens, self.api_key, spinner=self.spinner) + result = sgpt.openai_request( + self.prompt, self.model, self.max_tokens, self.api_key, self.temperature, self.top_p, spinner=self.spinner + ) self.assertEqual(result, self.response_text) - expected_json = {"prompt": self.prompt, "model": self.model, "max_tokens": self.max_tokens} + expected_json = { + "prompt": self.prompt, + "model": self.model, + "max_tokens": self.max_tokens, + "temperature": self.temperature, + "top_p": self.top_p, + } expected_headers = {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"} request = mock.request_history[0] self.assertEqual(request.json(), expected_json) @@ -33,7 +43,15 @@ class TestMain(unittest.TestCase): def test_openai_request_fail(self, mock): mock.post(sgpt.API_URL, status_code=400) with self.assertRaises(requests.exceptions.HTTPError): - sgpt.openai_request(self.prompt, self.model, self.max_tokens, self.api_key, spinner=self.spinner) + sgpt.openai_request( + self.prompt, + self.model, + self.max_tokens, + self.api_key, + self.temperature, + self.top_p, + spinner=self.spinner, + ) if __name__ == "__main__":