"""Main flow controller for the barber booking system.""" import json from typing import Dict, Any, Union, TypedDict from langchain.schema import HumanMessage, SystemMessage from langchain_core.messages import AIMessage import os from dotenv import load_dotenv import google.generativeai as genai from config.services import SERVICE_MAPPING from config.prompts import BOOKING_FLOW from config.validation import ( validate_email, validate_phone, validate_name, format_phone, validate_booking_info, get_closest_service ) # Load environment variables load_dotenv() # Initialize Gemini def get_llm(): api_key = os.getenv("GOOGLE_API_KEY") if not api_key: raise ValueError("GOOGLE_API_KEY environment variable is not set") try: genai.configure(api_key=api_key) model = genai.GenerativeModel('gemini-pro') return model except Exception as e: print(f"Error initializing Gemini: {str(e)}") raise try: llm = get_llm() except Exception as e: print(f"Error initializing LLM: {str(e)}") raise # Type definition for booking state class BookingState(TypedDict): current_node: str booking_info: Dict[str, Any] response: str messages: list def handle_greeting(message: str, booking_info: Dict) -> Dict: """Handle the greeting state.""" if validate_name(message): name = message.strip().title() return { "current_node": "ServiceSelection", "booking_info": {"name": name}, "response": f"Hi {name}, welcome to our barbershop! Here are our services:" } return { "current_node": "Greeting", "booking_info": booking_info, "response": "Could you please tell me your name?" } def handle_service_selection(message: str, booking_info: Dict) -> Dict: """Handle the service selection state.""" message = message.lower().strip() # Check for cancellation if message in ["no", "cancel", "actually no", "nevermind"]: return { "current_node": "ServiceSelection", "booking_info": booking_info, "response": "No problem. Which service would you like to book? We offer haircut, beard trim, and full service." } # Check for multiple services if "and" in message or "," in message: return { "current_node": "ServiceSelection", "booking_info": booking_info, "response": "I see you're interested in multiple services. I recommend booking our full service which includes both haircut and beard trim. Would you like that?" } # Check for completion keywords if message in ["no", "nothing", "nothing else", "that's all"]: if "service" in booking_info: return { "current_node": "ShowTopSlots", "booking_info": booking_info, "response": "Great! Let me show you our available time slots." } # Try to match service service = get_closest_service(message) if service: booking_info["service"] = service return { "current_node": "ShowTopSlots", "booking_info": booking_info, "response": f"Perfect! Let me show you our available time slots for your {service}." } return { "current_node": "ServiceSelection", "booking_info": booking_info, "response": "Which service would you like to book? We offer haircut, beard trim, and full service." } def handle_time_selection(message: str, booking_info: Dict) -> Dict: """Handle the time selection state.""" time_slots = ["9:00 AM", "10:00 AM", "11:00 AM", "2:00 PM", "3:00 PM"] message = message.upper().strip() # Check for cancellation or change if message.lower() in ["cancel", "change", "different time", "other time"]: slots_text = "\n".join([f"- {slot}" for slot in time_slots]) return { "current_node": "TimeSelection", "booking_info": booking_info, "response": f"Here are all our available times:\n{slots_text}\nWhich time works best for you?" } # Handle relative time references relative_times = { "MORNING": ["9:00 AM", "10:00 AM", "11:00 AM"], "AFTERNOON": ["2:00 PM", "3:00 PM"], "EARLY": ["9:00 AM", "10:00 AM"], "LATE": ["2:00 PM", "3:00 PM"], "FIRST": ["9:00 AM"], "LAST": ["3:00 PM"] } for keyword, options in relative_times.items(): if keyword in message: slots_text = "\n".join([f"- {slot}" for slot in options]) return { "current_node": "TimeSelection", "booking_info": booking_info, "response": f"Here are the {keyword.lower()} slots:\n{slots_text}\nWhich specific time would you prefer?" } # Normalize input message = message.replace(":", "").replace(" ", "") # Map variations to standard format time_map = { "9": "9:00 AM", "9AM": "9:00 AM", "900": "9:00 AM", "900AM": "9:00 AM", "NINE": "9:00 AM", "10": "10:00 AM", "10AM": "10:00 AM", "1000": "10:00 AM", "1000AM": "10:00 AM", "TEN": "10:00 AM", "11": "11:00 AM", "11AM": "11:00 AM", "1100": "11:00 AM", "1100AM": "11:00 AM", "ELEVEN": "11:00 AM", "2": "2:00 PM", "2PM": "2:00 PM", "200": "2:00 PM", "200PM": "2:00 PM", "TWO": "2:00 PM", "3": "3:00 PM", "3PM": "3:00 PM", "300": "3:00 PM", "300PM": "3:00 PM", "THREE": "3:00 PM" } selected_time = time_map.get(message) if selected_time and selected_time in time_slots: booking_info["time_slot"] = selected_time return { "current_node": "CustomerInfo", "booking_info": booking_info, "response": f"Perfect! I'll book you for {selected_time}. What's your email address?" } # If time not found, show all slots slots_text = "\n".join([f"- {slot}" for slot in time_slots]) return { "current_node": "TimeSelection", "booking_info": booking_info, "response": f"I couldn't understand that time. Please select from these available times:\n{slots_text}" } def handle_customer_info(message: str, booking_info: Dict) -> Dict: """Handle the customer information state.""" if "email" not in booking_info: if validate_email(message): booking_info["email"] = message return { "current_node": "CustomerInfo", "booking_info": booking_info, "response": "Great! Now, what's your phone number for appointment reminders?" } return { "current_node": "CustomerInfo", "booking_info": booking_info, "response": "Please provide a valid email address (e.g., name@example.com)." } if "phone" not in booking_info: if validate_phone(message): booking_info["phone"] = format_phone(message) return { "current_node": "Confirmation", "booking_info": booking_info, "response": "Perfect! Let me confirm your booking details..." } return { "current_node": "CustomerInfo", "booking_info": booking_info, "response": "Please provide a valid phone number (e.g., 1234567890)." } return { "current_node": "CustomerInfo", "booking_info": booking_info, "response": "I need your contact information. What's your email address?" } def handle_confirmation(message: str, booking_info: Dict) -> Dict: """Handle the confirmation state.""" message = message.lower().strip() # Format booking summary service = booking_info.get("service", "") price = SERVICE_MAPPING[service]["price"] if service else 0 duration = SERVICE_MAPPING[service]["duration"] if service else 0 summary = f"""Here's your booking summary: - Name: {booking_info.get('name', '')} - Service: {service.title()} (${price}, {duration} min) - Time: {booking_info.get('time_slot', '')} - Email: {booking_info.get('email', '')} - Phone: {booking_info.get('phone', '')} Is this correct? Please confirm (yes/no).""" if message in ["yes", "correct", "confirm", "y"]: return { "current_node": "BookingComplete", "booking_info": {**booking_info, "confirmation": True}, "response": "Great! Your appointment has been confirmed. We'll send you a confirmation email shortly." } if message in ["no", "wrong", "incorrect", "n"]: return { "current_node": "ServiceSelection", "booking_info": {"name": booking_info.get("name", "")}, "response": "I understand. Let's start over with the service selection." } return { "current_node": "Confirmation", "booking_info": booking_info, "response": summary } def process_node(state: Dict) -> Dict: """Process the current node in the booking flow.""" try: current_node = state["current_node"] booking_info = state.get("booking_info", {}) messages = state.get("messages", []) # Get the last user message last_message = "" if messages: last_msg = messages[-1] if isinstance(last_msg, dict): last_message = last_msg.get("content", "") elif isinstance(last_msg, (HumanMessage, AIMessage, SystemMessage)): last_message = last_msg.content else: last_message = str(last_msg) # Handle each state if current_node == "Greeting": new_state = handle_greeting(last_message, booking_info) elif current_node == "ServiceSelection": new_state = handle_service_selection(last_message, booking_info) elif current_node == "ShowTopSlots": time_slots = ["9:00 AM", "10:00 AM", "11:00 AM", "2:00 PM", "3:00 PM"] slots_text = "\n".join([f"- {slot}" for slot in time_slots]) new_state = { "current_node": "TimeSelection", "booking_info": booking_info, "response": f"Here are our available time slots:\n{slots_text}\n\nWhich time works best for you?" } elif current_node == "TimeSelection": new_state = handle_time_selection(last_message, booking_info) elif current_node == "CustomerInfo": new_state = handle_customer_info(last_message, booking_info) elif current_node == "Confirmation": new_state = handle_confirmation(last_message, booking_info) elif current_node == "BookingComplete": new_state = { "current_node": "Farewell", "booking_info": booking_info, "response": "Thank you for booking with us! We look forward to seeing you soon." } else: new_state = { "current_node": "Greeting", "booking_info": {}, "response": "Hello! Welcome to our barber shop. Could you please tell me your name?" } return {**new_state, "messages": messages} except Exception as e: print(f"Error in process_node: {str(e)}") return { "current_node": state.get("current_node", "Greeting"), "booking_info": state.get("booking_info", {}), "response": "I encountered an error. Could you please rephrase that?", "messages": messages } def run_booking_flow(user_input: str, state: Dict[str, Any] = None) -> Dict[str, Any]: """Run the booking workflow.""" if state is None: state = { "current_node": "Greeting", "booking_info": {}, "response": "", "messages": [] } try: # Process the current node new_state = process_node(state) return new_state except Exception as e: print(f"Error in run_booking_flow: {str(e)}") return { "current_node": "Greeting", "booking_info": {}, "response": "I encountered an error. Let's start over.", "messages": [] } # Export necessary components __all__ = ['run_booking_flow', 'SERVICE_MAPPING']