-
Notifications
You must be signed in to change notification settings - Fork 42
/
yolo.py
executable file
·198 lines (159 loc) · 6.46 KB
/
yolo.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
#!/usr/bin/env python3
# MIT License
# Copyright (c) 2023-2024 wunderwuzzi23
# Greetings from Seattle!
import os
import platform
from ai_model import AIModel, GroqModel, OpenAIModel, OllamaModel, AnthropicModel, AzureOpenAIModel
import sys
import subprocess
import dotenv
import distro
import yaml
import pyperclip
from termcolor import colored
from colorama import init
def read_config():
## 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)
config_file = os.path.join(prompt_path, "yolo.yaml")
with open(config_file, 'r') as file:
return yaml.safe_load(file)
def get_system_prompt(shell):
## 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)
## Load the prompt and prep it
prompt_file = os.path.join(prompt_path, "prompt.txt")
system_prompt = open(prompt_file,"r").read()
system_prompt = system_prompt.replace("{shell}", shell)
system_prompt = system_prompt.replace("{os}", get_os_friendly_name())
return system_prompt
def ensure_prompt_is_question(prompt):
if prompt[-1:] != "?" and prompt[-1:] != ".":
prompt+="?"
return prompt
def print_usage(config):
print("Yolo v0.5 - by @wunderwuzzi23 (June 29, 2024)")
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()
print("Current configuration per yolo.yaml:")
print("* API : " + str(config["api"]))
print("* Model : " + str(config["model"]))
print("* Temperature : " + str(config["temperature"]))
print("* Max. Tokens : " + str(config["max_tokens"]))
print("* Safety : " + str(bool(config["safety"])))
print("* Command Color: " + str(config["suggested_command_color"]))
def get_os_friendly_name():
os_name = platform.system()
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 chat_completion(client, query, config, shell):
if query == "":
print ("No user prompt specified.")
sys.exit(-1)
system_prompt = get_system_prompt(shell)
response = client.chat(
model=config["model"],
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": query}
],
temperature=config["temperature"],
max_tokens=config["max_tokens"])
return response
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)
def check_for_markdown(response):
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)
def missing_posix_display():
return 'DISPLAY' not in os.environ or not os.environ["DISPLAY"]
def prompt_user_for_action(config, ask_flag, response):
print("Command: " + colored(response, config["suggested_command_color"], attrs=['bold']))
modify_snippet = ""
if bool(config["modify"]) == True:
modify_snippet = " [m]odify"
copy_to_clipboard_snippet = " [c]opy to clipboard"
if os.name == "posix" and missing_posix_display():
if get_os_friendly_name() != "Darwin/macOS":
copy_to_clipboard_snippet = ""
if bool(config["safety"]) == True or ask_flag == True:
prompt_text = f"Execute command? [Y]es [n]o{modify_snippet}{copy_to_clipboard_snippet} ==> "
print(prompt_text, end = '')
user_input = input()
return user_input
if bool(config["safety"]) == False:
return "Y"
def eval_user_intent_and_execute(client, config, user_input, command, shell, ask_flag):
if user_input.upper() not in ["", "Y", "C", "M"]:
print("No action taken.")
return
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)
if bool(config["modify"]) and user_input.upper() == "M":
print("Modify prompt: ", end = '')
modded_query = input()
modded_response = chat_completion(client, modded_query, config, shell)
check_for_issue(modded_response)
check_for_markdown(modded_response)
user_intent = prompt_user_for_action(config, ask_flag, modded_response)
print()
eval_user_intent_and_execute(client, config, user_intent, modded_response, shell, ask_flag)
if user_input.upper() == "C":
if os.name == "posix" and missing_posix_display():
if get_os_friendly_name() != "Darwin/macOS":
return
pyperclip.copy(command)
print("Copied command to clipboard.")
def main():
init() #Enable color output on Windows using colorama
dotenv.load_dotenv()
config = read_config()
client = AIModel.get_model_client(config)
# Unix based SHELL (/bin/bash, /bin/zsh), otherwise assuming it's Windows
shell = os.environ.get("SHELL", "powershell.exe")
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
# Parse arguments and make sure we have at least a single word
if len(sys.argv) < 2:
print_usage(config)
sys.exit(-1)
# 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
# To allow easy/natural use we don't require the input to be a single string.
# 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)
## core prompting loop logic
result = chat_completion(client, user_prompt, config, shell)
check_for_issue(result)
check_for_markdown(result)
users_intent = prompt_user_for_action(config, ask_flag, result)
print()
eval_user_intent_and_execute(client, config, users_intent, result, shell, ask_flag)
if __name__ == "__main__":
main()