Update app.py

app.py

@@ -28,7 +28,7 @@ def generate_response(endpoint, api_key, model, user_prompt):
             system_prompt_path = __file__.replace("app.py", "json.txt")
         else:
             print("Invalid API key")
-            return "Invalid API key"
+            return "Invalid API key", "Invalid API key", None
     else:
         if model in claude_models:
             # Set the Anthropic API key
@@ -45,24 +45,32 @@ def generate_response(endpoint, api_key, model, user_prompt):
 
     if model in claude_models:
         # Generate a response using the selected Anthropic model
-        response = client.messages.create(
-            system=system_prompt,
-            messages=[{"role": "user", "content": user_prompt}],
-            model=model,
-            max_tokens=4096
-        )
-        response_text = response.content[0].text
+        try:
+            response = client.messages.create(
+                system=system_prompt,
+                messages=[{"role": "user", "content": user_prompt}],
+                model=model,
+                max_tokens=4096
+            )
+            response_text = response.content[0].text
+        except Exception as e:
+            print(e)
+            response_text = "An error occurred while generating the response. Check that your API key is correct!"
     else:
-        # Generate a response using the selected OpenAI model
-        response = client.chat.completions.create(
-            model=model,
-            messages=[
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": user_prompt}
-            ],
-            max_tokens=4096
-        )
-        response_text = response.choices[0].message.content
+        try:
+            # Generate a response using the selected OpenAI model
+            response = client.chat.completions.create(
+                model=model,
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": user_prompt}
+                ],
+                max_tokens=4096
+            )
+            response_text = response.choices[0].message.content
+        except Exception as e:
+            print(e)
+            response_text = "An error occurred while generating the response. Check that your API key is correct!"
 
     json_string, json_json = extract_json(response_text)
     json_file = json_string if json_string else None
@@ -116,7 +124,7 @@ def generate_second_response(endpoint, api_key, model, generated_output):
             system_prompt_path = __file__.replace("app.py", "diffusion.txt")
         else:
             print("Invalid API key")
-            return "Invalid API key"
+            return "Invalid API key", "Invalid API key", None
     else:
         if model in claude_models:
             # Set the Anthropic API key
@@ -132,25 +140,33 @@ def generate_second_response(endpoint, api_key, model, generated_output):
         system_prompt = file.read()
 
     if model in claude_models:
-        # Generate a second response using the selected Anthropic model and the previously generated output
-        response = client.messages.create(
-            system=system_prompt,
-            messages=[{"role": "user", "content": generated_output}],
-            model=model,
-            max_tokens=4096
-        )
-        response_text = response.content[0].text
+        try:
+            # Generate a second response using the selected Anthropic model and the previously generated output
+            response = client.messages.create(
+                system=system_prompt,
+                messages=[{"role": "user", "content": generated_output}],
+                model=model,
+                max_tokens=4096
+            )
+            response_text = response.content[0].text
+        except Exception as e:
+            print(e)
+            response_text = "An error occurred while generating the response. Check that your API key is correct!"
     else:
-        # Generate a response using the selected OpenAI model
-        response = client.chat.completions.create(
-            model=model,
-            messages=[
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": generated_output}
-            ],
-            max_tokens=4096
-        )
-        response_text = response.choices[0].message.content
+        try:
+            # Generate a response using the selected OpenAI model
+            response = client.chat.completions.create(
+                model=model,
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": generated_output}
+                ],
+                max_tokens=4096
+            )
+            response_text = response.choices[0].message.content
+        except Exception as e:
+            print(e)
+            response_text = "An error occurred while generating the response. Check that your API key is correct!"
 
     return response_text