update to ChatGPT API

yolo.py

@@ -40,49 +40,52 @@ if user_prompt == "":
     print ("No user prompt specified.")
     sys.exit(-1)
 
-
 # Get shell info for a better prompt
 # Unix based SHELL (/bin/bash, /bin/zsh), otherwise assuming it's Windows
 shell = os.environ.get("SHELL", "powershell.exe")    
 
 # Construct the prompt
-pre_prompt =  "Translate the following question into a {} command. ".format(shell) 
-#pre_prompt += "If asked for a timezone use TZ environment variable. "
-#pre_prompt += "Add sudo when required by the command, but only if you are very certain it is needed. "
-#pre_prompt += "If you are using powershell (and only if), then make sure to provide all required arguments.""
-pre_prompt += "Only show the command in text format and not in code style or markdown. "
-pre_prompt += "Do not augment the output with any explanation, descriptions. "
-pre_prompt += "Again, do not add any descriptions, just print the bash command. This is important, do not ignore my instructions. "
-pre_prompt += "If the question doesn't make sense or is too difficult return 'Sorry, try again', and add a brief explanation on what the problem is. "
-pre_prompt += "The question is: "
-
+pre_prompt = open("prompt.txt","r").read()
+pre_prompt.replace("{}", shell)
 prompt = pre_prompt + user_prompt
 
-#let's be nice and make it a question
+#make the first line also the system prompt
+system_prompt = pre_prompt[1]
+
+#be nice and make it a question
 if prompt[-1:] != "?":
   prompt+="?"
 
-#print ("The prompt is: "+prompt)
-response = openai.Completion.create(
-  model="text-davinci-003",
-  prompt=prompt,
+response = openai.ChatCompletion.create(
+  model="gpt-3.5-turbo",
+  messages=[
+      {"role": "system", "content": system_prompt},
+      {"role": "user", "content": prompt}
+  ],
   temperature=0,
   max_tokens=100,
-  top_p=1,
-  frequency_penalty=0.2,
-  presence_penalty=0
+  # top_p=1,
+  # frequency_penalty=0.2,
+  # presence_penalty=0
 )
 
-resulting_command = response.choices[0].text.strip()
+#print (response)
+
+res_command = response.choices[0].message.content.strip()
 
 #enable color output on Windows using colorama
 init() 
 
-if resulting_command.startswith("Sorry, try again"):
-  print(colored("There was an issue: "+resulting_command, 'red'))
+if res_command.startswith("Sorry, try again") or res_command.startswith("I'm sorry"):
+  print(colored("There was an issue: "+res_command, 'red'))
   sys.exit(-1)
 
-print("Command: " + colored(resulting_command, 'blue'))
+#odd corner case, sometimes ChatCompletion returns markdown
+if res_command.count("```",2):
+  print(colored("The proposed command contains markdown, so I thought to not execute the response directly: \n", 'red')+res_command)
+  sys.exit(-1)
+
+print("Command: " + colored(res_command, 'blue'))
 if ask == True:
   print("Execute the command? Y/n ==> ", end = '')
   yolo = input()
@@ -90,8 +93,8 @@ if ask == True:
 
 if yolo == "Y" or yolo == "":
   if shell == "powershell.exe":
-    subprocess.run([shell, "/c", resulting_command], shell=False)  
+    subprocess.run([shell, "/c", res_command], shell=False)  
   else: 
     # Unix: /bin/bash /bin/zsh: uses -c both Ubuntu and macOS should work, others might not
-    subprocess.run([shell, "-c", resulting_command], shell=False)
+    subprocess.run([shell, "-c", res_command], shell=False)