Update app.py

app.py

@@ -1,3 +1,6 @@
+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
@@ -5,195 +8,216 @@ from llama_index.llms.openai import OpenAI
 from dotenv import load_dotenv
 import gradio as gr
 import os
-import time
-import random
+import json
+from datetime import datetime
 
 # Load environment variables
 load_dotenv()
 
-class CalendarWrappedAPI:
+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 = {}
+        self.connections: Dict[str, Dict[str, Any]] = {}
         
-    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"
-            
+    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
             )
             
-            # Store connection info
             self.connections[entity_id] = {
-                'status': connection_request.connectionStatus,
-                'redirect_url': connection_request.redirectUrl
-            }
-            
-            # Wait for random time between 60-100 seconds
-            wait_time = random.randint(60, 100)
-            
-            # Return redirect URL and wait time immediately
-            initial_response = {
-                'status': 'initiated',
+                'status': ConnectionStatus.PENDING.value,
                 'redirect_url': connection_request.redirectUrl,
-                'wait_time': wait_time,
-                'message': f'Please click the link below to authenticate. Waiting {wait_time} seconds for completion...'
+                'created_at': datetime.now().isoformat()
             }
             
-            
-            # Check final status after waiting
-            final_status = connection_request.connectionStatus  # You would typically check the actual status here
-            self.connections[entity_id]['status'] = 'active'
-            
-            return {
-                'status': final_status,
-                'redirect_url': connection_request.redirectUrl,
-                'message': f'Authentication process completed after {wait_time} seconds. Status: {final_status}'
-            }
+            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 {
-                '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:
-            wait_time = random.randint(60, 100)
-            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(actions=[Action.GOOGLECALENDAR_GET_CALENDAR, 
-                                                Action.GOOGLECALENDAR_FIND_EVENT, 
-                                                Action.GOOGLECALENDAR_GET_CURRENT_DATE_TIME,
-                                                Action.GOOGLECALENDAR_LIST_CALENDARS
-                                               ])
-        
-        prefix_messages = [
-            ChatMessage(
-                role="system",
-                content="""
-                Use the Get calendar action, find event action from jan 1 to current datetime. Based on the results from these two actions create a Google Calendar Wrapped for the user.
-                I want the following things: 1. No. of meetings this year, 2. Overall time spent throughout meetings 3. Busiest month 4. Busiest Day 5. Most meetings held with person
-                2 more details that you choose.
-                
-                Be funny and creative and roast with the inferences drawn from the outputs. Also assign a greek god to the player at the end based on the meetings
-                """
+            return APIResponse(
+                success=False,
+                error=f"Failed to initiate connection: {str(e)}"
             )
-        ]
-        
-        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}. Find actual data about the user using Google Calendar Actions at your disposal and then generate the wrapped.")
-            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")
+    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"
+                )
             
-            # 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)
+            # In a real implementation, you would check the actual status
+            # For now, we'll simulate the status check
+            connection = self.connections[entity_id]
             
-            connect_btn.click(
-                fn=handle_connection,
-                inputs=[entity_id_input, redirect_url_input],
-                outputs=[status_message, redirect_link, connection_status]
+            return APIResponse(
+                success=True,
+                data={
+                    'status': connection['status'],
+                    'message': f"Connection status: {connection['status']}"
+                }
             )
             
-        with gr.Tab("Status Check"):
-            status_entity_id = gr.Textbox(
-                label="Entity ID (Username)",
-                placeholder="Enter your username as entity ID"
+        except Exception as e:
+            return APIResponse(
+                success=False,
+                error=f"Failed to check status: {str(e)}"
             )
-            check_btn = gr.Button("Check Status")
-            status_output = gr.Textbox(label="Connection Status")
+    
+    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."
+                )
             
-            check_btn.click(
-                fn=check_status,
-                inputs=status_entity_id,
-                outputs=status_output
+            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."
             )
             
-        with gr.Tab("Generate Wrapped"):
-            wrapped_entity_id = gr.Textbox(
-                label="Entity ID (Username)",
-                placeholder="Enter your username as entity ID"
+            return APIResponse(
+                success=True,
+                data={
+                    'wrapped_summary': str(response)
+                }
             )
-            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
+        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 demo
+    return api
 
 if __name__ == "__main__":
-    demo = create_gradio_api()
-    demo.launch(share=True)  # Set share=False in production
+    api = create_gradio_api()
+    api.launch(server_name="0.0.0.0", server_port=7860)