app.py

import streamlit as st
from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
import numpy as np
import pandas as pd
import time
import random
import re
import matplotlib.pyplot as plt
import altair as alt

# Set page configuration
st.set_page_config(page_title="China Mobile Customer Service", layout="wide")

# Define constants
USER = "user"
ASSISTANT = "assistant"
SYSTEM = "system"

# Custom CSS for chat interface
st.markdown("""
<style>
    .chat-container {
        border-radius: 10px;
        margin-bottom: 10px;
        padding: 10px;
    }
    .user-message {
        background-color: #e6f7ff;
        border-left: 3px solid #1890ff;
    }
    .assistant-message {
        background-color: #f6ffed;
        border-left: 3px solid #52c41a;
    }
    .system-message {
        background-color: #fffbe6;
        border-left: 3px solid #faad14;
        font-style: italic;
    }
    .emotion-container {
        border-radius: 10px;
        background-color: #f9f9f9;
        padding: 15px;
        margin-top: 10px;
        box-shadow: 0 1px 3px rgba(0,0,0,0.12), 0 1px 2px rgba(0,0,0,0.24);
    }
    .suggested-reply {
        background-color: #f0f5ff;
        border-radius: 10px;
        padding: 15px;
        margin-top: 10px;
        box-shadow: 0 1px 3px rgba(0,0,0,0.12), 0 1px 2px rgba(0,0,0,0.24);
    }
    .buttons-container {
        display: flex;
        gap: 10px;
        margin-top: 10px;
    }
    .chat-title {
        text-align: center;
        margin-bottom: 20px;
    }
    .feature-tag {
        display: inline-block;
        background-color: #f0f0f0;
        padding: 5px 10px;
        border-radius: 15px;
        margin-right: 5px;
        margin-bottom: 5px;
        font-size: 0.8em;
    }
    .intent-container {
        background-color: #f0f5ff;
        border-radius: 10px;
        padding: 15px;
        margin-top: 10px;
        box-shadow: 0 1px 3px rgba(0,0,0,0.12), 0 1px 2px rgba(0,0,0,0.24);
    }
    .status-indicator {
        display: flex;
        align-items: center;
        margin-bottom: 10px;
    }
    .status-dot {
        width: 10px;
        height: 10px;
        border-radius: 50%;
        margin-right: 5px;
    }
    .active {
        background-color: #52c41a;
    }
    .inactive {
        background-color: #d9d9d9;
    }
            
    .stChatFloatingInputContainer {
        position: relative !important;
        bottom: auto !important;
        width: 100%;
        padding: 10px;
    }
            
</style>
""", unsafe_allow_html=True)

# Initialize session state
if "messages" not in st.session_state:
    st.session_state.messages = []

if "chat_started" not in st.session_state:
    st.session_state.chat_started = False

if "current_emotion" not in st.session_state:
    st.session_state.current_emotion = None

if "current_features" not in st.session_state:
    st.session_state.current_features = []

if "suggested_reply" not in st.session_state:
    st.session_state.suggested_reply = ""

if "waiting_for_response" not in st.session_state:
    st.session_state.waiting_for_response = False

# Expanded candidate tasks for better intent classification
candidate_tasks = [
    "change mobile plan",
    "top up balance",
    "report service outage",
    "ask for billing support",
    "reactivate service",
    "cancel subscription",
    "check account status",
    "upgrade device",
    "report lost/stolen device",
    "request network coverage information",
    "set up international roaming",
    "inquire about data plans",
    "inquire about family plans",
    "request technical support",
    "apply for discounts or promotions",
    "dispute a charge",
    "update personal information",
    "reset password or PIN",
    "request paper bill",
    "manage automatic payments"
]

# Sample customer inquiries for simulation
SAMPLE_INQUIRIES = [
    "My data plan seems to be used up too quickly. I'm very frustrated with the service.",
    "I've been trying to upgrade my plan for days but your website keeps giving errors. This is really annoying!",
    "Thanks for the quick resolution of my billing issue. I'm very satisfied with the service.",
    "I'm confused about the different family plan options. Could you explain them to me?",
    "Your network coverage is terrible in my area. I'm considering switching to another provider.",
    "I'm curious if there's a student discount available for monthly plans?",
    "I need to report my phone as stolen. Please help me block my SIM card immediately.",
    "Can you help me set up international roaming for my trip to Europe next week?",
    "I want to cancel my current subscription. What's the process for this?"
]

