app.py

import streamlit as st
from transformers import pipeline
import os
from huggingface_hub import login

login(os.environ["HF_TOKEN"])

emotion_classifier = pipeline("text-classification", model="shengqizhao0124/emotion_trainer", return_all_scores=True)
intent_classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
text_generator = pipeline("text2text-generation", model="declare-lab/flan-alpaca-base")

candidate_tasks = [
    "change mobile plan", "top up balance", "report service outage",
    "ask for billing support", "reactivate service", "cancel subscription",
    "check account status", "upgrade device"
]

urgent_emotions = {"anger", "frustration", "anxiety", "urgency", "afraid", "annoyed"}
moderate_emotions = {"confused", "sad", "tired", "concerned", "sadness"}

def refine_emotion_label(text, model_emotion):
    text_lower = text.lower()
    urgent_keywords = ["fix", "now", "immediately", "urgent", "can't", "need", "asap"]
    exclamations = text.count("!")
    upper_words = sum(1 for w in text.split() if w.isupper())
    signal_score = sum([
        any(word in text_lower for word in urgent_keywords),
        exclamations >= 2,
        upper_words >= 1
    ])
    if model_emotion.lower() in {"joy", "neutral", "sadness"} and signal_score >= 2:
        return "urgency"
    return model_emotion

def get_emotion_label(result, text):
    sorted_emotions = sorted(result[0], key=lambda x: x['score'], reverse=True)
    return refine_emotion_label(text, sorted_emotions[0]['label'])

def get_emotion_score(emotion):
    if emotion.lower() in urgent_emotions:
        return 1.0
    elif emotion.lower() in moderate_emotions:
        return 0.6
    else:
        return 0.2

def generate_response(intent, human=True):
    if human:
        prompt = (
            f"You are a telecom agent. The customer intends to '{intent}'. "
            "Give a 3-part polite reply: 1) Greeting, 2) Mention current plan (fictional) and suggest better one, 3) Ask if want to proceed."
        )
        result = text_generator(prompt, max_new_tokens=150, do_sample=False)
        return result[0]['generated_text'].strip()
    else:
        return f"[Below is a link to the service you need：{intent} → https://support.example.com/{intent.replace(' ', '_')}]\\n[If your problem still can not be solved, welcome to continue to consult, we will continue to serve you!]"

st.set_page_config(page_title="Smart Customer Support Assistant", layout="wide")
st.sidebar.title("📁 Customer Selector")

if "customers" not in st.session_state:
    st.session_state.customers = {"Customer A": [], "Customer B": [], "Customer C": []}
if "chat_sessions" not in st.session_state:
    st.session_state.chat_sessions = {}

selected_customer = st.sidebar.selectbox("Choose a customer:", list(st.session_state.customers.keys()))

if selected_customer not in st.session_state.chat_sessions:
    st.session_state.chat_sessions[selected_customer] = {
        "chat": [], "system_result": None,
        "agent_reply": "", "support_required": "", "user_input": ""
    }

session = st.session_state.chat_sessions[selected_customer]
st.title("Smart Customer Support Assistant (for Agents Only)")

# Conversation UI
st.markdown("### Conversation")
for msg in session["chat"]:
    avatar = "👤" if msg['role'] == 'user' else ("🤖" if msg.get("auto") else "👨‍💼")
    with st.chat_message(msg['role'], avatar=avatar):
        if msg["role"] == "user" and "emotion" in msg:
            st.markdown(f"<div style='text-align:right;font-size:0.9em;color:gray;'>Emotion: {msg['emotion'].capitalize()}</div>", unsafe_allow_html=True)
        st.markdown(msg['content'])

# Input & Analyze
col1, col2 = st.columns([6, 1])
with col1:
    user_input = st.text_input("Enter customer message:", key="customer_input")
with col2:
    if st.button("Analyze"):
        if user_input.strip():
            emotion_result = emotion_classifier(user_input)
            emotion_label = get_emotion_label(emotion_result, user_input)
            emotion_score = get_emotion_score(emotion_label)

            session["chat"].append({"role": "user", "content": user_input, "emotion": emotion_label})

            intent_result = intent_classifier(user_input, candidate_tasks)
            top_intents = [label for label, score in zip(intent_result['labels'], intent_result['scores']) if score > 0.15][:3]

            content_score = 0.0
            if any(x in user_input.lower() for x in ["out of service", "can't", "urgent", "immediately"]):
                content_score += 0.4
            if any(label in ["top up balance", "reactivate service"] for label in top_intents):
                content_score += 0.4

            final_score = 0.5 * emotion_score + 0.5 * content_score

            if final_score < 0.5 and top_intents:
                intent = top_intents[0]
                response = generate_response(intent, human=False)
                session["chat"].append({"role": "assistant", "content": response, "auto": True})
                session["system_result"] = None
                session["support_required"] = "🟢 Automated response handled this request."
            else:
                session["system_result"] = {
                    "emotion": emotion_label,
                    "tone": "Urgent" if emotion_score > 0.8 else "Concerned" if emotion_score > 0.5 else "Calm",
                    "intents": top_intents
                }
                session["support_required"] = "🔴 Human support required."
                session["agent_reply"] = ""
            st.rerun()

# Agent panel
if session["support_required"]:
    st.markdown(f"### {session['support_required']}")

st.subheader("Agent Response Console")
session["agent_reply"] = st.text_area("Compose your reply:", value=session["agent_reply"], key="agent_reply_box")
if st.button("Send Reply"):
    if session["agent_reply"].strip():
        session["chat"].append({"role": "assistant", "content": session["agent_reply"], "auto": False})
        session["agent_reply"] = ""
        session["system_result"] = None
        session["support_required"] = ""
        st.rerun()

# Show customer analysis
if session["system_result"] is not None:
    st.markdown("#### Customer Status")
    st.markdown(f"- **Emotion:** {session['system_result']['emotion'].capitalize()}")
    st.markdown(f"- **Tone:** {session['system_result']['tone']}")
    st.markdown("#### Detected Customer Needs")
    for intent in session["system_result"]["intents"]:
        suggestion = generate_response(intent, human=True)
        st.markdown(f"**• {intent.capitalize()}**")
        st.code(suggestion)

if st.button("End Conversation"):
    session["chat"] = []
    session["system_result"] = None
    session["agent_reply"] = ""
    session["support_required"] = ""
    session["user_input"] = ""
    st.success("Conversation ended and cleared.")
    st.rerun()