Update app.py

app.py

@@ -4,7 +4,7 @@ from typing import Optional, Dict, Any
 from composio_llamaindex import ComposioToolSet, App, Action
 from datetime import datetime, timedelta
 from collections import defaultdict, Counter
-from llama_index.llms.openai import OpenAI
+from llama_index_llms_openai import OpenAI
 import gradio as gr
 import os
 import json
@@ -117,80 +117,36 @@ class CalendarService:
         
         return stats
         
-    def generate_wrapped(self, entity_id: str) -> APIResponse:
+    def initiate_connection(self, entity_id: str, redirect_url: Optional[str] = None) -> APIResponse:
         try:
-            # Get current year's start and end dates
-            current_year = datetime.now().year
-            time_min = f"{current_year},1,1,0,0,0"
-            time_max = f"{current_year},12,31,23,59,59"
+            if not redirect_url:
+                redirect_url = "https://yourwebsite.com/connection/success"
+                
+            connection_request = self.toolset.initiate_connection(
+                entity_id=entity_id,
+                app=App.GOOGLECALENDAR
+            )
             
-            request_params = {
-                "calendar_id": "primary",
-                "timeMin": time_min,
-                "timeMax": time_max,
-                "single_events": True,
-                "max_results": 2500,
-                "order_by": "startTime"
+            self.connections[entity_id] = {
+                'status': ConnectionStatus.PENDING.value,
+                'redirect_url': connection_request.redirectUrl,
+                'created_at': datetime.now().isoformat()
             }
             
-            events_response = self.toolset.execute_action(
-                action=Action.GOOGLECALENDAR_FIND_EVENT,
-                params=request_params,
-                entity_id=entity_id
+            return APIResponse(
+                success=True,
+                data={
+                    'status': ConnectionStatus.PENDING.value,
+                    'redirect_url': connection_request.redirectUrl,
+                    'wait_time': 60,
+                    'message': "Please authenticate using the provided link."
+                }
             )
             
-            if events_response["successfull"]:
-                stats = self.analyze_calendar_events(events_response)
-                
-                # Create a prompt for the LLM with the stats
-                prompt = f"""Based on the following calendar statistics, analyze which tech billionaire this person's schedule most resembles and provide brief comments for each metric:
-
-Stats:
-- Total meetings this year: {stats['total_meetings_this_year']}
-- Total time in meetings: {stats['total_time_spent']}
-- Busiest month: {stats['busiest_month']}
-- Busiest day: {stats['busiest_day']}
-- Average meeting duration: {stats['average_meeting_duration']}
-- Most common meeting time: {stats['most_common_meeting_time']}
-- Most frequent collaborator: {stats['most_frequent_participant']}
-
-Please provide:
-1. Which tech billionaire's schedule this most resembles and why
-2. A one-sentence comment for each of the above metrics
-Format your response as JSON with keys: 'billionaire_match' and 'metric_comments'
-
-Dont make any extra comments before or after. Just give the required output and shut up. Dont send any sentences saying like here's your response, or the output is generated"""
-
-                llm_analysis = llm.complete(prompt)
-                
-                try:
-                    llm_json = json.loads(llm_analysis)
-                except json.JSONDecodeError:
-                    llm_json = {
-                        "billionaire_match": "Analysis unavailable",
-                        "metric_comments": "Comments unavailable"
-                    }
-                
-                # Add LLM analysis to stats
-                stats.update({
-                    "schedule_analysis": llm_json["billionaire_match"],
-                    "metric_insights": llm_json["metric_comments"]
-                })
-                
-                return APIResponse(
-                    success=True,
-                    data=stats
-                )
-            else:
-                return APIResponse(
-                    success=False,
-                    error=events_response["error"] or "Failed to fetch calendar events"
-                )
-            
         except Exception as e:
             return APIResponse(
                 success=False,
-                error=f"Failed to generate wrapped: {str(e)}"
+                error=f"Failed to initiate connection: {str(e)}"
             )
     
     def check_status(self, entity_id: str) -> APIResponse:
@@ -218,50 +174,92 @@ Dont make any extra comments before or after. Just give the required output and
             )
     
     def generate_wrapped(self, entity_id: str) -> APIResponse:
-        try:
-            # Get current year's start and end dates
-            current_year = datetime.now().year
-            time_min = f"{current_year},1,1,0,0,0"
-            time_max = f"{current_year},12,31,23,59,59"
+    try:
+        # Get current year's start and end dates
+        current_year = datetime.now().year
+        time_min = f"{current_year},1,1,0,0,0"
+        time_max = f"{current_year},12,31,23,59,59"
+        
+        request_params = {
+            "calendar_id": "primary",
+            "timeMin": time_min,
+            "timeMax": time_max,
+            "single_events": True,
+            "max_results": 2500,
+            "order_by": "startTime"
+        }
+        
+        events_response = self.toolset.execute_action(
+            action=Action.GOOGLECALENDAR_FIND_EVENT,
+            params=request_params,
+            entity_id=entity_id
+        )
+        
+        if events_response["successfull"]:
+            stats = self.analyze_calendar_events(events_response)
             
-            request_params = {
-                "calendar_id": "primary",
-                "timeMin": time_min,
-                "timeMax": time_max,
-                "single_events": True,
-                "max_results": 2500,
-                "order_by": "startTime"
-            }
+            # Get tech billionaire comparison
+            billionaire_prompt = f"""Based on these calendar stats, which tech billionaire's schedule does this most resemble and why?
+            Stats:
+            - {stats['total_meetings_this_year']} total meetings
+            - {stats['total_time_spent']} total time in meetings
+            - Most active on {stats['busiest_day']}s
+            - Busiest month is {stats['busiest_month']}
+            - Average meeting duration: {stats['average_meeting_duration']}
             
-            events_response = self.toolset.execute_action(
-                action=Action.GOOGLECALENDAR_FIND_EVENT,
-                params=request_params,
-                entity_id=entity_id
-            )
-            llm_response = llm.complete(f"{str(events_response)} is the event response. Based on this what tech billionaire are they most similar to and why??")
-            if events_response["successfull"]:
-                stats = self.analyze_calendar_events(events_response)
-                return APIResponse(
-                    success=True,
-                    data=stats
-                )
-            else:
-                return APIResponse(
-                    success=False,
-                    error=events_response["error"] or "Failed to fetch calendar events"
-                )
+            Return as JSON with format: {{"name": "billionaire name", "reason": "explanation"}}
+            """
+            
+            # Get comments for each stat
+            stats_prompt = f"""Analyze these calendar stats and write a brief, insightful one-sentence comment for each metric:
+            - Total meetings: {stats['total_meetings_this_year']}
+            - Total time in meetings: {stats['total_time_spent']}
+            - Busiest month: {stats['busiest_month']}
+            - Busiest day: {stats['busiest_day']}
+            - Average meeting duration: {stats['average_meeting_duration']}
+            - Most common meeting time: {stats['most_common_meeting_time']}
+            - Most frequent participant: {stats['most_frequent_participant']}
+
+            Return as JSON with format: {{"total_meetings_comment": "", "time_spent_comment": "", "busiest_times_comment": "", "collaborator_comment": "", "habits_comment": ""}}
+            """
+            
+            # Make LLM calls
+            try:
+                billionaire_response = json.loads(llm.complete(billionaire_prompt))
+                stats_comments = json.loads(llm.complete(stats_prompt))
+                
+                # Add new fields to stats
+                stats["schedule_analysis"] = billionaire_response
+                stats["metric_insights"] = stats_comments
+            except (json.JSONDecodeError, Exception) as e:
+                print(f"Error processing LLM responses: {e}")
+                # Add empty defaults if LLM processing fails
+                stats["schedule_analysis"] = {"name": "Unknown", "reason": "Analysis unavailable"}
+                stats["metric_insights"] = {"total_meetings_comment": "", "time_spent_comment": "", 
+                                          "busiest_times_comment": "", "collaborator_comment": "", 
+                                          "habits_comment": ""}
             
-        except Exception as e:
+            return APIResponse(
+                success=True,
+                data=stats
+            )
+        else:
             return APIResponse(
                 success=False,
-                error=f"Failed to generate wrapped: {str(e)}"
+                error=events_response["error"] or "Failed to fetch calendar events"
             )
+        
+    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.toolset.initiate_connection(entity_id, redirect_url)
+        response = service.initiate_connection(entity_id, redirect_url)
         return response.to_json()
     
     def check_status(entity_id: str) -> str: