Update autotemp.py

autotemp.py

@@ -1,9 +1,16 @@
 import openai
 from termcolor import colored
+from dotenv import load_dotenv
+import os
+import re
+
+# Load environment variables from .env file
+load_dotenv()
 
 class AutoTemp:
-    def __init__(self, api_key, default_temp=0.5, alt_temps=None, auto_select=True):
-        openai.api_key = api_key
+    def __init__(self, default_temp=0.5, alt_temps=None, auto_select=True):
+        self.api_key = os.getenv('OPENAI_API_KEY')
+        openai.api_key = self.api_key
         self.default_temp = default_temp
         self.alt_temps = alt_temps if alt_temps else [0.2, 0.4, 0.6, 0.8, 1.0, 1.2, 1.4]
         self.auto_select = auto_select
@@ -21,13 +28,22 @@ class AutoTemp:
         chosen_temp = float(input(colored("Choose the temperature of the output you like: ", "green")))
         return outputs.get(chosen_temp, "Invalid temperature chosen.")
 
-    def run(self, prompt):
+    def generate_with_chat(self, prompt, temperature):
         try:
-            initial_output = openai.Completion.create(engine="text-davinci-003", prompt=prompt, temperature=self.default_temp)
-            initial_output_text = initial_output['choices'][0]['text'].strip()
+            response = openai.ChatCompletion.create(
+                model="gpt-3.5-turbo",
+                messages=[{"role": "system", "content": "You are a helpful assistant."},
+                          {"role": "user", "content": prompt}],
+                temperature=temperature
+            )
+            message = response['choices'][0]['message']['content']
+            return message.strip()
         except Exception as e:
-            print(colored(f"Error generating initial output: {e}", "red"))
-            initial_output_text = None
+            print(colored(f"Error generating text at temperature {temperature}: {e}", "red"))
+            return None
+
+    def run(self, prompt):
+        initial_output_text = self.generate_with_chat(prompt, self.default_temp)
 
         user_satisfied = self.ask_user_feedback(initial_output_text)
 
@@ -37,31 +53,31 @@ class AutoTemp:
             outputs = {}
             scores = {}
             for temp in self.alt_temps:
-                try:
-                    output = openai.Completion.create(engine="text-davinci-003", prompt=prompt, temperature=temp)
-                    output_text = output['choices'][0]['text'].strip()
+                output_text = self.generate_with_chat(prompt, temp)
+                if output_text:
                     outputs[temp] = output_text
                     eval_prompt = f"You are a task output evaluator. Rate the quality of the following output on a scale from 0 to 1000. Output only an integer. The output is: {output_text}"
-                    score = openai.Completion.create(engine="text-davinci-003", prompt=eval_prompt, temperature=0)
-                    scores[temp] = int(score['choices'][0]['text'].strip())
-                    print(colored(f"Score for temperature {temp}: {scores[temp]}", "yellow"))  # Print the score
-                except Exception as e:
-                    print(colored(f"Error generating text at temperature {temp}: {e}", "red"))
-                    outputs[temp] = None
-                    scores[temp] = 0
+                    score_text = self.generate_with_chat(eval_prompt, 0)
+                    score_match = re.search(r'\d+', score_text)
+                    if score_match:
+                        scores[temp] = int(score_match.group())
+                        print(colored(f"Score for temperature {temp}: {scores[temp]}", "yellow"))
+                    else:
+                        print(colored(f"Unable to parse score for temperature {temp}. Received: {score_text}", "red"))
+                        scores[temp] = 0
 
             if self.auto_select:
-                best_temp = max(scores, key=scores.get)
+                best_temp = max(scores, key=scores.get, default=self.default_temp)
+                chosen_output = outputs.get(best_temp, "No valid outputs generated.")
                 print(colored(f"Automatically selected output from Temperature {best_temp}:", "green"))
-                print(colored(outputs[best_temp], "white"))
-                return outputs[best_temp]
+                print(colored(chosen_output, "white"))
+                return chosen_output
             else:
                 chosen_output = self.present_options_to_user(outputs)
                 return chosen_output
 
 if __name__ == "__main__":
-    api_key = "your-api-key"  # Replace with your actual OpenAI API key
-    agent = AutoTemp(api_key, auto_select=True)  # Set auto_select to False if you want manual selection
+    agent = AutoTemp(auto_select=True)  # Set auto_select to False if you want manual selection
     prompt = "Code a simple new innovative video game that I can play in browser"
     final_output = agent.run(prompt)
     print(final_output)