Spaces:
Sleeping
Sleeping
| """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'] | |