Update app.py

app.py

@@ -1,14 +1,16 @@
 import gradio as gr
 import requests
 import os
+import time
+import threading
+import queue
 
-os.environ.get("TIMEOUT_KEEP_ALIVE")
-
-def generate_email(company, recipient_name, product):
-    api_key = os.environ.get("BASETEN_API_KEY")
+def generate_email_async(company, recipient_name, product, progress=gr.Progress()):
+    api_key = os.environ.get("API_KEY")
     if not api_key:
-        return "Error: API_KEY not found in environment variables"
-    
+        yield "Error: API_KEY not found in environment variables"
+        return
+
     url = "https://model-4w56xy7q.api.baseten.co/development/predict"
     
     headers = {
@@ -21,20 +23,50 @@ def generate_email(company, recipient_name, product):
         "recipient_name": recipient_name,
         "product_name": product
     }
-    
-    response = requests.post(url, headers=headers, json=data)
-    
-    if response.status_code == 200:
-        result = response.json()
-        # Extract the part after [/INST] from the response
-        email_content = result['response'].split('[/INST]')[-1].strip()
-        return email_content
+
+    yield "Initiating email generation..."
+    progress(0.1)
+
+    def make_request(result_queue):
+        try:
+            response = requests.post(url, headers=headers, json=data, timeout=300)  # 5 minutes timeout
+            if response.status_code == 200:
+                result = response.json()
+                email_content = result['response'].split('[/INST]')[-1].strip()
+                result_queue.put(("success", email_content))
+            else:
+                result_queue.put(("error", f"Error: {response.status_code} - {response.text}"))
+        except requests.Timeout:
+            result_queue.put(("error", "Error: Request timed out. Please try again later."))
+        except Exception as e:
+            result_queue.put(("error", f"Error: An unexpected error occurred - {str(e)}"))
+
+    result_queue = queue.Queue()
+    thread = threading.Thread(target=make_request, args=(result_queue,))
+    thread.start()
+
+    start_time = time.time()
+    while thread.is_alive():
+        elapsed_time = time.time() - start_time
+        if elapsed_time > 300:  # 5 minutes timeout
+            yield "Error: Process took too long. Please try again later."
+            return
+        yield f"Still working... (Elapsed time: {elapsed_time:.0f} seconds)"
+        progress(min(0.1 + (elapsed_time / 300) * 0.8, 0.9))  # Progress up to 90%
+        time.sleep(5)  # Update every 5 seconds
+
+    thread.join()
+    status, result = result_queue.get()
+
+    if status == "success":
+        progress(1.0)
+        yield f"Email generated successfully:\n\n{result}"
     else:
-        return f"Error: {response.status_code} - {response.text}"
+        yield result
 
 # Create the Gradio interface
 iface = gr.Interface(
-    fn=generate_email,
+    fn=generate_email_async,
     inputs=[
         gr.Textbox(label="Company Name"),
         gr.Textbox(label="Recipient Name"),
@@ -42,8 +74,10 @@ iface = gr.Interface(
     ],
     outputs=gr.Textbox(label="Generated Email"),
     title="AI Email Generator",
-    description="Generate personalized emails using AI"
+    description="Generate personalized emails using AI (This process may take several minutes)",
+    allow_flagging="never",
+    live=True
 )
 
 # Launch the interface
-iface.launch()
+iface.queue().launch()