# Load models (these would be replaced with actual implementations)
@st.cache_resource
def load_emotion_model():
    # In a real app, this would load your fine-tuned model
    return pipeline("text-classification", model="shengqizhao0124/emotion_trainer", return_all_scores=True)

@st.cache_resource
def load_feature_extraction_model():
    return pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

@st.cache_resource
def load_text_generation_model():
    return pipeline("text2text-generation", model="declare-lab/flan-alpaca-base")

# Function to extract intent using zero-shot classification
def extract_features(text):
    feature_model = load_feature_extraction_model()
    results = feature_model(text, candidate_tasks)
    # Return the most likely intent
    return results['labels'][0]

# Function to analyze emotion
def analyze_emotion(text):
    emotion_classifier = load_emotion_model()
    results = emotion_classifier(text)[0]
    # Convert to dictionary for easier handling
    emotion_scores = {item['label']: item['score'] for item in results}
    return emotion_scores

# Improved function to generate response based on emotion and intent
def generate_response(text, emotion, intent):
    # Get the dominant emotion
    dominant_emotion = max(emotion.items(), key=lambda x: x[1])[0]
    
    # Intent-specific templates
    intent_templates = {
        "change mobile plan": {
            "context": "Based on your usage patterns, your current 'Standard 5GB' plan might not be the best fit.",
            "action": "I can recommend our 'Unlimited Data' plan at ¥198/month or our 'Smart 20GB' plan at ¥158/month that would better suit your needs."
        },
        "top up balance": {
            "context": "Your current balance is ¥45.20, which is below our recommended minimum of ¥50.",
            "action": "You can top up through our China Mobile app, website, or at any of our physical stores. Would you like me to guide you through the online top-up process?"
        },
        "report service outage": {
            "context": "I've checked our system and can confirm there's a temporary network maintenance in your area that may be affecting your service.",
            "action": "Our technical team is already working on this issue, and service should be restored within 2-3 hours. I've registered your report which will help us prioritize the area."
        },
        "ask for billing support": {
            "context": "I can see your most recent bill was ¥178.50, issued on May 15th, which includes your monthly plan (¥158) and additional data usage (¥20.50).",
            "action": "I'd be happy to explain any specific charges or help set up a payment arrangement if needed."
        },
        "reactivate service": {
            "context": "I see that your service was deactivated on May 18th due to payment delay.",
            "action": "We can easily reactivate your service after processing the outstanding payment of ¥178.50. Would you like to make this payment now?"
        },
        "cancel subscription": {
            "context": "You currently have our 'Premium Entertainment Package' subscription for ¥30/month, which provides access to streaming services and is valid until June 15th.",
            "action": "I can process your cancellation immediately, but I'd like to understand if there's a specific reason you're canceling, as we might have alternatives that better meet your needs."
        },
        "check account status": {
            "context": "Your account is in good standing. You've been with China Mobile for 2 years and 3 months, and you're currently on our 'Family Share 50GB' plan.",
            "action": "You've used 32GB of your 50GB data allocation this month, with 10 days remaining in your billing cycle. Is there any specific aspect of your account you'd like to know more about?"
        },
        "upgrade device": {
            "context": "As a loyal customer, you're eligible for our VIP device upgrade program with a 20% discount on new flagship phones.",
            "action": "We currently have special offers on the latest smartphone models. Would you like me to tell you about our installment plans with 0% interest?"
        },
        "report lost/stolen device": {
            "context": "I'm sorry to hear about your device. This is an urgent matter that we need to address immediately to protect your account.",
            "action": "I've placed a temporary block on your SIM card to prevent unauthorized use. We can issue a replacement SIM that maintains your current number. Would you like to proceed with this?"
        },
        "request network coverage information": {
            "context": "Our coverage map shows excellent 5G signal strength in most of your city, with some potential limitations in remote suburban areas.",
            "action": "I can provide detailed coverage information for specific locations you frequent. Is there a particular area you're concerned about?"
        },
        "set up international roaming": {
            "context": "For European travel, we offer our 'Global Traveler' package at ¥58/day, which includes 500MB of daily data and 30 minutes of calls.",
            "action": "I can activate international roaming on your account right now, and it will be effective immediately. Would you like me to proceed with the activation?"
        },
        "inquire about data plans": {
            "context": "We offer various data plans ranging from our Basic 5GB at ¥88/month to our Unlimited Premium at ¥258/month.",
            "action": "Based on average user patterns similar to yours, I would recommend our Smart 20GB plan at ¥158/month. Would you like to hear more details about this or other plans?"
        },
        "inquire about family plans": {
            "context": "Our family plans allow up to 5 members to share data, with each member receiving individual voice minutes and messaging allowances.",
            "action": "Our most popular option is the Family Share 50GB at ¥298/month for up to 3 members, with each additional member at ¥50/month. Would this meet your household's needs?"
        },
        "request technical support": {
            "context": "I understand you're experiencing connectivity issues with your mobile data service.",
            "action": "Let's try some troubleshooting steps: 1) Toggle airplane mode on and off, 2) Restart your device, 3) Check your APN settings. Would you like me to guide you through these steps?"
        },
        "apply for discounts or promotions": {
            "context": "Based on your profile, you qualify for our Student Discount Program (20% off monthly plans) and our Loyalty Rewards (additional 5GB data).",
            "action": "I can apply these discounts to your account immediately. All we need is verification of your student status. Would you like to proceed?"
        },
        "dispute a charge": {
            "context": "I see a charge of ¥50 for 'Premium Content Access' on May 10th that you're questioning.",
            "action": "I'll investigate this charge thoroughly. If it was made in error, we'll process a refund within 3-5 business days. Can you provide any additional information about this charge?"
        },
        "update personal information": {
            "context": "Your current contact information shows your email as user@example.com and your address as 123 Main St, Beijing.",
            "action": "I can update your personal details right now. For security purposes, we'll send a verification code to your registered mobile number. What information would you like to update?"
        },
        "reset password or PIN": {
            "context": "For security purposes, we'll need to verify your identity before resetting your access credentials.",
            "action": "I can guide you through our secure password/PIN reset process. First, we'll send a verification code to your registered mobile number. Would you like to proceed?"
        },
        "request paper bill": {
            "context": "You're currently set up for electronic billing sent to user@example.com on the 15th of each month.",
            "action": "I can change your preference to receive paper bills at your registered address. Please note there's a small environmental fee of ¥5 per month for paper bills. Would you like to proceed?"
        },
        "manage automatic payments": {
            "context": "Your automatic payment is currently active using your credit card ending in *1234, set to process on the 18th of each month.",
            "action": "I can help modify your payment method, change the processing date, or cancel automatic payments altogether. What changes would you like to make?"
        }
    }
    
    # Default template if the intent isn't recognized
    default_template = {
        "context": "Thank you for reaching out to China Mobile customer service.",
        "action": "I'm here to assist you with your inquiry. Could you provide more details about your specific needs?"
    }
    
    # Get the appropriate template based on the detected intent
    template = intent_templates.get(intent, default_template)
    
    # Emotion-specific phrasing
    emotion_phrasing = {
        "joy": {
            "opening": "I'm delighted to hear from you today! Thank you for your positive feedback.",
            "closing": "It's been a pleasure assisting you. Is there anything else I can help with? Have a wonderful day!"
        },
        "sadness": {
            "opening": "I understand this situation can be disappointing, and I'm here to help resolve it for you.",
            "closing": "Please don't hesitate to reach out if you need any further assistance. We value your patience."
        },
        "anger": {
            "opening": "I sincerely apologize for any frustration this has caused. Your concern is our top priority, and I'm committed to resolving this for you.",
            "closing": "Thank you for bringing this to our attention. We'll use your feedback to improve our service. Is there anything else I can address for you today?"
        },
        "fear": {
            "opening": "I understand your concerns, and I want to reassure you that we'll find a solution together.",
            "closing": "I hope I've been able to address your concerns. We're here to support you whenever you need us."
        },
        "surprise": {
            "opening": "I understand this may be unexpected. Let me help clarify the situation for you.",
            "closing": "Thank you for your patience as we worked through this. Is there anything else you'd like to know?"
        },
        "disgust": {
            "opening": "I sincerely apologize for your negative experience. We take your feedback very seriously and will address this immediately.",
            "closing": "Thank you for giving us the opportunity to address this situation. Your feedback is valuable and helps us improve."
        },
        "neutral": {
            "opening": "Thank you for contacting China Mobile customer service.",
            "closing": "Is there anything else I can help you with today?"
        }
    }
    
    # Default to neutral if emotion not recognized
    emotion_phrases = emotion_phrasing.get(dominant_emotion.lower(), emotion_phrasing["neutral"])
    
    # Construct the prompt for text generation
    prompt = (
        f"You are a professional China Mobile customer service agent. "
        f"The customer has contacted you with an intent to '{intent}' and is expressing '{dominant_emotion}' emotion in their message: '{text}'. "
        f"Respond with a structured message that includes: "
        f"1) An empathetic opening: '{emotion_phrases['opening']}' "
        f"2) Context specific to their intent: '{template['context']}' "
        f"3) A helpful action or solution: '{template['action']}' "
        f"4) A polite closing: '{emotion_phrases['closing']}'"
    )
    
    try:
        # In a real application, you would use your text generation model here
        # For this demo, we'll simulate a response
        # result = load_text_generation_model()(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)[0]['generated_text'].strip()
        
        # Fallback response - structured combination of the templates
        result = f"{emotion_phrases['opening']} {template['context']} {template['action']} {emotion_phrases['closing']}"
        return result
    except Exception as e:
        # Fallback response if something goes wrong
        return f"{emotion_phrases['opening']} {template['context']} {template['action']} {emotion_phrases['closing']}"

