Update app.py

app.py

@@ -17,11 +17,14 @@ class CalendarWrappedAPI:
         self.llm = OpenAI(model="gpt-4", api_key=os.getenv('OPENAI_API_KEY'))
         self.connections = {}
         
-    def initiate_connection(self, entity_id):
+    def initiate_connection(self, entity_id, redirect_url=None):
         """Initialize connection using entity_id (username)"""
-
+        if redirect_url is None:
+            redirect_url = "https://yourwebsite.com/connection/success"
+            
         try:
             connection_request = self.toolset.initiate_connection(
+                redirect_url=redirect_url,
                 entity_id=entity_id,
                 app=App.GOOGLECALENDAR
             )
@@ -31,24 +34,27 @@ class CalendarWrappedAPI:
                 'status': connection_request.connectionStatus,
                 'redirect_url': connection_request.redirectUrl
             }
-            return {
-                'redirect_url':f'<a href="{connection_request.redirectUrl}" target="_blank">{connection_request.redirectUrl}</a>'
-            }
+            
             # Wait for random time between 60-100 seconds
             wait_time = random.randint(60, 100)
-            print(f"Waiting for {wait_time} seconds before checking connection status...")
-            time.sleep(wait_time)
             
+            # Return redirect URL and wait time immediately
+            initial_response = {
+                'status': 'initiated',
+                'redirect_url': connection_request.redirectUrl,
+                'wait_time': wait_time,
+                'message': f'Please click the link below to authenticate. Waiting {wait_time} seconds for completion...'
+            }            
             # Check final status after waiting
             final_status = 'active'  # You would typically check the actual status here
             self.connections[entity_id]['status'] = final_status
             
             return {
-                'status': 'success',
-                'redirect_url': f'<a href="{connection_request.redirectUrl}" target="_blank">{connection_request.redirectUrl}</a>',
-                'connection_status': final_status,
-                'wait_time': wait_time
+                'status': final_status,
+                'redirect_url': connection_request.redirectUrl,
+                'message': f'Authentication process completed after {wait_time} seconds. Status: {final_status}'
             }
+            
         except Exception as e:
             return {
                 'status': 'error',
@@ -58,12 +64,17 @@ class CalendarWrappedAPI:
     def check_connection_status(self, entity_id):
         """Check the connection status using entity_id"""
         if entity_id in self.connections:
-            # Add a delay here too for subsequent status checks
             wait_time = random.randint(60, 100)
-            print(f"Waiting for {wait_time} seconds before returning status...")
             time.sleep(wait_time)
-            return self.connections[entity_id]['status']
-        return 'not_found'
+            return {
+                'status': self.connections[entity_id]['status'],
+                'wait_time': wait_time,
+                'message': f'Status checked after {wait_time} seconds'
+            }
+        return {
+            'status': 'not_found',
+            'message': 'No connection found for this entity ID'
+        }
     
     def generate_wrapped(self, entity_id):
         """Generate Calendar Wrapped summary using entity_id"""
@@ -104,16 +115,17 @@ class CalendarWrappedAPI:
 def create_gradio_api():
     api = CalendarWrappedAPI()
     
-    def handle_connection(entity_id):
-        result = api.initiate_connection(entity_id)
-        return {
-            **result,
-            'message': f"Connection initiated. Waited for {result.get('wait_time', 0)} seconds."
-        }
+    def handle_connection(entity_id, redirect_url):
+        result = api.initiate_connection(entity_id, redirect_url)
+        return (
+            result['message'],
+            result['redirect_url'],
+            result.get('status', 'unknown')
+        )
     
     def check_status(entity_id):
-        status = api.check_connection_status(entity_id)
-        return f"Status after waiting: {status}"
+        result = api.check_connection_status(entity_id)
+        return f"Status: {result['status']}\n{result['message']}"
     
     def generate(entity_id):
         return api.generate_wrapped(entity_id)
@@ -121,7 +133,7 @@ def create_gradio_api():
     with gr.Blocks() as demo:
         gr.Markdown("""
         # Google Calendar Wrapped Generator
-        **Note**: After initiating connection, check the connection status and generate your wrapped summary.
+        Connect your Google Calendar and generate your personalized wrapped summary.
         """)
         
         with gr.Tab("Connection"):
@@ -129,13 +141,21 @@ def create_gradio_api():
                 label="Entity ID (Username)",
                 placeholder="Enter your username as entity ID"
             )
+            redirect_url_input = gr.Textbox(
+                label="Redirect URL",
+                placeholder="https://yourwebsite.com/connection/success"
+            )
             connect_btn = gr.Button("Initialize Connection")
-            connection_output = gr.HTML(label="Connection Status")
+            
+            # New outputs for better visibility
+            status_message = gr.Textbox(label="Status Message", interactive=False)
+            redirect_link = gr.HTML(label="Authentication Link")
+            connection_status = gr.Textbox(label="Connection Status", interactive=False)
             
             connect_btn.click(
                 fn=handle_connection,
-                inputs=entity_id_input,
-                outputs=connection_output
+                inputs=[entity_id_input, redirect_url_input],
+                outputs=[status_message, redirect_link, connection_status]
             )
             
         with gr.Tab("Status Check"):
@@ -170,4 +190,4 @@ def create_gradio_api():
 
 if __name__ == "__main__":
     demo = create_gradio_api()
-    demo.launch(share=True)
+    demo.launch(share=True)  # Set share=False in production