c_langgraph_flow_new.py

"""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']