# App title
st.title("China Mobile Customer Automated Reply System")

# Add status indicator
chat_status = "Active" if st.session_state.chat_started else "Inactive"
status_color = "active" if st.session_state.chat_started else "inactive"
st.markdown(f"""
    <div class="status-indicator">
        <div class="status-dot {status_color}"></div>
        <span>Chat Status: {chat_status}</span>
    </div>
""", unsafe_allow_html=True)

# Create two columns for the layout
col1, col2 = st.columns([2, 1])

# Customer selection in sidebar
with st.sidebar:
    st.markdown("<h4>Start a new conversation</h4>", unsafe_allow_html=True)
    
    # Option to select a specific inquiry or use random
    selected_inquiry = st.selectbox(
        "Select a customer inquiry or use random:",
        ["Random inquiry"] + SAMPLE_INQUIRIES
    )
    
    start_button_label = "Start with Random Inquiry" if selected_inquiry == "Random inquiry" else "Start with Selected Inquiry"
    
    if st.button(start_button_label, use_container_width=True):
        # Get inquiry based on selection
        inquiry = random.choice(SAMPLE_INQUIRIES) if selected_inquiry == "Random inquiry" else selected_inquiry
        
        # Reset any existing conversation
        st.session_state.messages = []
        
        # Add system message
        st.session_state.messages.append({"role": SYSTEM, "content": "A new customer has connected to chat."})
        
        # Add customer message
        st.session_state.messages.append({"role": USER, "content": inquiry})
        
        # Set chat as started
        st.session_state.chat_started = True
        
        # Analyze emotion
        st.session_state.current_emotion = analyze_emotion(inquiry)
        
        # Extract features (intent)
        st.session_state.current_features = extract_features(inquiry)
        
        # Generate suggested reply
        st.session_state.suggested_reply = generate_response(
            inquiry, 
            st.session_state.current_emotion, 
            st.session_state.current_features
        )
        
        st.session_state.waiting_for_response = True
        
        # Force rerun to update the UI
        st.rerun()

