app.py

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 gradio as gr
import os
import time
import random

# Load environment variables
load_dotenv()

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
            }
            
            
            # Return redirect URL and wait time immediately
            initial_response = {
                'status': 'initiated',
                'redirect_url': '<a>'+connection_request.redirectUrl+'</a>',
                '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': 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',
                'message': str(e)
            }
    
    def check_connection_status(self, entity_id):
        """Check the connection status using entity_id"""
        if entity_id in self.connections:
            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"""
        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)}"

def create_gradio_api():
    api = CalendarWrappedAPI()
    
    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):
        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)
    
    with gr.Blocks() as demo:
        gr.Markdown("""
        # Google Calendar Wrapped Generator
        Connect your Google Calendar and generate your personalized wrapped summary.
        """)
        
        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")
            
            # New outputs for better visibility
            redirect_link = gr.HTML(label="Authentication Link")
            
            connect_btn.click(
                fn=handle_connection,
                inputs=[entity_id_input, redirect_url_input],
                outputs=[redirect_link]
            )
            
        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

if __name__ == "__main__":
    demo = create_gradio_api()
    demo.launch(share=True)  # Set share=False in production