Update ai_generate.py

ai_generate.py

@@ -33,6 +33,17 @@ def generate_groq(text, model):
             response += chunk.choices[0].delta.content or ""
     return response
 
+def generate_openai(text, model):
+    message=[{"role": "user", "content": text}]
+    response = openai_client.chat.completions.create(
+        model=model,
+        messages = message,
+        temperature=0.2,
+        max_tokens=800,
+        frequency_penalty=0.0
+    )
+    return response.choices[0].message
+
 def generate(text, model, api):
     if model == "Llama 3":
         return generate_groq(text, "llama3-70b-8192")
@@ -43,41 +54,20 @@ def generate(text, model, api):
     elif model == "Gemma":
         return generate_groq(text, "gemma2-9b-it")    
     elif model == "OpenAI GPT 3.5":
-        client = OpenAI(
-            api_key=api,
-        )
-        message=[{"role": "user", "content": text}]
-        response = client.chat.completions.create(
-            model="gpt-3.5-turbo",
-            messages = message,
-            temperature=0.2,
-            max_tokens=800,
-            frequency_penalty=0.0
-        )
-        return response[0].message.content
+        try:
+            openai_client = OpenAI(api_key=api)
+            return generate_openai(text, "gpt-3.5-turbo")
+        except:
+            return "Please add a valid API key"
     elif model == "OpenAI GPT 4":
-        client = OpenAI(
-            api_key=api,
-        )
-        message=[{"role": "user", "content": text}]
-        response = client.chat.completions.create(
-            model="gpt-4-turbo",
-            messages = message,
-            temperature=0.2,
-            max_tokens=800,
-            frequency_penalty=0.0
-        )
-        return response.choices[0].message
+        try:
+            openai_client = OpenAI(api_key=api)
+            return generate_openai(text, "gpt-4-turbo")
+        except:
+            return "Please add a valid API key"
     elif model == "OpenAI GPT 4o":
-        client = OpenAI(
-            api_key=api,
-        )
-        message=[{"role": "user", "content": text}]
-        response = client.chat.completions.create(
-            model="gpt-4o",
-            messages = message,
-            temperature=0.2,
-            max_tokens=800,
-            frequency_penalty=0.0
-        )
-        return response.choices[0].message
+        try:
+            openai_client = OpenAI(api_key=api)
+            return generate_openai(text, "gpt-4o")
+        except:
+            return "Please add a valid API key"