Update app.py

app.py

@@ -1,65 +1,185 @@
-import gradio as gr
 from composio_llamaindex import ComposioToolSet, App, Action
 from llama_index.core.agent import FunctionCallingAgentWorker
 from llama_index.core.llms import ChatMessage
 from llama_index.llms.openai import OpenAI
 from dotenv import load_dotenv
-import os 
+import gradio as gr
+import os
+import time
+import random
+
 # Load environment variables
 load_dotenv()
 
-# Initialize ComposioToolSet and OpenAI LLM
-toolset = ComposioToolSet(api_key=os.getenv('COMPOSIO_API_KEY')) 
-tools = toolset.get_tools(apps=[App.GOOGLECALENDAR])
-
-llm = OpenAI(model="gpt-4o", api_key=os.getenv('OPENAI_API_KEY'))
-
-# Set up prefix messages for the agent
-prefix_messages = [
-    ChatMessage(
-        role="system",
-        content=(
-            f"""
-            You are a GOOGLE CALENDAR wrapped generator. Based on the GOOGLE CALENDAR username provided, analyze the user's profile, recent tweets, and engagement data.
-            Be extremely creative and funny about it.
-            Create a personalized "GOOGLE CALENDAR Wrapped" summary highlighting their key insights.
-            Include fields that you deem necessary be creative.
-            Use the tools you have to get the info
-            """
-        ),
-    )
-]
-
-# Initialize the agent
-agent = FunctionCallingAgentWorker(
-    tools=tools,
-    llm=llm,
-    prefix_messages=prefix_messages,
-    max_function_calls=10,
-    allow_parallel_tool_calls=False,
-    verbose=True,
-).as_agent()
-
-def generate_wrapped(username):
-    """
-    Function to generate a "GOOGLE CALENDAR Wrapped" summary based on the GOOGLE CALENDAR username provided by the user.
-    """
-    user_input = f"Create a GOOGLE CALENDAR Wrapped summary for the username: {username}"
-    response = agent.chat(user_input)
-    return response
-
-# Create Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("""### GOOGLE CALENDAR Wrapped Generator
-    Enter a GOOGLE CALENDAR username below to generate your personalized GOOGLE CALENDAR Wrapped summary.
-    """)
-
-    username_input = gr.Textbox(label="GOOGLE CALENDAR Username", placeholder="e.g., elonmusk")
-    output = gr.Textbox(label="Output", placeholder="Your GOOGLE CALENDAR Wrapped summary and Google Sheet link will appear here.", lines=10)
-
-    generate_button = gr.Button("Generate Wrapped")
+class CalendarWrappedAPI:
+    def __init__(self):
+        self.toolset = ComposioToolSet(api_key=os.getenv('COMPOSIO_API_KEY'))
+        self.llm = OpenAI(model="gpt-4", api_key=os.getenv('OPENAI_API_KEY'))
+        self.connections = {}
+        
+    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
+            )
+            
+            # Store connection info
+            self.connections[entity_id] = {
+                'status': connection_request.connectionStatus,
+                'redirect_url': connection_request.redirectUrl
+            }
+            
+            # Print the redirect URL prominently
+            print("\n" + "="*50)
+            print("REDIRECT URL:")
+            print(connection_request.redirectUrl)
+            print("="*50 + "\n")
+            
+            # 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)
+            
+            # 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': connection_request.redirectUrl,
+                'connection_status': final_status,
+                'wait_time': wait_time
+            }
+        except Exception as e:
+            return {
+                'status': 'error',
+                'message': str(e)
+            }
+    
+    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'
+    
+    def generate_wrapped(self, entity_id):
+        """Generate Calendar Wrapped summary using entity_id"""
+        if entity_id not in self.connections:
+            return "Please authenticate first by initiating a connection."
+        
+        if self.connections[entity_id]['status'] != 'active':
+            return "Connection not active. Please complete the authentication process."
+        
+        tools = self.toolset.get_tools(apps=[App.GOOGLECALENDAR])
+        
+        prefix_messages = [
+            ChatMessage(
+                role="system",
+                content="""
+                You are a GOOGLE CALENDAR wrapped generator. Based on the user's calendar data,
+                analyze the events and create a personalized "Calendar Wrapped" summary.
+                Be extremely creative and funny about it. Include interesting statistics and patterns.
+                """
+            )
+        ]
+        
+        agent = FunctionCallingAgentWorker(
+            tools=tools,
+            llm=self.llm,
+            prefix_messages=prefix_messages,
+            max_function_calls=10,
+            allow_parallel_tool_calls=False,
+            verbose=True
+        ).as_agent()
+        
+        try:
+            response = agent.chat(f"Create a Calendar Wrapped summary for user with entity_id: {entity_id}")
+            return response
+        except Exception as e:
+            return f"Error generating wrapped: {str(e)}"
 
-    generate_button.click(fn=generate_wrapped, inputs=username_input, outputs=output)
+def create_gradio_api():
+    api = CalendarWrappedAPI()
+    
+    def handle_connection(entity_id, redirect_url):
+        result = api.initiate_connection(entity_id, redirect_url)
+        return {
+            **result,
+            'message': f"Connection initiated. Waited for {result.get('wait_time', 0)} seconds. Please check the console for the redirect URL."
+        }
+    
+    def check_status(entity_id):
+        status = api.check_connection_status(entity_id)
+        return f"Status after waiting: {status}"
+    
+    def generate(entity_id):
+        return api.generate_wrapped(entity_id)
+    
+    with gr.Blocks() as demo:
+        gr.Markdown("""
+        # Google Calendar Wrapped Generator
+        **Note**: After initiating connection, check the console/terminal for the redirect URL.
+        The status will be checked automatically after 60-100 seconds.
+        """)
+        
+        with gr.Tab("Connection"):
+            entity_id_input = gr.Textbox(
+                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.JSON(label="Connection Status")
+            
+            connect_btn.click(
+                fn=handle_connection,
+                inputs=[entity_id_input, redirect_url_input],
+                outputs=connection_output
+            )
+            
+        with gr.Tab("Status Check"):
+            status_entity_id = gr.Textbox(
+                label="Entity ID (Username)",
+                placeholder="Enter your username as entity ID"
+            )
+            check_btn = gr.Button("Check Status")
+            status_output = gr.Textbox(label="Connection Status")
+            
+            check_btn.click(
+                fn=check_status,
+                inputs=status_entity_id,
+                outputs=status_output
+            )
+            
+        with gr.Tab("Generate Wrapped"):
+            wrapped_entity_id = gr.Textbox(
+                label="Entity ID (Username)",
+                placeholder="Enter your username as entity ID"
+            )
+            generate_btn = gr.Button("Generate Wrapped")
+            wrapped_output = gr.Textbox(label="Wrapped Summary", lines=10)
+            
+            generate_btn.click(
+                fn=generate,
+                inputs=wrapped_entity_id,
+                outputs=wrapped_output
+            )
+    
+    return demo
 
-# Launch the Gradio app
-demo.launch()
+if __name__ == "__main__":
+    demo = create_gradio_api()
+    demo.launch(share=True)  # Set share=False in production