Update app.py

app.py

@@ -0,0 +1,636 @@
+import gradio as gr
+import torch
+import datetime
+import pytz
+import uuid
+import re
+import json
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from google.oauth2 import service_account
+from googleapiclient.discovery import build
+import os
+import gc
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Log startup
+logger.info("Starting appointment booking application...")
+
+# Set up timezone
+IST = pytz.timezone('Asia/Kolkata')
+
+# ===== CONFIGURATION =====
+
+# Model ID on Hugging Face
+MODEL_ID = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+
+# Google Calendar API Configuration
+SCOPES = ['https://www.googleapis.com/auth/calendar']
+SERVICE_ACCOUNT_FILE = 'service-account-key.json'
+CALENDAR_ID = '26f5856049fab3d6648a2f1dea57c70370de6bc1629a5182be1511b0e75d11d3@group.calendar.google.com'  # Update with your calendar ID if not using primary
+
+# Local appointments database (for backup)
+appointments_db = {}
+
+# ===== GOOGLE CALENDAR FUNCTIONS =====
+
+def get_calendar_service():
+    """Get Google Calendar service"""
+    try:
+        # Check if Google credentials are stored in env variable
+        google_credentials = os.environ.get('GOOGLE_CREDENTIALS')
+        if google_credentials:
+            logger.info("Using Google credentials from environment variable")
+            # Write the credentials to a temporary file
+            with open('temp_credentials.json', 'w') as f:
+                f.write(google_credentials)
+            temp_file_path = 'temp_credentials.json'
+            credentials = service_account.Credentials.from_service_account_file(
+                temp_file_path, scopes=SCOPES)
+        elif os.path.exists(SERVICE_ACCOUNT_FILE):
+            logger.info(f"Using Google credentials from file: {SERVICE_ACCOUNT_FILE}")
+            # Use the file on disk
+            credentials = service_account.Credentials.from_service_account_file(
+                SERVICE_ACCOUNT_FILE, scopes=SCOPES)
+        else:
+            logger.warning("No Google Calendar credentials found")
+            return None
+
+        service = build('calendar', 'v3', credentials=credentials)
+        return service
+    except Exception as e:
+        logger.error(f"Error getting calendar service: {e}")
+        return None
+
+def add_to_google_calendar(appointment_details):
+    """Add an appointment to Google Calendar"""
+    try:
+        service = get_calendar_service()
+        if not service:
+            return None
+
+        # Format start and end time
+        date_str = appointment_details["date"]
+        time_str = appointment_details["time"]
+        
+        # Parse date and time
+        date_parts = date_str.split('-')
+        year, month, day = int(date_parts[0]), int(date_parts[1]), int(date_parts[2])
+
+        time_parts = time_str.split(' ')
+        time_val = time_parts[0]
+        meridian = time_parts[1] if len(time_parts) > 1 else 'AM'
+
+        hours, minutes = map(int, time_val.split(':'))
+
+        if meridian.upper() == 'PM' and hours != 12:
+            hours += 12
+        if meridian.upper() == 'AM' and hours == 12:
+            hours = 0
+
+        # Create datetime objects
+        start_time = datetime.datetime(year, month, day, hours, minutes, 0, tzinfo=IST)
+        end_time = start_time + datetime.timedelta(hours=1)  # Default 1 hour appointment
+
+        # Create event
+        event = {
+            'summary': f"Appointment with {appointment_details['name']}",
+            'location': 'Office',
+            'description': 'Appointment booked via AI Assistant',
+            'start': {
+                'dateTime': start_time.isoformat(),
+                'timeZone': 'Asia/Kolkata',
+            },
+            'end': {
+                'dateTime': end_time.isoformat(),
+                'timeZone': 'Asia/Kolkata',
+            },
+            'reminders': {
+                'useDefault': False,
+                'overrides': [
+                    {'method': 'email', 'minutes': 24 * 60},
+                    {'method': 'popup', 'minutes': 10},
+                ],
+            },
+        }
+
+        # Add unique ID to track for cancellation
+        appointment_id = appointment_details.get('appointment_id', str(uuid.uuid4()))
+        event['extendedProperties'] = {
+            'private': {
+                'appointment_id': appointment_id
+            }
+        }
+
+        # Insert event
+        created_event = service.events().insert(calendarId=CALENDAR_ID, body=event).execute()
+        return created_event['id']
+
+    except Exception as e:
+        logger.error(f"Error adding to Google Calendar: {e}")
+        return None
+
+# ===== FUNCTION DEFINITIONS =====
+
+function_definitions = [
+  {
+    "name": "book_appointment",
+    "description": "Book an appointment",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "name": {
+          "type": "string",
+          "description": "The name of the person"
+        },
+        "date": {
+          "type": "string",
+          "description": "The date in YYYY-MM-DD format"
+        },
+        "time": {
+          "type": "string",
+          "description": "The time of the appointment (e.g., '10:00 AM')"
+        }
+      },
+      "required": ["name", "date", "time"]
+    }
+  }
+]
+
+# ===== FUNCTION IMPLEMENTATIONS =====
+
+def book_appointment(appointment_details):
+    """Book an appointment with just name, date and time"""
+    try:
+        # Generate a unique appointment ID
+        appointment_id = str(uuid.uuid4())[:8]  # Shorter ID for simplicity
+        
+        # Add appointment ID to details
+        appointment_details['appointment_id'] = appointment_id
+        
+        # Store in local database
+        appointments_db[appointment_id] = appointment_details
+        
+        # Add to Google Calendar
+        calendar_event_id = add_to_google_calendar(appointment_details)
+        
+        if calendar_event_id:
+            # Store the calendar event ID
+            appointments_db[appointment_id]['calendar_event_id'] = calendar_event_id
+            
+            return {
+                "success": True,
+                "appointment_id": appointment_id,
+                "message": "Appointment successfully booked and added to calendar",
+                "details": {
+                    "name": appointment_details["name"],
+                    "date": appointment_details["date"],
+                    "time": appointment_details["time"],
+                    "location": "Office"
+                }
+            }
+        else:
+            return {
+                "success": True,
+                "appointment_id": appointment_id,
+                "message": "Appointment booked but failed to add to calendar (offline mode)",
+                "details": {
+                    "name": appointment_details["name"],
+                    "date": appointment_details["date"],
+                    "time": appointment_details["time"],
+                    "location": "Office"
+                }
+            }
+    except Exception as e:
+        logger.error(f"Error in book_appointment: {e}")
+        return {
+            "success": False,
+            "message": f"Failed to book appointment: {str(e)}"
+        }
+
+# ===== MODEL MANAGEMENT =====
+
+# Global model and tokenizer - SINGLETON PATTERN
+model = None
+tokenizer = None
+
+def free_memory():
+    """Free memory by clearing cache and running garbage collection"""
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        logger.info(f"GPU memory allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
+        logger.info(f"GPU memory reserved: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")
+
+def load_llama_model():
+    """Load the Llama 3.1 model and tokenizer using singleton pattern"""
+    global model, tokenizer
+
+    # If model already loaded, return the existing instances
+    if model is not None and tokenizer is not None:
+        return True
+
+    logger.info("Loading Llama 3.1 model and tokenizer...")
+    free_memory()
+
+    try:
+        # Set up quantization config for better memory efficiency
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_use_double_quant=True
+        )
+
+        # Load tokenizer
+        tokenizer_local = AutoTokenizer.from_pretrained(MODEL_ID)
+        logger.info("Tokenizer loaded successfully")
+
+        # Load model with optimized settings
+        model_local = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            quantization_config=quantization_config,
+            device_map="auto",
+            torch_dtype=torch.float16,
+            low_cpu_mem_usage=True
+        )
+        logger.info("Model loaded successfully")
+
+        # Store in global variables
+        model = model_local
+        tokenizer = tokenizer_local
+
+        free_memory()
+        logger.info("Model and tokenizer initialization complete")
+        return True
+
+    except Exception as e:
+        logger.error(f"Error loading model: {e}")
+        return False
+
+# ===== CHAT PROCESSING =====
+
+def format_prompt_with_functions(messages, system_prompt):
+    """Format the prompt for Llama 3.1 with function definitions"""
+    # Add function definitions to system prompt
+    full_system_prompt = system_prompt + "\n\n"
+    full_system_prompt += "You have access to the following functions that you MUST use for specific user queries:\n"
+
+    for func in function_definitions:
+        full_system_prompt += f"- {func['name']}: {func['description']}\n"
+        full_system_prompt += "  Parameters:\n"
+        for param_name, param_info in func['parameters']['properties'].items():
+            required = "required" if param_name in func['parameters'].get('required', []) else "optional"
+            full_system_prompt += f"    - {param_name} ({required}): {param_info.get('description', '')}\n"
+
+    full_system_prompt += "\nIMPORTANT: When a user asks to book an appointment, you MUST respond using the following JSON format:\n"
+    full_system_prompt += '```json\n{"function_call": {"name": "function_name", "arguments": {"arg1": "value1", "arg2": "value2"}}}\n```\n'
+    full_system_prompt += "You MUST collect all required information first: name, date, and time."
+    full_system_prompt += "\n\nFor non-function-calling queries, respond in a conversational manner."
+
+    # Format conversation history
+    formatted_messages = [
+        {"role": "system", "content": full_system_prompt}
+    ]
+
+    # Add conversation history
+    for message in messages:
+        if message["role"] == "function":
+            # Convert function results to assistant format for Llama 3.1
+            formatted_messages.append({
+                "role": "assistant",
+                "content": f"I'll process the function result: {message['content']}"
+            })
+        else:
+            formatted_messages.append(message)
+
+    return formatted_messages
+
+def extract_function_call(response_text):
+    """Extract function call from model response"""
+    # Look for JSON block in the response
+    json_pattern = r'```json\s*(.*?)\s*```'
+    json_matches = re.findall(json_pattern, response_text, re.DOTALL)
+
+    if not json_matches:
+        # Try alternative pattern without markdown
+        json_pattern = r'({.*"function_call".*})'
+        json_matches = re.findall(json_pattern, response_text, re.DOTALL)
+
+    if json_matches:
+        try:
+            for json_str in json_matches:
+                parsed_json = json.loads(json_str.strip())
+                if "function_call" in parsed_json:
+                    function_call = parsed_json["function_call"]
+                    return {
+                        "id": str(uuid.uuid4()),
+                        "name": function_call["name"],
+                        "arguments": function_call["arguments"]
+                    }
+        except json.JSONDecodeError:
+            logger.error(f"Failed to parse JSON: {json_matches[0]}")
+
+    return None
+
+def safe_generate(inputs, max_new_tokens=512):
+    """Safely generate text with error handling and memory management"""
+    global model, tokenizer
+
+    try:
+        free_memory()
+
+        # Generate with appropriate settings
+        outputs = model.generate(
+            inputs,
+            max_new_tokens=max_new_tokens,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+
+        response_text = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
+        free_memory()
+        return response_text
+    except Exception as e:
+        logger.error(f"Error in generation: {e}")
+        free_memory()
+        return f"Error generating response: {str(e)}"
+
+def process_chat(message, chat_history):
+    """Process a chat message, calling functions when necessary"""
+    global model, tokenizer
+
+    if model is None or tokenizer is None:
+        error_msg = "Model not loaded properly. Please click 'Reload Model' and try again."
+        new_history = chat_history + [(message, error_msg)]
+        return new_history, new_history
+
+    try:
+        # Create system prompt
+        system_prompt = """You are a friendly appointment booking assistant. You help users book appointments by collecting their name, preferred date, and time.
+
+Be polite and helpful. For any appointment request, make sure to collect the person's name, date (in YYYY-MM-DD format), and time (e.g., '10:00 AM').
+
+If the user doesn't specify all the required information, politely ask for the missing details before booking."""
+
+        # Convert Gradio chat history to message format
+        messages = []
+
+        # Limit history to last 3 exchanges to save memory
+        limited_chat_history = chat_history[-3:] if len(chat_history) > 3 else chat_history
+
+        for user_msg, bot_msg in limited_chat_history:
+            messages.append({"role": "user", "content": user_msg})
+            messages.append({"role": "assistant", "content": bot_msg})
+
+        # Add current message
+        messages.append({"role": "user", "content": message})
+
+        # Format messages with function calling info
+        formatted_messages = format_prompt_with_functions(messages, system_prompt)
+
+        # Generate model response with error handling
+        try:
+            inputs = tokenizer.apply_chat_template(
+                formatted_messages,
+                tokenize=True,
+                add_generation_prompt=True,
+                return_tensors="pt"
+            ).to(model.device)
+
+            # First generation
+            response_text = safe_generate(inputs, max_new_tokens=512)
+            logger.info(f"Model response: {response_text[:100]}...")
+
+            # Check if response contains a function call
+            function_call = extract_function_call(response_text)
+            
+            if function_call and function_call["name"] == "book_appointment":
+                # Execute the booking function
+                function_result = book_appointment(function_call["arguments"])
+                logger.info(f"Function result: {json.dumps(function_result)[:200]}...")
+
+                # Add the function result to messages
+                messages.append({
+                    "role": "assistant",
+                    "content": response_text,
+                })
+
+                messages.append({
+                    "role": "function",
+                    "name": "book_appointment",
+                    "content": json.dumps(function_result)
+                })
+
+                # Format messages for second call
+                formatted_messages = format_prompt_with_functions(messages, system_prompt)
+
+                # Generate second response
+                inputs = tokenizer.apply_chat_template(
+                    formatted_messages,
+                    tokenize=True,
+                    add_generation_prompt=True,
+                    return_tensors="pt"
+                ).to(model.device)
+
+                second_response = safe_generate(inputs, max_new_tokens=512)
+                logger.info(f"Second model response: {second_response[:100]}...")
+
+                # Update chat history
+                new_chat_history = chat_history + [(message, second_response)]
+                return new_chat_history, new_chat_history
+            else:
+                # No function call, just return the response
+                new_chat_history = chat_history + [(message, response_text)]
+                return new_chat_history, new_chat_history
+                
+        except Exception as e:
+            logger.error(f"Error in generation: {e}")
+            error_msg = f"Sorry, I couldn't generate a response. Please try a simpler question or try again later."
+            new_chat_history = chat_history + [(message, error_msg)]
+            return new_chat_history, new_chat_history
+    except Exception as e:
+        logger.error(f"Error in process_chat: {e}")
+        error_msg = f"Sorry, I encountered an error. Please try again."
+        new_chat_history = chat_history + [(message, error_msg)]
+        return new_chat_history, new_chat_history
+
+# ===== GRADIO INTERFACE =====
+
+def create_gradio_interface():
+    """Create the Gradio interface for the chatbot"""
+    logger.info("Creating Gradio interface...")
+
+    with gr.Blocks(css="""
+        .gradio-container {max-width: 800px !important}
+        .chat-window {height: 600px !important; overflow-y: auto}
+    """) as demo:
+        gr.Markdown("# Simple Appointment Booking Assistant")
+        gr.Markdown("### Tell me your name, date and time to book an appointment")
+
+        # Model status indicator
+        with gr.Row():
+            model_status = gr.Textbox(
+                label="Model Status",
+                value="Loading model...",
+                interactive=False
+            )
+
+        # Calendar integration status
+        with gr.Row():
+            calendar_status = gr.Textbox(
+                label="Calendar Integration Status",
+                value="Checking Google Calendar integration...",
+                interactive=False
+            )
+
+        # Function to check Google Calendar connectivity
+        def check_calendar_integration():
+            try:
+                service = get_calendar_service()
+                if service:
+                    return "Google Calendar integration is active. Appointments will be saved to calendar."
+                else:
+                    return "Google Calendar integration is not available. Appointments will only be stored in memory."
+            except Exception as e:
+                logger.error(f"Error checking calendar integration: {str(e)}")
+                return f"Error checking calendar integration: {str(e)}"
+
+        # Chatbot interface
+        chatbot = gr.Chatbot(
+            [],
+            elem_id="chatbot",
+            label="Chat with Appointment Assistant",
+            height=500
+        )
+
+        with gr.Row():
+            msg = gr.Textbox(
+                show_label=False,
+                placeholder="Type your message here...",
+                container=False
+            )
+            submit = gr.Button("Send")
+
+        with gr.Row():
+            clear = gr.Button("Clear Conversation")
+            reload_model = gr.Button("Reload Model")
+
+        # Provide instructions
+        with gr.Accordion("Instructions", open=False):
+            gr.Markdown("""
+            ## How to use this appointment booking assistant:
+
+            Simply tell the assistant you want to book an appointment and provide:
+            1. Your name
+            2. The date you want (in YYYY-MM-DD format)
+            3. The time you want (like "10:00 AM")
+
+            ### Example messages:
+            - "I'd like to book an appointment"
+            - "Book an appointment for John Smith on 2025-05-20 at 2:30 PM"
+            - "Can I schedule a meeting tomorrow at 10 AM?"
+            """)
+
+        chat_history = gr.State([])
+
+        def initialize_model():
+            """Initialize the model on app startup"""
+            success = load_llama_model()
+            status = "Model loaded successfully!" if success else "Error loading model. Try clicking 'Reload Model'."
+            cal_status = check_calendar_integration()
+            return status, [], cal_status
+
+        def reload_model_click():
+            """Force reload the model and free memory"""
+            global model, tokenizer
+            # Clear global variables
+            model = None
+            tokenizer = None
+            # Free memory
+            free_memory()
+            # Reload model
+            success = load_llama_model()
+            status = "Model reloaded successfully!" if success else "Error reloading model. Check logs for details."
+            cal_status = check_calendar_integration()
+            return status, [], cal_status
+
+        # Set up event handlers
+        submit.click(
+            process_chat,
+            inputs=[msg, chat_history],
+            outputs=[chatbot, chat_history]
+        ).then(
+            lambda: "",
+            None,
+            msg
+        )
+
+        msg.submit(
+            process_chat,
+            inputs=[msg, chat_history],
+            outputs=[chatbot, chat_history]
+        ).then(
+            lambda: "",
+            None,
+            msg
+        )
+
+        clear.click(
+            lambda: [],
+            inputs=None,
+            outputs=[chat_history]
+        ).then(
+            lambda: [],
+            inputs=None,
+            outputs=[chatbot]
+        )
+
+        reload_model.click(
+            reload_model_click,
+            inputs=None,
+            outputs=[model_status, chat_history, calendar_status]
+        ).then(
+            lambda: [],
+            inputs=None,
+            outputs=[chatbot]
+        )
+
+        # Initial welcome message
+        demo.load(
+            initialize_model,
+            inputs=None,
+            outputs=[model_status, chat_history, calendar_status]
+        ).then(
+            lambda: [("", "Hello! I'm your appointment booking assistant. I can help you schedule an appointment. Just let me know your name, the date, and time you prefer.")],
+            inputs=None,
+            outputs=[chatbot]
+        )
+
+    return demo
+
+# ===== MAIN EXECUTION =====
+
+if __name__ == "__main__":
+    logger.info("===== Simple Appointment Booking Assistant =====")
+    logger.info("Using Llama 3.1-8B-Instruct")
+
+    # Set PyTorch environment variables for memory efficiency
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128,garbage_collection_threshold:0.8"
+
+    try:
+        # Create and launch the Gradio interface
+        logger.info("Creating demo...")
+        demo = create_gradio_interface()
+        logger.info("Demo created, launching...")
+        demo.launch(share=False, debug=True)
+        logger.info("Gradio interface launched successfully")
+    except Exception as e:
+        logger.error(f"Error launching Gradio interface: {e}")
+        import traceback
+        logger.error(traceback.format_exc())