minor cleanup: sorted imports, doc string

chatgpt.py

@@ -1,23 +1,31 @@
-import os
 import json
+import os
 from pathlib import Path
 
-import requests
-
 import openai
+import requests
 
-openai.organization = "org-pyHC73IrTntMeKFF1ogHPSb4"
+openai.organization = os.getenv("OPENAI_ORGANIZATION")
 openai.api_key = os.getenv("OPENAI_API_KEY")
-openai.Model.list()
+# openai.Model.list()
 
 
 URL = "https://api.openai.com/v1/chat/completions"
 
 
-def ask_gpt(prompt, n_calls, save_path_base_filename):
+def ask_gpt(prompt, model, n_calls, save_path_base_filename):
+    """make a request to the openai api with a prompt for a given model
+
+    Args:
+        prompt (str): prompt to pass to chatgpt
+        model (str): "gpt-3.5-turbo" or "gpt-4"
+            ref for latest options: https://platform.openai.com/docs/models/gpt-4
+        n_calls (int): number of times to call the openai api
+        save_path_base_filename (str): where to save the output
+    """
 
     payload = {
-    "model": "gpt-3.5-turbo",
+    "model": model,
     "messages": [{"role": "user", "content": prompt}],
     "temperature" : 1.0,
     "top_p":1.0,
@@ -45,10 +53,12 @@ def ask_gpt(prompt, n_calls, save_path_base_filename):
 
 if __name__ == "__main__":
 
+    model = "gpt-3.5-turbo"
     prompt = "Write 10 haikus about cats"
     n_calls = 10
     base_filename = "cat_haikus_10"
     save_path = "./data/gpt_generated/cat_haikus"
+
     Path(save_path).mkdir(parents=True, exist_ok=True)
 
     save_path_base_filename = f"{save_path}/{base_filename}"