# Main chat area in column 1
with col1:
    st.markdown("<h3 class='chat-title'>Customer Service Chat</h3>", unsafe_allow_html=True)
    
    # Chat container with border and fixed height
    chat_container = st.container(border=True, height=600)
    with chat_container:
        # Display all messages in the chat history
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])
    
        # Chat input area - only shown when we're not waiting for a response
        if st.session_state.chat_started :
            text_area = st.text_area(label = '',value = st.session_state.suggested_reply)
            user_input = st.chat_input('Say something')
            if prompt := user_input:
                # Add the user's message to the chat
                st.session_state.messages.append({"role": USER, "content": user_input})
                
                # Process this new message
                st.session_state.current_emotion = analyze_emotion(user_input)
                st.session_state.current_features = extract_features(user_input)
                st.session_state.suggested_reply = generate_response(
                    user_input,
                    st.session_state.current_emotion,
                    st.session_state.current_features
                )
                
                st.session_state.waiting_for_response = True
                
                # Rerun to update UI
                st.rerun()
    
    # Suggested reply section (only shown if waiting for response)
    if st.session_state.waiting_for_response and st.session_state.suggested_reply:
        # Accept/Reject buttons in a single row
        cols = st.columns(2)
        with cols[0]:
            if st.button("✅ Accept and Send", use_container_width=True):
                # Add assistant message with suggested reply
                st.session_state.messages.append({"role": ASSISTANT, "content": text_area})
                
                # Clear suggested reply
                st.session_state.suggested_reply = ""
                
                # No longer waiting for response
                st.session_state.waiting_for_response = False
                
                # Force rerun to update the UI
                st.rerun()
                
        with cols[1]:
            if st.button("✏️ Modify Reply", use_container_width=True):
                # Show text area for manual editing
                modified_reply = text_area
                
                if st.button("Send Modified Reply", use_container_width=True):
                    # Add assistant message with modified reply
                    st.session_state.messages.append({"role": ASSISTANT, "content": modified_reply})
                    
                    # Clear suggested reply
                    st.session_state.suggested_reply = ""
                    
                    # No longer waiting for response
                    st.session_state.waiting_for_response = False
                    
                    # Force rerun to update the UI
                    st.rerun()

