Update app.py

app.py

@@ -1,13 +1,10 @@
 # Smart Customer Support Assistant (Enhanced UI Version)
-# Note: Core analysis logic remains unchanged, now with text generation and customer selection
+# Note: Enhanced UI with role avatars, structured suggestions, and end chat functionality
 
 import streamlit as st
 from transformers import pipeline
 import re
 
-# ------------------------------
-# Load models (now includes 3rd: text generation)
-# ------------------------------
 emotion_classifier = pipeline(
     "text-classification",
     model="j-hartmann/emotion-english-distilroberta-base",
@@ -16,9 +13,6 @@ emotion_classifier = pipeline(
 intent_classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
 text_generator = pipeline("text2text-generation", model="declare-lab/flan-alpaca-base")
 
-# ------------------------------
-# Candidate tasks / prompts
-# ------------------------------
 candidate_tasks = [
     "change mobile plan",
     "top up balance",
@@ -30,22 +24,23 @@ candidate_tasks = [
     "upgrade device"
 ]
 
-def generate_response(intent):
-    prompt = (
-        f"Generate a customer service reply for: '{intent}'. "
-        "Structure your answer into: 1. Greeting, 2. Summary of customer's current situation (e.g., current plan is Plan X costing ¥X/month or balance is ¥X), "
-        "3. Your suggested recommendation (e.g., upgrade to Plan Y with XX GB), and 4. End with a helpful and polite question. "
-        "Use anonymized or placeholder values for all specific customer data."
-    )
-    output = text_generator(prompt, max_new_tokens=100, do_sample=True)[0]['generated_text']
-    return output
+def generate_response(intent, human=True):
+    if human:
+        prompt = (
+            f"Write a customer service message for intent '{intent}'. Structure into: "
+            "1. Greeting, 2. Description of customer's current service (Plan X ¥X/GB, etc.) and a suitable new option, "
+            "3. A polite closing question. Use placeholder data (Plan X, ¥X, etc.)."
+        )
+    else:
+        prompt = (
+            f"Reply to the following user request automatically: '{intent}'. Directly resolve the request in one helpful message. "
+            "If it's plan change, suggest suitable options. Use placeholders (Plan A ¥X/5GB, Plan B ¥Y/15GB, etc.)."
+        )
+    return text_generator(prompt, max_new_tokens=100, do_sample=True)[0]['generated_text']
 
 urgent_emotions = {"anger", "frustration", "anxiety", "urgency", "afraid", "annoyed"}
 moderate_emotions = {"confused", "sad", "tired", "concerned", "sadness"}
 
-# ------------------------------
-# Emotion processing
-# ------------------------------
 def refine_emotion_label(text, model_emotion):
     text_lower = text.lower()
     urgent_keywords = ["fix", "now", "immediately", "urgent", "can't", "need", "asap"]
@@ -72,9 +67,6 @@ def get_emotion_score(emotion):
     else:
         return 0.2
 
-# ------------------------------
-# UI: Sidebar for customer selection
-# ------------------------------
 st.set_page_config(page_title="Smart Customer Support Assistant", layout="wide")
 st.sidebar.title("📁 Customer Selector")
 if "customers" not in st.session_state:
@@ -82,7 +74,6 @@ if "customers" not in st.session_state:
 customer_names = list(st.session_state.customers.keys())
 selected_customer = st.sidebar.selectbox("Choose a customer:", customer_names)
 
-# Load or init selected customer's session
 if "chat_sessions" not in st.session_state:
     st.session_state.chat_sessions = {}
 if selected_customer not in st.session_state.chat_sessions:
@@ -94,14 +85,12 @@ if selected_customer not in st.session_state.chat_sessions:
     }
 session = st.session_state.chat_sessions[selected_customer]
 
-# ------------------------------
-# Main Interface
-# ------------------------------
 st.title("Smart Customer Support Assistant (for Agents Only)")
 
 st.markdown("### Conversation")
 for msg in session["chat"]:
-    with st.chat_message(msg['role']):
+    avatar = "👤" if msg['role'] == 'user' else ("🤖" if msg.get("auto") else "👨‍💼")
+    with st.chat_message(msg['role'], avatar=avatar):
         st.markdown(msg['content'])
 
 col1, col2 = st.columns([6,1])
@@ -129,8 +118,8 @@ if analyze_clicked and user_input.strip():
 
     if final_score < 0.5 and top_intents:
         intent = top_intents[0]
-        response = generate_response(intent)
-        session["chat"].append({"role": "assistant", "content": response})
+        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:
@@ -148,7 +137,7 @@ st.subheader("Agent Response Console")
 session["agent_reply"] = st.text_area("Compose your reply:", value=session["agent_reply"])
 if st.button("Send Reply"):
     if session["agent_reply"].strip():
-        session["chat"].append({"role": "assistant", "content": session["agent_reply"]})
+        session["chat"].append({"role": "assistant", "content": session["agent_reply"], "auto": False})
         session["agent_reply"] = ""
         session["system_result"] = None
         session["support_required"] = ""
@@ -160,6 +149,13 @@ if session["system_result"] is not None:
 
     st.markdown("#### Detected Customer Needs")
     for intent in session['system_result']['intents']:
-        suggestion = generate_response(intent)
+        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"] = ""
+    st.success("Conversation ended and cleared.")