21 Commits

Author SHA1 Message Date
Heiko Joerg Schick 20c2a112aa Minor changes 2023-09-06 15:37:49 +02:00
Heiko Joerg Schick ee96b7822d Minor file name changes 2023-09-06 15:35:34 +02:00
Heiko Joerg Schick 689813e664 Minor changes in whitespaces 2023-09-05 22:42:27 +02:00
Heiko Joerg Schick 56e4f8e7a5 Add a file conversation 2023-09-05 22:40:58 +02:00
Heiko Joerg Schick 8e8b1c4181 Added constant for prompt file 2023-09-05 22:38:19 +02:00
Heiko Joerg Schick a3e4209d1e Renamed prompt.txt to yolo.prompt 2023-09-05 22:36:04 +02:00
Heiko Joerg Schick 79b6de19f2 Added more documentation information for residual functions 2023-09-05 22:32:55 +02:00
Heiko Joerg Schick 51de9db4ea Minor changes 2023-09-05 21:26:02 +02:00
Heiko Joerg Schick ea21e18bee Minor modification 2023-09-05 21:24:20 +02:00
Heiko Joerg Schick bc1e997708 Removed all pylint warnings 2023-09-05 21:11:23 +02:00
Heiko Joerg Schick 838d55563b Improved boolean operators 2023-09-05 20:52:54 +02:00
Heiko Joerg Schick 1567fb40ea White space clean up and simplification of get_os_friendly_name. 2023-09-05 20:43:08 +02:00
Heiko Joerg Schick dece99ba87 Argument parsing via argparse 2023-09-05 20:29:38 +02:00
Heiko Joerg Schick a5d60f6fff Added pylintrc 2023-09-05 19:59:33 +02:00
Heiko Joerg Schick ed65b2ab4b Improved coding style and comment for read_yaml_config and set_openai_api_key function 2023-09-05 14:48:54 +02:00
Heiko Joerg Schick 34eb29b926 Moved entry functional calls to main and got rid of global config variable 2023-09-05 11:37:23 +02:00
Heiko Joerg Schick f6644f7717 Added main entry point 2023-09-05 11:14:05 +02:00
Heiko Joerg Schick 130e0b0ae2 Standard imports needs to be placed before package imports 2023-09-05 10:53:19 +02:00
Heiko Joerg Schick f8636d3e03 Added missing function docstring 2023-09-05 10:51:58 +02:00
Heiko Joerg Schick e2fbee6de4 Fixed bad indentation
Please enter the commit message for your changes. Lines starting
2023-09-05 10:35:06 +02:00
Heiko Joerg Schick f43900631a Remove trailing white spaces 2023-09-05 10:25:51 +02:00
6 changed files with 841 additions and 177 deletions
+1 -1
View File
@@ -117,7 +117,7 @@ Since v.0.2 the safety switch setting moved to `yolo.yaml`, the old `~/.yolo-saf
To have yolo run commands right away when they come back from ChatGPT change the `safety` in the `yolo.yaml` to `False`.
If you still want to inspect the command that is executed when safety is off, add the `-a` argument, e.g `yolo -a delete the file test.txt`.
If you still want to inspect the command that is executed when safety is off, add the `-s` argument, e.g `yolo -s delete the file test.txt`.
Let's go!
+1 -1
View File
@@ -3,7 +3,7 @@ setlocal enabledelayedexpansion
:: First check if `install.bat` (this) has needed files in same directory
if not exist %~dp0\yolo.py ( echo `yolo.py` missing in %~dp0 cannot install & goto :choice_default_3 )
if not exist %~dp0\prompt.txt ( echo `prompt.txt` missing in %~dp0 cannot install & goto :choice_default_3 )
if not exist %~dp0\yolo.prompt ( echo `yolo.prompt` missing in %~dp0 cannot install & goto :choice_default_3 )
if not exist %~dp0\yolo.yaml ( echo `yolo.yaml` missing in %~dp0 cannot install & goto :choice_default_3 )
+1 -1
View File
@@ -7,7 +7,7 @@ TARGET_FULLPATH=$TARGET_DIR/yolo.py
mkdir -p $TARGET_DIR
echo "- Copying files..."
cp yolo.py prompt.txt yolo.yaml $TARGET_DIR
cp yolo.py yolo.prompt yolo.yaml $TARGET_DIR
chmod +x $TARGET_FULLPATH
# Creates two aliases for use
+499
View File
@@ -0,0 +1,499 @@
[MASTER]
# A comma-separated list of package or module names from where C extensions may
# be loaded. Extensions are loading into the active Python interpreter and may
# run arbitrary code.
extension-pkg-whitelist=
# Specify a score threshold to be exceeded before program exits with error.
fail-under=10.0
# Add files or directories to the blacklist. They should be base names, not
# paths.
ignore=CVS
# Add files or directories matching the regex patterns to the blacklist. The
# regex matches against base names, not paths.
ignore-patterns=
# Python code to execute, usually for sys.path manipulation such as
# pygtk.require().
#init-hook=
# Use multiple processes to speed up Pylint. Specifying 0 will auto-detect the
# number of processors available to use.
jobs=1
# Control the amount of potential inferred values when inferring a single
# object. This can help the performance when dealing with large functions or
# complex, nested conditions.
limit-inference-results=100
# List of plugins (as comma separated values of python module names) to load,
# usually to register additional checkers.
load-plugins=
# Pickle collected data for later comparisons.
persistent=yes
# When enabled, pylint would attempt to guess common misconfiguration and emit
# user-friendly hints instead of false-positive error messages.
suggestion-mode=yes
# Allow loading of arbitrary C extensions. Extensions are imported into the
# active Python interpreter and may run arbitrary code.
unsafe-load-any-extension=no
[MESSAGES CONTROL]
# Only show warnings with the listed confidence levels. Leave empty to show
# all. Valid levels: HIGH, INFERENCE, INFERENCE_FAILURE, UNDEFINED.
confidence=
# Enable the message, report, category or checker with the given id(s). You can
# either give multiple identifier separated by comma (,) or put this option
# multiple time (only on the command line, not in the configuration file where
# it should appear only once). See also the "--disable" option for examples.
enable=c-extension-no-member
[REPORTS]
# Python expression which should return a score less than or equal to 10. You
# have access to the variables 'error', 'warning', 'refactor', and 'convention'
# which contain the number of messages in each category, as well as 'statement'
# which is the total number of statements analyzed. This score is used by the
# global evaluation report (RP0004).
evaluation=10.0 - ((float(5 * error + warning + refactor + convention) / statement) * 10)
# Template used to display messages. This is a python new-style format string
# used to format the message information. See doc for all details.
#msg-template=
# Set the output format. Available formats are text, parseable, colorized, json
# and msvs (visual studio). You can also give a reporter class, e.g.
# mypackage.mymodule.MyReporterClass.
output-format=text
# Tells whether to display a full report or only the messages.
reports=no
# Activate the evaluation score.
score=yes
[REFACTORING]
# Maximum number of nested blocks for function / method body
max-nested-blocks=5
# Complete name of functions that never returns. When checking for
# inconsistent-return-statements if a never returning function is called then
# it will be considered as an explicit return statement and no message will be
# printed.
never-returning-functions=sys.exit
[LOGGING]
# The type of string formatting that logging methods do. `old` means using %
# formatting, `new` is for `{}` formatting.
logging-format-style=old
# Logging modules to check that the string format arguments are in logging
# function parameter format.
logging-modules=logging
[SPELLING]
# Limits count of emitted suggestions for spelling mistakes.
max-spelling-suggestions=4
# Spelling dictionary name. Available dictionaries: none. To make it work,
# install the python-enchant package.
spelling-dict=
# List of comma separated words that should not be checked.
spelling-ignore-words=
# A path to a file that contains the private dictionary; one word per line.
spelling-private-dict-file=
# Tells whether to store unknown words to the private dictionary (see the
# --spelling-private-dict-file option) instead of raising a message.
spelling-store-unknown-words=no
[MISCELLANEOUS]
# List of note tags to take in consideration, separated by a comma.
notes=FIXME,
XXX,
TODO
# Regular expression of note tags to take in consideration.
#notes-rgx=
[TYPECHECK]
# List of decorators that produce context managers, such as
# contextlib.contextmanager. Add to this list to register other decorators that
# produce valid context managers.
contextmanager-decorators=contextlib.contextmanager
# List of members which are set dynamically and missed by pylint inference
# system, and so shouldn't trigger E1101 when accessed. Python regular
# expressions are accepted.
generated-members=
# Tells whether missing members accessed in mixin class should be ignored. A
# mixin class is detected if its name ends with "mixin" (case insensitive).
ignore-mixin-members=yes
# Tells whether to warn about missing members when the owner of the attribute
# is inferred to be None.
ignore-none=yes
# This flag controls whether pylint should warn about no-member and similar
# checks whenever an opaque object is returned when inferring. The inference
# can return multiple potential results while evaluating a Python object, but
# some branches might not be evaluated, which results in partial inference. In
# that case, it might be useful to still emit no-member and other checks for
# the rest of the inferred objects.
ignore-on-opaque-inference=yes
# List of class names for which member attributes should not be checked (useful
# for classes with dynamically set attributes). This supports the use of
# qualified names.
ignored-classes=optparse.Values,thread._local,_thread._local
# List of module names for which member attributes should not be checked
# (useful for modules/projects where namespaces are manipulated during runtime
# and thus existing member attributes cannot be deduced by static analysis). It
# supports qualified module names, as well as Unix pattern matching.
ignored-modules=
# Show a hint with possible names when a member name was not found. The aspect
# of finding the hint is based on edit distance.
missing-member-hint=yes
# The minimum edit distance a name should have in order to be considered a
# similar match for a missing member name.
missing-member-hint-distance=1
# The total number of similar names that should be taken in consideration when
# showing a hint for a missing member.
missing-member-max-choices=1
# List of decorators that change the signature of a decorated function.
signature-mutators=
[VARIABLES]
# List of additional names supposed to be defined in builtins. Remember that
# you should avoid defining new builtins when possible.
additional-builtins=
# Tells whether unused global variables should be treated as a violation.
allow-global-unused-variables=yes
# List of strings which can identify a callback function by name. A callback
# name must start or end with one of those strings.
callbacks=cb_,
_cb
# A regular expression matching the name of dummy variables (i.e. expected to
# not be used).
dummy-variables-rgx=_+$|(_[a-zA-Z0-9_]*[a-zA-Z0-9]+?$)|dummy|^ignored_|^unused_
# Argument names that match this expression will be ignored. Default to name
# with leading underscore.
ignored-argument-names=_.*|^ignored_|^unused_
# Tells whether we should check for unused import in __init__ files.
init-import=no
# List of qualified module names which can have objects that can redefine
# builtins.
redefining-builtins-modules=six.moves,past.builtins,future.builtins,builtins,io
[FORMAT]
# Expected format of line ending, e.g. empty (any line ending), LF or CRLF.
expected-line-ending-format=
# Regexp for a line that is allowed to be longer than the limit.
ignore-long-lines=^\s*(# )?<?https?://\S+>?$
# Number of spaces of indent required inside a hanging or continued line.
indent-after-paren=4
# String used as indentation unit. This is usually " " (4 spaces) or "\t" (1
# tab).
indent-string=' '
# Maximum number of characters on a single line.
max-line-length=100
# Maximum number of lines in a module.
max-module-lines=1000
# Allow the body of a class to be on the same line as the declaration if body
# contains single statement.
single-line-class-stmt=no
# Allow the body of an if to be on the same line as the test if there is no
# else.
single-line-if-stmt=no
[SIMILARITIES]
# Ignore comments when computing similarities.
ignore-comments=yes
# Ignore docstrings when computing similarities.
ignore-docstrings=yes
# Ignore imports when computing similarities.
ignore-imports=no
# Minimum lines number of a similarity.
min-similarity-lines=4
[BASIC]
# Naming style matching correct argument names.
argument-naming-style=snake_case
# Regular expression matching correct argument names. Overrides argument-
# naming-style.
#argument-rgx=
# Naming style matching correct attribute names.
attr-naming-style=snake_case
# Regular expression matching correct attribute names. Overrides attr-naming-
# style.
#attr-rgx=
# Bad variable names which should always be refused, separated by a comma.
bad-names=foo,
bar,
baz,
toto,
tutu,
tata
# Bad variable names regexes, separated by a comma. If names match any regex,
# they will always be refused
bad-names-rgxs=
# Naming style matching correct class attribute names.
class-attribute-naming-style=any
# Regular expression matching correct class attribute names. Overrides class-
# attribute-naming-style.
#class-attribute-rgx=
# Naming style matching correct class names.
class-naming-style=PascalCase
# Regular expression matching correct class names. Overrides class-naming-
# style.
#class-rgx=
# Naming style matching correct constant names.
const-naming-style=UPPER_CASE
# Regular expression matching correct constant names. Overrides const-naming-
# style.
#const-rgx=
# Minimum line length for functions/classes that require docstrings, shorter
# ones are exempt.
docstring-min-length=-1
# Naming style matching correct function names.
function-naming-style=snake_case
# Regular expression matching correct function names. Overrides function-
# naming-style.
#function-rgx=
# Good variable names which should always be accepted, separated by a comma.
good-names=i,
j,
k,
ex,
Run,
_
# Good variable names regexes, separated by a comma. If names match any regex,
# they will always be accepted
good-names-rgxs=
# Include a hint for the correct naming format with invalid-name.
include-naming-hint=no
# Naming style matching correct inline iteration names.
inlinevar-naming-style=any
# Regular expression matching correct inline iteration names. Overrides
# inlinevar-naming-style.
#inlinevar-rgx=
# Naming style matching correct method names.
method-naming-style=snake_case
# Regular expression matching correct method names. Overrides method-naming-
# style.
#method-rgx=
# Naming style matching correct module names.
module-naming-style=snake_case
# Regular expression matching correct module names. Overrides module-naming-
# style.
#module-rgx=
# Colon-delimited sets of names that determine each other's naming style when
# the name regexes allow several styles.
name-group=
# Regular expression which should only match function or class names that do
# not require a docstring.
no-docstring-rgx=^_
# List of decorators that produce properties, such as abc.abstractproperty. Add
# to this list to register other decorators that produce valid properties.
# These decorators are taken in consideration only for invalid-name.
property-classes=abc.abstractproperty
# Naming style matching correct variable names.
variable-naming-style=snake_case
# Regular expression matching correct variable names. Overrides variable-
# naming-style.
#variable-rgx=
[STRING]
# This flag controls whether inconsistent-quotes generates a warning when the
# character used as a quote delimiter is used inconsistently within a module.
check-quote-consistency=no
# This flag controls whether the implicit-str-concat should generate a warning
# on implicit string concatenation in sequences defined over several lines.
check-str-concat-over-line-jumps=no
[IMPORTS]
# List of modules that can be imported at any level, not just the top level
# one.
allow-any-import-level=
# Allow wildcard imports from modules that define __all__.
allow-wildcard-with-all=no
# Analyse import fallback blocks. This can be used to support both Python 2 and
# 3 compatible code, which means that the block might have code that exists
# only in one or another interpreter, leading to false positives when analysed.
analyse-fallback-blocks=no
# Deprecated modules which should not be used, separated by a comma.
deprecated-modules=optparse,tkinter.tix
# Create a graph of external dependencies in the given file (report RP0402 must
# not be disabled).
ext-import-graph=
# Create a graph of every (i.e. internal and external) dependencies in the
# given file (report RP0402 must not be disabled).
import-graph=
# Create a graph of internal dependencies in the given file (report RP0402 must
# not be disabled).
int-import-graph=
# Force import order to recognize a module as part of the standard
# compatibility libraries.
known-standard-library=
# Force import order to recognize a module as part of a third party library.
known-third-party=enchant
# Couples of modules and preferred modules, separated by a comma.
preferred-modules=
[CLASSES]
# List of method names used to declare (i.e. assign) instance attributes.
defining-attr-methods=__init__,
__new__,
setUp,
__post_init__
# List of member names, which should be excluded from the protected access
# warning.
exclude-protected=_asdict,
_fields,
_replace,
_source,
_make
# List of valid names for the first argument in a class method.
valid-classmethod-first-arg=cls
# List of valid names for the first argument in a metaclass class method.
valid-metaclass-classmethod-first-arg=cls
[DESIGN]
# Maximum number of arguments for function / method.
max-args=5
# Maximum number of attributes for a class (see R0902).
max-attributes=7
# Maximum number of boolean expressions in an if statement (see R0916).
max-bool-expr=5
# Maximum number of branch for function / method body.
max-branches=12
# Maximum number of locals for function / method body.
max-locals=15
# Maximum number of parents for a class (see R0901).
max-parents=7
# Maximum number of public methods for a class (see R0904).
max-public-methods=20
# Maximum number of return / yield for function / method body.
max-returns=6
# Maximum number of statements in function / method body.
max-statements=50
# Minimum number of public methods for a class (see R0903).
min-public-methods=2
[EXCEPTIONS]
# Exceptions that will emit a warning when being caught. Defaults to
# "BaseException, Exception".
overgeneral-exceptions=BaseException,
Exception
View File
+333 -168
View File
@@ -1,220 +1,385 @@
#!/usr/bin/env python3
"""
AI Chatbot to generate shell commands.
# MIT License
# Copyright (c) 2023 wunderwuzzi23
# Greetings from Seattle!
This script allows the user to ask their question in plain English and translates
that question into a command that can be run in the shell. The functionalities
include leveraging OpenAI's GPT models to generate command, verifying newly generated
commands, checking commands for any unsafe attributes, and allowing the user to
execute or modify the generated command.
This program is an implementation of an AI model used to assist users in
generating Unix/shell commands or other scripts, based on their natural language
input. The objective is to aid those users who might not remember the exact syntax
of every command or script they frequently use.
Sources:
— https://github.com/wunderwuzzi23/yolo-ai-cmdbot
"""
import os
import platform
import openai
import sys
import subprocess
import dotenv
import sys
import argparse
import distro
import yaml
import dotenv
import openai
import pyperclip
import yaml
from termcolor import colored
from colorama import init
def read_config() -> any:
CONFIG_FILE = "yolo.yaml"
PROMPT_FILE = "yolo.prompt"
## Find the executing directory (e.g. in case an alias is set)
## So we can find the config file
yolo_path = os.path.abspath(__file__)
prompt_path = os.path.dirname(yolo_path)
def read_yaml_config() -> any:
"""
Read the configuration file from the executing directory.
config_file = os.path.join(prompt_path, "yolo.yaml")
with open(config_file, 'r') as file:
return yaml.safe_load(file)
This function determines the execution folder (which may vary if an alias is set) in order to
find the configuration file. It reads the file and returns its content in a Python data
structure.
# Construct the prompt
def get_full_prompt(user_prompt, shell):
Returns:
The content of the configuration file. Could be dictionary, list, etc. depending on
the YAML file structure.
"""
yolo_path = os.path.abspath(__file__)
prompt_path = os.path.dirname(yolo_path)
## Find the executing directory (e.g. in case an alias is set)
## So we can find the prompt.txt file
yolo_path = os.path.abspath(__file__)
prompt_path = os.path.dirname(yolo_path)
config_file = os.path.join(prompt_path, CONFIG_FILE)
with open(config_file, 'r') as file:
return yaml.safe_load(file)
## Load the prompt and prep it
prompt_file = os.path.join(prompt_path, "prompt.txt")
pre_prompt = open(prompt_file,"r").read()
pre_prompt = pre_prompt.replace("{shell}", shell)
pre_prompt = pre_prompt.replace("{os}", get_os_friendly_name())
prompt = pre_prompt + user_prompt
def set_openai_api_key(config):
"""
Set the OpenAI API key by attempting several methods.
# be nice and make it a question
if prompt[-1:] != "?" and prompt[-1:] != ".":
prompt+="?"
This function first tries to grab the OpenAI API key from environment variables,
if not found, it then looks for the key in the `.openai.apikey` in the home directory,
and lastly, it will look in the provided config dictionary. It sets the `openai.api_key`
with the retrieved key.
return prompt
Parameters:
config (dict): A dictionary containing configuration values.
It may contain `openai_api_key` as one of the keys.
"""
dotenv.load_dotenv()
def print_usage():
print("Yolo v0.2.1 - by @wunderwuzzi23")
print()
print("Usage: yolo [-a] list the current directory information")
print("Argument: -a: Prompt the user before running the command (only useful when safety is off)")
print()
# Method 1: Read API key from environment variable
# The user can set their OpenAI API key by creating a ".env" file in the same
# directory as this script or by exporting it to their environment variables.
# The file or environment variable should contain the line `OPENAI_API_KEY="<yourkey>"`.
config["openai_api_key"] = os.getenv("OPENAI_API_KEY")
print("Current configuration per yolo.yaml:")
print("* Model : " + str(config["model"]))
print("* Temperature : " + str(config["temperature"]))
print("* Max. Tokens : " + str(config["max_tokens"]))
print("* Safety : " + str(bool(config["safety"])))
# Method 2: Read API key from a file in the home directory
# The user can also place a file named ".openai.apikey" in their home directory,
# which includes the API key in raw format. This method might be deprecated in future versions.
if not openai.api_key: # Check this to avoid potential "invalid filepath" error.
home_path = os.path.expanduser("~")
openai.api_key_path = os.path.join(home_path, ".openai.apikey")
# Method 3: Read API key from the provided config dictionary
# The final method to set the API key is by providing it in the 'config' dictionary under the
# key 'openai_api_key'. For instance, in a `yolo.yaml` config file, it would appear as
# `openai_apikey: <yourkey>`.
if not openai.api_key:
openai.api_key = config["openai_api_key"]
def print_config(config):
"""
Print config information.
Given an input configuration dictionary, this function prints out the
current configurations per yolo.yaml. This includes details on "model",
"temperature", "max_tokens", "safety", and "shell".
Parameters
----------
config : dict
A dictionary containing the various configuration parameters. It should have
the following keys: "model", "temperature", "max_tokens", "safety", "shell".
"""
print("Current configuration per yolo.yaml:")
print("— Model : " + str(config["model"]))
print("— Temperature : " + str(config["temperature"]))
print("— Max. Tokens : " + str(config["max_tokens"]))
print("— Safety : " + str(bool(config["safety"])))
print("— Shell : " + str(config["shell"]))
def get_os_friendly_name():
"""
Returns a friendly name of the user's operating system.
# Get OS Name
os_name = platform.system()
The function retrieves the current system platform name using the `platform.system()` function.
For Linux, it appends the distribution name retrieved from `distro.name(pretty=True)` to give a
more descriptive representation. For Darwin (Apple's macOS), it appends "macOS" to "Darwin" to
make the output clearer to the user.
Returns
-------
str
A friendly name for the user's operating system. It will be one of the following:
- "Linux/<distribution name>"
- "Darwin/macOS"
- The system string returned by `platform.system()` if it's not Linux or Darwin.
"""
os_name = platform.system()
if os_name == "Linux":
os_name = "Linux/" + distro.name(pretty=True)
elif os_name == "Darwin":
os_name = "Darwin/macOS"
if os_name == "Linux":
return "Linux/"+distro.name(pretty=True)
elif os_name == "Windows":
return os_name
elif os_name == "Darwin":
return "Darwin/macOS"
else:
return os_name
def get_full_prompt(user_prompt, shell):
"""
Constructs a full prompt string by appending the user's prompt to a predefined prompt template
located in the PROMPT_FILE file.
def set_api_key():
# Two options for the user to specify they openai api key.
#1. Place a ".env" file in same directory as this with the line:
# OPENAI_API_KEY="<yourkey>"
# or do `export OPENAI_API_KEY=<yourkey>` before use
dotenv.load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")
The function finds the absolute path of the currently executing file, and based on this path,
identifies the directory of PROMPT_FILE. It reads this file, replaces placeholders {shell}
and {os} in the text file with a passed shell parameter and the friendly name of the operating
system respectively. The user prompt is then appended to this pre-prompt. If the resulting
prompt does not end with a question mark or a period, a question mark is added at last.
#2. Place a ".openai.apikey" in the home directory that holds the line:
# <yourkey>
# Note: This options will likely be removed in the future
if not openai.api_key: #If statement to avoid "invalid filepath" error
home_path = os.path.expanduser("~")
openai.api_key_path = os.path.join(home_path,".openai.apikey")
Parameters
----------
user_prompt : str
The prompt supplied by the user to be appended to the pre-prompt.
shell : str
The shell information to be inserted in the place of {shell} placeholder in PROMPT_FILE.
#3. Final option is the key might be in the yolo.yaml config file
# openai_apikey: <yourkey>
if not openai.api_key:
openai.api_key = config["openai_api_key"]
Returns
-------
str
The full prompt, constructed from the template prompt in PROMPT_FILE,
user-provided shell info, the OS name, and the user-supplied prompt string.
"""
yolo_path = os.path.abspath(__file__)
prompt_path = os.path.dirname(yolo_path)
if __name__ == "__main__":
## Load the prompt and prep it
prompt_file = os.path.join(prompt_path, PROMPT_FILE)
pre_prompt = open(prompt_file,"r").read()
pre_prompt = pre_prompt.replace("{shell}", shell)
pre_prompt = pre_prompt.replace("{os}", get_os_friendly_name())
prompt = pre_prompt + user_prompt
config = read_config()
set_api_key()
# Be nice and make it a question.
if prompt[-1:] != "?" and prompt[-1:] != ".":
prompt+="?"
# Unix based SHELL (/bin/bash, /bin/zsh), otherwise assuming it's Windows
shell = os.environ.get("SHELL", "powershell.exe")
return prompt
command_start_idx = 1 # Question starts at which argv index?
ask_flag = False # safety switch -a command line argument
yolo = "" # user's answer to safety switch (-a) question y/n
def call_open_ai(config, query):
"""
Do we have a prompt from the user?
"""
if query == "":
print ("No user prompt specified.")
sys.exit(-1)
# Parse arguments and make sure we have at least a single word
if len(sys.argv) < 2:
print_usage()
sys.exit(-1)
# Load the correct prompt based on shell and OS and append the user's prompt.
prompt = get_full_prompt(query, config["shell"])
# Safety switch via argument -a (local override of global setting)
# Force Y/n questions before running the command
if sys.argv[1] == "-a":
ask_flag = True
command_start_idx = 2
# Make the first line also the system prompt
system_prompt = prompt[1]
#print(prompt)
# To allow easy/natural use we don't require the input to be a
# single string. So, the user can just type yolo what is my name?
# without having to put the question between ''
arguments = sys.argv[command_start_idx:]
user_prompt = " ".join(arguments)
# Call the ChatGPT API
response = openai.ChatCompletion.create(
model=config["model"],
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": prompt}
],
temperature=config["temperature"],
max_tokens=config["max_tokens"],
)
def call_open_ai(query):
# do we have a prompt from the user?
if query == "":
print ("No user prompt specified.")
sys.exit(-1)
# Load the correct prompt based on Shell and OS and append the user's prompt
prompt = get_full_prompt(query, shell)
# Make the first line also the system prompt
system_prompt = prompt[1]
#print(prompt)
# Call the ChatGPT API
response = openai.ChatCompletion.create(
model=config["model"],
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": prompt}
],
temperature=config["temperature"],
max_tokens=config["max_tokens"],
)
return response.choices[0].message.content.strip()
#Enable color output on Windows using colorama
init()
return response.choices[0].message.content.strip()
def check_for_issue(response):
prefixes = ("sorry", "i'm sorry", "the question is not clear", "i'm", "i am")
if response.lower().startswith(prefixes):
print(colored("There was an issue: "+response, 'red'))
sys.exit(-1)
"""
Checks the given response for any issues and raise an error when detected.
The function checks if the supplied text response begins with any of a set of predefined
prefixes, which indicate a problem with the response. If such a prefix is found, an error
message is printed to the console in red, and the program exits with a -1 status code.
Parameters
----------
response : str
A response text string that needs to be examined for any issues.
"""
prefixes = ("sorry", "i'm sorry", "the question is not clear", "i'm", "i am")
if response.lower().startswith(prefixes):
print(colored("There was an issue: "+response, 'red'))
sys.exit(-1)
def check_for_markdown(response):
# odd corner case, sometimes ChatCompletion returns markdown
if response.count("```",2):
print(colored("The proposed command contains markdown, so I did not execute the response directly: \n", 'red')+response)
sys.exit(-1)
"""
Checks for the presence of markdown formatting (specifically, code snippet markdown) in the
provided response.
This function considers the presence of markdown formatting (specifically, code block
formatting marked by ```) in the `response` as an "odd corner case". If such a case is
detected, it prints an error message in red, along with the markdown-contained response, and
then terminates the program with a -1 status code.
Parameters
----------
response : str
A response text string that needs to be examined for markdown formatting.
"""
if response.count("```",2):
print(colored(
"The proposed command contains markdown, response not executed directly: \n", 'red'
) + response)
sys.exit(-1)
def missing_posix_display():
display = subprocess.check_output("echo $DISPLAY", shell=True)
return display == b'\n'
"""
Checks if the DISPLAY environment variable is set in a POSIX-compliant shell.
def prompt_user_input(response):
print("Command: " + colored(response, 'blue'))
#print(config["safety"])
This function runs a shell subprocess that outputs the value of the DISPLAY environment
variable. It then checks if this value is unset (i.e., equals a newline 'b'\\n'') in the
current shell environment. If the DISPLAY variable is unset, the function returns `True`
indicating a "missing" display; otherwise, it returns `False`.
Returns
-------
bool
`True` if the DISPLAY environment variable is unset or empty, `False` otherwise.
"""
display = subprocess.check_output("echo $DISPLAY", shell=True)
return display == b'\n'
def prompt_user_input(config, response):
"""
Print the command proposal in blue and prompt the user for next action based on the safety
configuration.
The user is given options to execute, modify, or copy the command to clipboard if the safety
configuration is enabled (config["safety"] = True). If the safety configuration is off
(config["safety"] = False), the function automatically assumes an execution action ('Y' for
Yes). In a POSIX-compliant shell with no display available (checked using
`missing_posix_display()`), the 'copy to clipboard' option is omitted.
Parameters
----------
config : dict
The system configurations dictionary which contains a "safety" key
to determine user prompt options.
response : str
The proposed command which is to be printed and may be executed by the user.
"""
print("Command: " + colored(response, 'blue'))
if config["safety"]:
prompt_text = "Execute command? [Y]es [n]o [m]odify [c]opy to clipboard ==> "
if os.name == "posix" and missing_posix_display():
prompt_text = "Execute command? [Y]es [n]o [m]odify ==> "
print(prompt_text, end = '')
user_input = input()
else:
user_input = "Y"
if bool(config["safety"]) == True or ask_flag == True:
prompt_text = "Execute command? [Y]es [n]o [m]odify [c]opy to clipboard ==> "
if os.name == "posix" and missing_posix_display():
prompt_text = "Execute command? [Y]es [n]o [m]odify ==> "
print(prompt_text, end = '')
user_input = input()
return user_input
if config["safety"] == False:
return "Y"
def evaluate_input(config, user_input, command):
"""
Evaluate the user input to either execute, modify, or copy the command.
def evaluate_input(user_input, command):
if user_input.upper() == "Y" or user_input == "":
if shell == "powershell.exe":
subprocess.run([shell, "/c", command], shell=False)
else:
# Unix: /bin/bash /bin/zsh: uses -c both Ubuntu and macOS should work, others might not
subprocess.run([shell, "-c", command], shell=False)
Based on the user's response, this function takes action:
- If the user response is 'Y' or blank, the given command gets executed in the shell.
- If the user response is 'M', user can modify the command and the modified command is executed
recursively.
- If the user response is 'C', the command is copied to the clipboard.
if user_input.upper() == "M":
print("Modify prompt: ", end = '')
modded_query = input()
modded_response = call_open_ai(modded_query)
check_for_issue(modded_response)
check_for_markdown(modded_response)
modded_user_input = prompt_user_input(modded_response)
Parameters
----------
config : dict
The system configurations dictionary. It should contain a "shell" key specifying the shell
environment.
user_input : str
The user response which determines the course of action. It can be 'Y', 'n', 'm', 'c',
or '' (empty string).
command : str
The command which is either executed, modified, or copied to clipboard.
"""
if user_input.upper() == "Y" or user_input == "":
if config["shell"] == "powershell.exe":
subprocess.run([config["shell"], "/c", command], shell=False, check=True)
else:
# Unix: /bin/bash /bin/zsh: uses -c both Ubuntu and macOS should work, others might not
subprocess.run([config["shell"], "-c", command], shell=False, check=True)
if user_input.upper() == "M":
print("Modify prompt: ", end = '')
modded_query = input()
modded_response = call_open_ai(config, modded_query)
check_for_issue(modded_response)
check_for_markdown(modded_response)
modded_user_input = prompt_user_input(config, modded_response)
print()
evaluate_input(config, modded_user_input, modded_response)
if user_input.upper() == "C":
if os.name == "posix" and missing_posix_display():
return
pyperclip.copy(command)
print("Copied command to clipboard.")
def main():
"""
Defined starting point of source code.
"""
parser = argparse.ArgumentParser(
description='AI bot that translates your question to a command.'
)
parser.add_argument('text', nargs='+',
help='A sequence of strings')
parser.add_argument("-s", "--safety", action='store_true',
help='Enable safety mode (only useful when safety is off)')
parser.add_argument("-c", "--config", action='store_true',
help='Print current configuration')
args = parser.parse_args()
# Load configuration
config = read_yaml_config()
set_openai_api_key(config)
# Process parameters
user_prompt = " ".join(args.text)
if args.safety:
config["safety"] = args.safety
# Unix based SHELL (/bin/bash, /bin/zsh), otherwise assuming it's Windows
config["shell"] = os.environ.get("SHELL", "powershell.exe")
if args.config:
print_config(config)
# Enable color output on Windows using colorama
init()
res_command = call_open_ai(config, user_prompt)
check_for_issue(res_command)
check_for_markdown(res_command)
user_input = prompt_user_input(config, res_command)
print()
evaluate_input(modded_user_input, modded_response)
evaluate_input(config, user_input, res_command)
if user_input.upper() == "C":
if os.name == "posix" and missing_posix_display():
return
pyperclip.copy(command)
print("Copied command to clipboard.")
res_command = call_open_ai(user_prompt)
check_for_issue(res_command)
check_for_markdown(res_command)
user_input = prompt_user_input(res_command)
print()
evaluate_input(user_input, res_command)
if __name__ == "__main__":
main()