# Analysis panel in column 2
with col2:
    st.markdown("<h3>Analysis Dashboard</h3>", unsafe_allow_html=True)
    
    # Intent analysis with icon
    if st.session_state.current_features:
        # Map intents to icons
        intent_icons = {
            "change mobile plan": "📱",
            "top up balance": "💰",
            "report service outage": "🔌",
            "ask for billing support": "💲",
            "reactivate service": "🔄",
            "cancel subscription": "❌",
            "check account status": "📊",
            "upgrade device": "📲",
            "report lost/stolen device": "🚨",
            "request network coverage information": "📶",
            "set up international roaming": "🌍",
            "inquire about data plans": "📊",
            "inquire about family plans": "👪",
            "request technical support": "🔧",
            "apply for discounts or promotions": "🏷️",
            "dispute a charge": "⚠️",
            "update personal information": "📝",
            "reset password or PIN": "🔒",
            "request paper bill": "📄",
            "manage automatic payments": "💳"
        }
        
        # Get icon for intent or default
        intent_icon = intent_icons.get(st.session_state.current_features, "🔍")
        
        st.markdown(f"""
            <div class="intent-container">
                <h4>Detected Intent {intent_icon}</h4>
                <p style="font-size: 1.2em; font-weight: bold;">{st.session_state.current_features}</p>
                <div style="background-color: #e6f7ff; height: 8px; border-radius: 4px; margin: 10px 0;">
                    <div style="background-color: #1890ff; width: 85%; height: 8px; border-radius: 4px;"></div>
                </div>
                <p style="text-align: right; font-size: 0.8em;">Confidence: 85%</p>
            </div>
        """, unsafe_allow_html=True)
    
    # Emotion analysis visualization
    if st.session_state.current_emotion:
        st.markdown("<h4>Emotion Analysis</h4>", unsafe_allow_html=True)
        
        # Create a dataframe for the emotion scores
        emotion_df = pd.DataFrame({
            'Emotion': list(st.session_state.current_emotion.keys()),
            'Score': list(st.session_state.current_emotion.values())
        })
        
        # Sort by score in descending order
        emotion_df = emotion_df.sort_values('Score', ascending=False)
        
        # Display dominant emotion with emoji
        emotion_emojis = {
            "anger": "😠",
            "disgust": "🤢",
            "fear": "😨",
            "joy": "😀",
            "neutral": "😐",
            "sadness": "😢",
            "surprise": "😲"
        }
        
        dominant_emotion = emotion_df.iloc[0]['Emotion']
        dominant_score = emotion_df.iloc[0]['Score']
        emoji = emotion_emojis.get(dominant_emotion.lower(), "")
        
        st.markdown(f"""
            <div class="intent-container">
                <h4>{emoji} {dominant_emotion}</h4>
                <p style="font-size: 1.2em; font-weight: bold;">Confidence: {dominant_score:.1%}</p>
            </div>
        """, unsafe_allow_html=True)
        
        # Create a horizontal bar chart
        chart = alt.Chart(emotion_df).mark_bar().encode(
            x=alt.X('Score:Q', axis=alt.Axis(format='.0%')),
            y=alt.Y('Emotion:N', sort='-x'),
            color=alt.Color('Emotion:N', legend=None, 
                           scale=alt.Scale(range=['#f5222d', '#fa8c16', '#faad14', '#a0d911', '#13c2c2', '#1890ff', '#722ed1']))
        ).properties(
            title='Emotion Distribution',
            width=300,
            height=200
        )
        
        st.altair_chart(chart, use_container_width=True)
    


# Button to reset the conversation
if st.button("Reset Conversation", use_container_width=True):
    # Clear session state
    st.session_state.messages = []
    st.session_state.chat_started = False
    st.session_state.current_emotion = None
    st.session_state.current_features = []
    st.session_state.suggested_reply = ""
    st.session_state.waiting_for_response = False
    
    # Force rerun to update the UI
    st.rerun()