app.py

from dataclasses import dataclass
from enum import Enum
from typing import Optional, Dict, Any
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 json
from datetime import datetime

# Load environment variables
load_dotenv()

class ConnectionStatus(Enum):
    PENDING = "pending"
    ACTIVE = "active"
    FAILED = "failed"
    NOT_FOUND = "not_found"

@dataclass
class APIResponse:
    success: bool
    data: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    
    def to_json(self) -> str:
        return json.dumps({
            "success": self.success,
            "data": self.data,
            "error": self.error
        })

class CalendarService:
    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: Dict[str, Dict[str, Any]] = {}
        
    def initiate_connection(self, entity_id: str, redirect_url: Optional[str] = None) -> APIResponse:
        try:
            if not redirect_url:
                redirect_url = "https://yourwebsite.com/connection/success"
                
            connection_request = self.toolset.initiate_connection(
                entity_id=entity_id,
                app=App.GOOGLECALENDAR
            )
            
            self.connections[entity_id] = {
                'status': ConnectionStatus.PENDING.value,
                'redirect_url': connection_request.redirectUrl,
                'created_at': datetime.now().isoformat()
            }
            
            return APIResponse(
                success=True,
                data={
                    'status': ConnectionStatus.PENDING.value,
                    'redirect_url': connection_request.redirectUrl,
                    'wait_time': 60,
                    'message': "Please authenticate using the provided link."
                }
            )
            
        except Exception as e:
            return APIResponse(
                success=False,
                error=f"Failed to initiate connection: {str(e)}"
            )
    
    def check_status(self, entity_id: str) -> APIResponse:
        try:
            if entity_id not in self.connections:
                return APIResponse(
                    success=False,
                    error="No connection found for this entity ID"
                )
            
            # In a real implementation, you would check the actual status
            # For now, we'll simulate the status check
            connection = self.connections[entity_id]
            
            return APIResponse(
                success=True,
                data={
                    'status': connection['status'],
                    'message': f"Connection status: {connection['status']}"
                }
            )
            
        except Exception as e:
            return APIResponse(
                success=False,
                error=f"Failed to check status: {str(e)}"
            )
    
    def generate_wrapped(self, entity_id: str) -> APIResponse:
        try:
            if entity_id not in self.connections:
                return APIResponse(
                    success=False,
                    error="Please authenticate first by initiating a connection."
                )
            
            if self.connections[entity_id]['status'] != ConnectionStatus.ACTIVE.value:
                return APIResponse(
                    success=False,
                    error="Connection not active. Please complete the authentication process."
                )
            
            tools = self.toolset.get_tools([
                Action.GOOGLECALENDAR_GET_CALENDAR,
                Action.GOOGLECALENDAR_FIND_EVENT,
                Action.GOOGLECALENDAR_GET_CURRENT_DATE_TIME,
                Action.GOOGLECALENDAR_LIST_CALENDARS
            ])
            
            agent = FunctionCallingAgentWorker(
                tools=tools,
                llm=self.llm,
                prefix_messages=[
                    ChatMessage(
                        role="system",
                        content="""
                        Generate a creative Google Calendar Wrapped summary including:
                        1. Total meetings this year
                        2. Overall time spent in meetings
                        3. Busiest month
                        4. Busiest day
                        5. Most frequent meeting participant
                        6. Average meeting duration
                        7. Most common meeting time
                        
                        Be entertaining and witty with the analysis. Conclude by assigning 
                        a Greek god persona based on their meeting patterns.
                        """
                    )
                ],
                max_function_calls=10,
                allow_parallel_tool_calls=False,
                verbose=True
            ).as_agent()
            
            response = agent.chat(
                f"Create a Calendar Wrapped summary for user with entity_id: {entity_id}. "
                "Use the available Google Calendar actions to gather real data about the user's calendar usage."
            )
            
            return APIResponse(
                success=True,
                data={
                    'wrapped_summary': str(response)
                }
            )
            
        except Exception as e:
            return APIResponse(
                success=False,
                error=f"Failed to generate wrapped: {str(e)}"
            )

def create_gradio_api():
    service = CalendarService()
    
    def handle_connection(entity_id: str, redirect_url: Optional[str] = None) -> str:
        response = service.initiate_connection(entity_id, redirect_url)
        return response.to_json()
    
    def check_status(entity_id: str) -> str:
        response = service.check_status(entity_id)
        return response.to_json()
    
    def generate_wrapped(entity_id: str) -> str:
        response = service.generate_wrapped(entity_id)
        return response.to_json()
    
    # Create API endpoints
    connection_api = gr.Interface(
        fn=handle_connection,
        inputs=[
            gr.Textbox(label="Entity ID"),
            gr.Textbox(label="Redirect URL", placeholder="https://yourwebsite.com/connection/success")
        ],
        outputs=gr.JSON(),
        title="Initialize Calendar Connection",
        description="Start a new calendar connection for an entity",
        examples=[["user123", "https://example.com/callback"]]
    )
    
    status_api = gr.Interface(
        fn=check_status,
        inputs=gr.Textbox(label="Entity ID"),
        outputs=gr.JSON(),
        title="Check Connection Status",
        description="Check the status of an existing connection",
        examples=[["user123"]]
    )
    
    wrapped_api = gr.Interface(
        fn=generate_wrapped,
        inputs=gr.Textbox(label="Entity ID"),
        outputs=gr.JSON(),
        title="Generate Calendar Wrapped",
        description="Generate a calendar wrapped summary for an entity",
        examples=[["user123"]]
    )
    
    # Combine all interfaces
    api = gr.TabbedInterface(
        [connection_api, status_api, wrapped_api],
        ["Connect", "Check Status", "Generate Wrapped"],
        title="Calendar Wrapped API",
        description="API endpoints for Calendar Wrapped functionality"
    )
    
    return api

if __name__ == "__main__":
    api = create_gradio_api()
    api.launch(server_name="0.0.0.0", server_port=7860)