Upload schedule_ai.py

src/schedule_ai.py

@@ -0,0 +1,214 @@
+import google.generativeai as genai
+import os
+from dotenv import load_dotenv
+import re
+import joblib
+import pandas as pd
+import numpy as np
+from sklearn.preprocessing import StandardScaler
+
+class SimpleScheduler:
+    def __init__(self):
+        load_dotenv()
+        api_key = os.getenv('GEMINI_API_KEY')
+        genai.configure(api_key=api_key)
+        self.model = genai.GenerativeModel('gemini-2.0-flash-exp')
+        
+        # Load ML model dan scaler
+        try:
+            self.rf_model = joblib.load('Models/random_forest_model.pkl')
+            self.scaler = joblib.load('Models/scaler.pkl')
+        except Exception as e:
+            print(f"Warning: Could not load ML models: {str(e)}")
+            self.rf_model = None
+            self.scaler = None
+
+    def generate_schedule(self, duration, goals, available_hours, considerations):
+        """Generate schedule using Gemini with improved formatting"""
+        duration_lower = duration.lower()
+        
+        # Select appropriate prompt based on duration
+        if "year" in duration_lower:
+            prompt = self._create_yearly_prompt(goals, available_hours, considerations)
+        elif "month" in duration_lower:
+            prompt = self._create_monthly_prompt(goals, available_hours, considerations)
+        else:  # weekly
+            prompt = self._create_weekly_prompt(goals, available_hours, considerations)
+
+        try:
+            response = self.model.generate_content(prompt)
+            return self._format_output(response.text, duration, goals, available_hours, considerations)
+        except Exception as e:
+            raise Exception(f"Error generating schedule: {str(e)}")
+
+    def _get_priority(self, hour):
+        """Determine priority based on hour of day"""
+        if hour < 12:
+            return "HIGH"
+        elif hour < 15:
+            return "MEDIUM"
+        else:
+            return "LOW"
+
+    def _format_time(self, time_str):
+        """Format time string consistently"""
+        # Remove any spaces and AM/PM
+        time_str = time_str.replace(" ", "").upper()
+        
+        # Convert to 24-hour format if needed
+        if "AM" in time_str or "PM" in time_str:
+            hour = int(time_str.split(":")[0])
+            if "PM" in time_str and hour != 12:
+                hour += 12
+            return f"{hour:02d}:00"
+        
+        return time_str
+
+    def _format_output(self, schedule, duration, goals, available_hours, considerations):
+        """Format the output with clean and consistent structure"""
+        
+        # Create header
+        header = f"""SCHEDULE OVERVIEW
+-------------------------------------------
+Duration: {duration}
+Goals: {goals}
+Available Hours: {available_hours} hours/day
+Considerations: {considerations}
+-------------------------------------------\n"""
+
+        # Process schedule content
+        formatted_schedule = ""
+        current_section = None
+        
+        for line in schedule.strip().split('\n'):
+            line = line.strip()
+            
+            # Skip unnecessary lines
+            if not line or "[Continue" in line or "rest" in line.lower():
+                continue
+                
+            # Process QUARTER sections for yearly schedule
+            if "QUARTER" in line.upper():
+                formatted_schedule += f"\n{line}\n{'=' * len(line)}\n"
+                continue
+                
+            # Process Week sections
+            elif "WEEK" in line.upper() and "-" in line:
+                week_info = line.replace('**', '').strip()
+                formatted_schedule += f"\n{week_info}\n{'-' * len(week_info)}\n"
+                continue
+                
+            # Process day headers for weekly schedule
+            elif any(day in line.upper() for day in ['MONDAY', 'TUESDAY', 'WEDNESDAY', 'THURSDAY', 'FRIDAY', 'SATURDAY', 'SUNDAY']):
+                day = line.split(':')[0].strip()
+                formatted_schedule += f"\n{day}:\n"
+                continue
+            
+            # Process time blocks
+            elif ':' in line:
+                try:
+                    # Split into time and task
+                    parts = line.split(':', 1)
+                    if len(parts) < 2:
+                        continue
+                        
+                    time_str = self._format_time(parts[0])
+                    task = parts[1].strip()
+                    
+                    # Clean up task text
+                    task = task.split('[')[0].strip()  # Remove existing priority if any
+                    
+                    # Get hour for priority
+                    hour = int(time_str.split(':')[0])
+                    priority = self._get_priority(hour)
+                    
+                    # Format schedule entry
+                    formatted_schedule += f"  {time_str}: {task} [{priority}]\n"
+                except Exception:
+                    # If processing fails, add line as is
+                    formatted_schedule += f"  {line}\n"
+            
+            # Add other lines that might be important
+            elif line.strip() and not line.startswith('*'):
+                formatted_schedule += f"{line}\n"
+
+        return header + formatted_schedule
+
+    def _create_weekly_prompt(self, goals, available_hours, considerations):
+        return f"""
+        Create a WEEKLY schedule for: {goals}
+        Available Hours: {available_hours}/day
+        Consider: {considerations}
+        
+        Format exactly as follows:
+        
+        MONDAY:
+        09:00: [Morning activity]
+        10:00: [Morning activity]
+        11:00: [Morning activity]
+        12:00: Lunch Break
+        13:00: [Afternoon activity]
+        14:00: [Afternoon activity]
+        15:00: [Late activity]
+        16:00: [Break/Gym if needed]
+        
+        TUESDAY:
+        [Same format]
+        
+        Continue for each day including SATURDAY and SUNDAY.
+        Use 24-hour format (09:00, 14:00).
+        Keep descriptions brief and specific.
+        Include breaks and considerations.
+        """
+
+    def _create_monthly_prompt(self, goals, available_hours, considerations):
+        return f"""
+        Create a MONTHLY schedule for: {goals}
+        Available Hours: {available_hours}/day
+        Consider: {considerations}
+        
+        Format as follows:
+        
+        WEEK 1:
+        Monday-Friday:
+        09:00: [Morning activity]
+        10:00: [Morning activity]
+        11:00: [Morning activity]
+        12:00: Lunch Break
+        13:00: [Afternoon activity]
+        14:00: [Afternoon activity]
+        15:00: [Late activity]
+        16:00: [Break/Gym if needed]
+        
+        WEEK 2:
+        [Same format]
+        
+        Continue for 4 weeks.
+        Use 24-hour format.
+        Keep descriptions brief and specific.
+        Include breaks and considerations.
+        """
+
+    def _create_yearly_prompt(self, goals, available_hours, considerations):
+        return f"""
+        Create a YEARLY schedule for: {goals}
+        Available Hours: {available_hours}/day
+        Consider: {considerations}
+        
+        Format exactly as:
+        
+        QUARTER 1 (January-March)
+        
+        Week 1-4 - [Phase Name]
+        09:00: [Primary morning activity]
+        11:00: [Late morning activity]
+        14:00: [Afternoon activity]
+        
+        Week 5-8 - [Next Phase]
+        [Same format]
+        
+        Continue for all quarters.
+        Use 24-hour format.
+        Focus on major milestones.
+        Keep descriptions brief and specific.
+        """