Update app.py

app.py

@@ -41,43 +41,45 @@ def query_gpt4o_mini(query, images, api_key):
     """Calls OpenAI's GPT-4o-mini with the query and image data."""
 
     if api_key and api_key.startswith("sk"):
+        try:
+            from openai import OpenAI
         
-        from openai import OpenAI
-    
-        base64_images = [encode_image_to_base64(image[0]) for image in images]
-        client = OpenAI(api_key=api_key)
-        PROMPT = """
-        You are a smart assistant designed to answer questions about a PDF document.
-        You are given relevant information in the form of PDF pages. Use them to construct a short response to the question, and cite your sources (page numbers, etc).
-        If it is not possible to answer using the provided pages, do not attempt to provide an answer and simply say the answer is not present within the documents.
-        Give detailed and extensive answers, only containing info in the pages you are given.
-        You can answer using information contained in plots and figures if necessary.
-        Answer in the same language as the query.
+            base64_images = [encode_image_to_base64(image[0]) for image in images]
+            client = OpenAI(api_key=api_key.strip())
+            PROMPT = """
+            You are a smart assistant designed to answer questions about a PDF document.
+            You are given relevant information in the form of PDF pages. Use them to construct a short response to the question, and cite your sources (page numbers, etc).
+            If it is not possible to answer using the provided pages, do not attempt to provide an answer and simply say the answer is not present within the documents.
+            Give detailed and extensive answers, only containing info in the pages you are given.
+            You can answer using information contained in plots and figures if necessary.
+            Answer in the same language as the query.
+            
+            Query: {query}
+            PDF pages:
+            """
         
-        Query: {query}
-        PDF pages:
-        """
-    
-        response = client.chat.completions.create(
-        model="gpt-4o-mini",
-        messages=[
-            {
-              "role": "user",
-              "content": [
+            response = client.chat.completions.create(
+            model="gpt-4o-mini",
+            messages=[
                 {
-                  "type": "text",
-                  "text": PROMPT.format(query=query)
-                }] + [{
-                  "type": "image_url",
-                  "image_url": {
-                    "url": f"data:image/jpeg;base64,{im}"
-                    },
-                } for im in base64_images]
-            }
-          ],
-          max_tokens=500,
-        )
-        return response.choices[0].message.content
+                  "role": "user",
+                  "content": [
+                    {
+                      "type": "text",
+                      "text": PROMPT.format(query=query)
+                    }] + [{
+                      "type": "image_url",
+                      "image_url": {
+                        "url": f"data:image/jpeg;base64,{im}"
+                        },
+                    } for im in base64_images]
+                }
+              ],
+              max_tokens=500,
+            )
+            return response.choices[0].message.content
+        except Exception as e:
+            return "OpenAI API connection failure. Verify the provided key is correct (sk-***)."
         
     return "Enter your OpenAI API key to get a custom response"