Update app.py

app.py

@@ -11,47 +11,50 @@ import time
 import os
 
 # --- import your architecture ---
-from model import DeBERTaLSTMClassifier 
+# Make sure this file is in the repo (e.g., models/deberta_lstm_classifier.py)
+# and update the import path accordingly.
+from model import DeBERTaLSTMClassifier  # <-- your class
 
 # --- Import RAG modules ---
-# Đảm bảo bạn đã có file rag_engine.py và llm_client.py trong cùng thư mục
-try:
-    from rag_engine import RAGEngine
-    from llm_client import LLMClient
-    RAG_AVAILABLE = True
-except ImportError:
-    print("⚠️ Warning: RAG modules not found. The LLM tab might not work.")
-    RAG_AVAILABLE = False
+from rag_engine import RAGEngine
+from llm_client import LLMClient
 
 # --------- Config ----------
-REPO_ID = "dungeon29/phishing-deberta-lstm"       
-CKPT_NAME = "pytorch_model.bin"                
-MODEL_NAME = "microsoft/deberta-base"         
-LABELS = ["benign", "phishing"]               
+REPO_ID = "dungeon29/phishing-deberta-lstm"       # HF repo that holds the checkpoint
+CKPT_NAME = "pytorch_model.bin"               # the .pt file name
+MODEL_NAME = "microsoft/deberta-base"         # base tokenizer/backbone
+LABELS = ["benign", "phishing"]               # adjust to your classes
+
+# If your checkpoint contains hyperparams, you can fetch them like:
+# checkpoint.get("config") or checkpoint.get("model_args")
+# and pass into DeBERTaLSTMClassifier(**model_args)
 
 # --------- Load model/tokenizer once (global) ----------
 device = "cuda" if torch.cuda.is_available() else "cpu"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-print(f"Loading DeBERTa model from {REPO_ID}...")
 ckpt_path = hf_hub_download(repo_id=REPO_ID, filename=CKPT_NAME)
 checkpoint = torch.load(ckpt_path, map_location=device)
 
-model_args = checkpoint.get("model_args", {}) 
+# If you saved hyperparams in the checkpoint, use them:
+model_args = checkpoint.get("model_args", {})  # e.g., {"lstm_hidden":256, "num_labels":2, ...}
 model = DeBERTaLSTMClassifier(**model_args)
 
 # Load weights
 try:
     state_dict = torch.load(ckpt_path, map_location=device)
+    
+    # Xử lý nếu file lưu dạng checkpoint đầy đủ (có key "model_state_dict")
     if "model_state_dict" in state_dict:
         state_dict = state_dict["model_state_dict"]
 
     model.load_state_dict(state_dict, strict=False)
     
+    # Kiểm tra layer attention
     if hasattr(model, 'attention') and 'attention.weight' not in state_dict:
          print("⚠️ Loaded model without attention layer, using newly initialized attention weights")
     else:
-         print("✅ Load DeBERTa weights successfully!")
+         print("✅ Load weights successfully!")
          
 except Exception as e:
     print(f"❌ Error when loading weights: {e}")
@@ -60,33 +63,29 @@ except Exception as e:
 model.to(device).eval()
 
 # --------- Initialize RAG & LLM ----------
-rag_engine = None
-llm_client = None
+print("Initializing RAG Engine (LangChain)...")
+rag_engine = RAGEngine()
+print("RAG Engine ready.")
 
-if RAG_AVAILABLE:
-    print("Initializing RAG Engine...")
-    try:
-        rag_engine = RAGEngine()
-        print("RAG Engine ready.")
-        
-        print("Initializing Qwen2.5-3B LLM (this may take a minute)...")
-        llm_client = LLMClient()
-        print("LLM ready.")
-    except Exception as e:
-        print(f"❌ Error initializing RAG/LLM: {e}")
+print("Initializing Qwen2.5-3B LLM (LangChain)...")
+# Pass vector_store to LLMClient for RetrievalQA
+llm_client = LLMClient(vector_store=rag_engine.vector_store)
+print("LLM ready.")
 
 # --------- Helper functions ----------
 def is_url(text):
+    """Check if text is a URL"""
     url_pattern = re.compile(
-        r'^https?://' 
-        r'(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+[A-Z]{2,6}\.?|' 
-        r'localhost|' 
-        r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})' 
-        r'(?::\d+)?' 
+        r'^https?://'  # http:// or https://
+        r'(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+[A-Z]{2,6}\.?|'  # domain...
+        r'localhost|'  # localhost...
+        r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})'  # ...or ip
+        r'(?::\d+)?'  # optional port
         r'(?:/?|[/?]\S+)$', re.IGNORECASE)
     return url_pattern.match(text) is not None
 
 def fetch_html_content(url, timeout=10):
+    """Fetch HTML content from URL"""
     try:
         headers = {
             'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
@@ -95,10 +94,12 @@ def fetch_html_content(url, timeout=10):
         response.raise_for_status()
         
         soup = BeautifulSoup(response.text, 'html.parser')
+
         for script in soup(["script", "style", "meta", "noscript", "header", "footer"]):
             script.decompose() 
 
         text = soup.get_text(separator=' ')
+        
         clean_text = " ".join(text.split())
         return clean_text, response.status_code
     except requests.exceptions.RequestException as e:
@@ -107,6 +108,8 @@ def fetch_html_content(url, timeout=10):
         return None, f"General error: {str(e)}"
 
 def predict_single_text(text, text_type="text"):
+    """Predict for a single text input"""
+    # Tokenize
     inputs = tokenizer(
         text,
         return_tensors="pt",
@@ -114,11 +117,14 @@ def predict_single_text(text, text_type="text"):
         padding=True,
         max_length=256
     )
+    # DeBERTa typically doesn't use token_type_ids
     inputs.pop("token_type_ids", None)
+    # Move to device
     inputs = {k: v.to(device) for k, v in inputs.items()}
 
     with torch.no_grad():
         try:
+            # Try to get predictions with attention weights
             result = model(**inputs, return_attention=True)
             if isinstance(result, tuple) and len(result) == 3:
                 logits, attention_weights, deberta_attentions = result
@@ -127,55 +133,73 @@ def predict_single_text(text, text_type="text"):
                 logits = result
                 has_attention = False
         except TypeError:
+            # Fallback for older model without return_attention parameter
             logits = model(**inputs)
             has_attention = False
             
         probs = F.softmax(logits, dim=-1).squeeze(0).tolist()
 
+    # Get tokens for visualization
     tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'].squeeze(0).tolist())
+    
     return probs, tokens, has_attention, attention_weights if has_attention else None
 
 def combine_predictions(url_probs, html_probs, url_weight=0.3, html_weight=0.7):
+    """Combine URL and HTML content predictions"""
     combined_probs = [
-        url_weight * url_probs[0] + html_weight * html_probs[0], 
-        url_weight * url_probs[1] + html_weight * html_probs[1] 
+        url_weight * url_probs[0] + html_weight * html_probs[0],  # benign
+        url_weight * url_probs[1] + html_weight * html_probs[1]   # phishing
     ]
     return combined_probs
 
-# --------- Inference function (DeBERTa) ----------
+# --------- Inference function ----------
 def predict_fn(text: str):
     if not text or not text.strip():
-        return "Please enter a URL or text."
+        return {"error": "Please enter a URL or text."}, ""
 
+    # Check if input is URL
     if is_url(text.strip()):
+        # Process URL
         url = text.strip()
+        
+        # Get prediction for URL itself
         url_probs, url_tokens, url_has_attention, url_attention = predict_single_text(url, "URL")
+        
+        # Try to fetch HTML content
         html_content, status = fetch_html_content(url)
         
         if html_content:
+            # Get prediction for HTML content
             html_probs, html_tokens, html_has_attention, html_attention = predict_single_text(html_content, "HTML")
+            
+            # Combine predictions
             combined_probs = combine_predictions(url_probs, html_probs)
             
+            # Use combined probabilities but show analysis for both
             probs = combined_probs
-            tokens = url_tokens + ["[SEP]"] + html_tokens[:50] 
+            tokens = url_tokens + ["[SEP]"] + html_tokens[:50]  # Limit HTML tokens for display
             has_attention = url_has_attention or html_has_attention
             attention_weights = url_attention if url_has_attention else html_attention
             
             analysis_type = "Combined URL + HTML Analysis"
             fetch_status = f"✅ Successfully fetched HTML content (Status: {status})"
+            
         else:
+            # Fallback to URL-only analysis
             probs = url_probs
             tokens = url_tokens
             has_attention = url_has_attention
             attention_weights = url_attention
+            
             analysis_type = "URL-only Analysis"
             fetch_status = f"⚠️ Could not fetch HTML content: {status}"
     else:
+        # Process as regular text
         probs, tokens, has_attention, attention_weights = predict_single_text(text, "text")
         analysis_type = "Text Analysis"
         fetch_status = ""
 
-    # Create HTML output
+    # Create detailed analysis
     predicted_class = "phishing" if probs[1] > probs[0] else "benign"
     confidence = max(probs)
     
@@ -194,6 +218,7 @@ def predict_fn(text: str):
     </div>
 </div>
 """
+
     if fetch_status:
         detailed_analysis += f"""
 <div style="background: #2d2d2d; padding: 15px; border-radius: 10px; margin: 15px 0; border-left: 4px solid #4caf50; color: #e0e0e0;">
@@ -203,39 +228,68 @@ def predict_fn(text: str):
     
     if has_attention and attention_weights is not None:
         attention_scores = attention_weights.squeeze(0).tolist()
+        
         token_analysis = []
         for i, (token, score) in enumerate(zip(tokens, attention_scores)):
+            # More lenient filtering - include more tokens for text analysis
             if token not in ['[CLS]', '[SEP]', '[PAD]', '<s>', '</s>'] and len(token.strip()) > 0 and score > 0.005:  
-                clean_token = token.replace(' ', '').replace('Ġ', '').strip() 
-                if clean_token: 
-                    token_analysis.append({'token': clean_token, 'importance': score})
+                clean_token = token.replace(' ', '').replace('Ġ', '').strip()  # Handle different tokenizer prefixes
+                if clean_token:  # Only add if token has content after cleaning
+                    token_analysis.append({
+                        'token': clean_token,
+                        'importance': score,
+                        'position': i
+                    })
         
+        # Sort by importance
         token_analysis.sort(key=lambda x: x['importance'], reverse=True)
         
         detailed_analysis += f"""
 ## Top important tokens:
 <div style="background: #2d2d2d; padding: 15px; border-radius: 10px; margin: 15px 0; border-left: 4px solid #4caf50; color: #e0e0e0;">
-    Found {len(token_analysis)} important tokens
+    <strong>Analysis Info:</strong> Found {len(token_analysis)} important tokens out of {len(tokens)} total tokens
 </div>
 <div style="font-family: Arial, sans-serif;">
 """
-        for i, token_info in enumerate(token_analysis[:10]): 
+        
+        for i, token_info in enumerate(token_analysis[:10]):  # Top 10 tokens
             bar_width = int(token_info['importance'] * 100)
             color = "#ff4444" if predicted_class == "phishing" else "#44ff44"
+            
             detailed_analysis += f"""
 <div style="margin: 8px 0; display: flex; align-items: center; background: #2d2d2d; padding: 8px; border-radius: 8px; border-left: 4px solid {color};">
-    <div style="width: 30px; text-align: right; margin-right: 10px; font-weight: bold; color: #ffffff;">{i+1}.</div>
-    <div style="width: 120px; margin-right: 10px; font-weight: bold; color: #e0e0e0; text-align: right;">{token_info['token']}</div>
+    <div style="width: 30px; text-align: right; margin-right: 10px; font-weight: bold; color: #ffffff;">
+        {i+1}.
+    </div>
+    <div style="width: 120px; margin-right: 10px; font-weight: bold; color: #e0e0e0; text-align: right;">
+        {token_info['token']}
+    </div>
     <div style="width: 300px; background-color: #404040; border-radius: 10px; overflow: hidden; margin-right: 10px; border: 1px solid #555;">
-        <div style="width: {bar_width}%; background-color: {color}; height: 20px; border-radius: 10px;"></div>
+        <div style="width: {bar_width}%; background-color: {color}; height: 20px; border-radius: 10px; transition: width 0.3s ease;"></div>
+    </div>
+    <div style="color: #cccccc; font-size: 12px; font-weight: bold;">
+        {token_info['importance']:.1%}
     </div>
-    <div style="color: #cccccc; font-size: 12px; font-weight: bold;">{token_info['importance']:.1%}</div>
 </div>
 """
-        detailed_analysis += "</div>"
         
-        # Add stats and simple bars
+        detailed_analysis += "</div>\n"
+        
         detailed_analysis += f"""
+## Detailed analysis:
+<div style="font-family: Arial, sans-serif; background: linear-gradient(135deg, #1a237e 0%, #3949ab 100%); padding: 20px; border-radius: 15px; color: white; margin: 15px 0; border: 2px solid #3f51b5;">
+    <h3 style="margin: 0 0 15px 0; color: white;">Statistical Overview</h3>
+    <div style="display: grid; grid-template-columns: repeat(2, 1fr); gap: 15px;">
+        <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px; border: 1px solid rgba(255,255,255,0.2);">
+            <div style="font-size: 24px; font-weight: bold; color: white;">{len([t for t in tokens if t not in ['[CLS]', '[SEP]', '[PAD]']])}</div>
+            <div style="font-size: 14px; color: #e0e0e0;">Total tokens</div>
+        </div>
+        <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px; border: 1px solid rgba(255,255,255,0.2);">
+            <div style="font-size: 24px; font-weight: bold, color: white;">{len([t for t in token_analysis if t['importance'] > 0.05])}</div>
+            <div style="font-size: 14px, color: #e0e0e0;">High impact tokens (>5%)</div>
+        </div>
+    </div>
+</div>
 <div style="font-family: Arial, sans-serif; margin: 15px 0; background: #2d2d2d; padding: 20px; border-radius: 15px; border: 1px solid #555;">
     <h3 style="color: #ffffff; margin-bottom: 15px;"> Prediction Confidence</h3>
     <div style="display: flex; justify-content: space-between; margin-bottom: 10px;">
@@ -247,36 +301,66 @@ def predict_fn(text: str):
             {probs[1]:.1%}
         </div>
     </div>
+    <div style="margin-top: 10px; text-align: center; color: #cccccc; font-size: 14px;">
+        Benign: {probs[0]:.1%}
+    </div>
 </div>
 """
     else:
-        # Fallback when no attention
+        # Fallback analysis without attention weights
         detailed_analysis += f"""
-<div style="background: #2d2d2d; padding: 20px; border-radius: 15px; margin: 15px 0; border: 1px solid #555;">
-    <h3 style="color: #ffffff; margin: 0 0 15px 0;">Prediction:</h3>
-    <p style="color:white">Phishing: {probs[1]:.1%} | Benign: {probs[0]:.1%}</p>
-    <p style="color: #ffcc02;">Note: Detailed attention weights not available.</p>
+<div style="background: linear-gradient(135deg, #1a237e 0%, #3949ab 100%); padding: 20px; border-radius: 15px; color: white; margin: 15px 0; border: 2px solid #3f51b5;">
+    <h3 style="margin: 0 0 15px 0; color: white;">Basic Analysis</h3>
+    <div style="display: grid; grid-template-columns: repeat(3, 1fr); gap: 15px;">
+        <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px; text-align: center; border: 1px solid rgba(255,255,255,0.2);">
+            <div style="font-size: 24px; font-weight: bold; color: white;">{probs[1]:.1%}</div>
+            <div style="font-size: 14px; color: #e0e0e0;">Phishing</div>
+        </div>
+        <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px; text-align: center; border: 1px solid rgba(255,255,255,0.2);">
+            <div style="font-size: 24px; font-weight: bold; color: white;">{probs[0]:.1%}</div>
+            <div style="font-size: 14px; color: #e0e0e0;">Benign</div>
+        </div>
+        <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px; text-align: center; border: 1px solid rgba(255,255,255,0.2);">
+            <div style="font-size: 24px; font-weight: bold; color: white;">{len([t for t in tokens if t not in ['[CLS]', '[SEP]', '[PAD]']])}</div>
+            <div style="font-size: 14px; color: #e0e0e0;">Tokens</div>
+        </div>
+    </div>
+</div>
+<div style="font-family: Arial, sans-serif; margin: 15px 0; background: #2d2d2d; padding: 20px; border-radius: 15px; border: 1px solid #555;">
+    <h3 style="color: #ffffff; margin: 0 0 15px 0;">🔤 Tokens in text:</h3>
+    <div style="display: flex; flex-wrap: wrap; gap: 8px;">""" + ''.join([f'<span style="background: #404040; color: #64b5f6; padding: 4px 8px; border-radius: 15px; font-size: 12px; border: 1px solid #666;">{token.replace(" ", "")}</span>' for token in tokens if token not in ['[CLS]', '[SEP]', '[PAD]']]) + f"""</div>
+    <div style="margin-top: 15px; padding: 10px; background: #3d2914; border-radius: 8px; border-left: 4px solid #ff9800;">
+        <strong style="color: #ffcc02;">Debug info:</strong> <span style="color: #e0e0e0;">Found {len(tokens)} total tokens, {len([t for t in tokens if t not in ['[CLS]', '[SEP]', '[PAD]']])} content tokens</span>
+    </div>
+</div>
+<div style="background: #3d2914; padding: 15px; border-radius: 10px; border-left: 4px solid #ff9800; margin: 15px 0;">
+    <p style="margin: 0; color: #ffcc02; font-size: 14px;">
+        <strong>Note:</strong> Detailed attention weights analysis is not available for the current model.
+    </p>
 </div>
 """
+
+    # Build label->prob mapping for Gradio Label output
+    if len(LABELS) == len(probs):
+        prediction_result = {LABELS[i]: float(probs[i]) for i in range(len(LABELS))}
+    else:
+        prediction_result = {f"class_{i}": float(p) for i, p in enumerate(probs)}
     
-    return detailed_analysis
+    return prediction_result, detailed_analysis
 
 # --------- RAG Inference function ----------
 def rag_predict_fn(text: str):
     if not text or not text.strip():
         return "Please enter text to analyze."
     
-    if not RAG_AVAILABLE or rag_engine is None or llm_client is None:
-        return "⚠️ RAG/LLM system is not initialized properly. Check console logs."
+    # Call LLM (which now handles retrieval internally via LangChain)
+    response = llm_client.analyze(text)
+    
+    return response
 
-    try:
-        # 1. Retrieve context
-        context = rag_engine.retrieve(text)
-        # 2. Call LLM
-        response = llm_client.analyze(text, context)
-        return response
-    except Exception as e:
-        return f"Error during RAG analysis: {str(e)}"
+# --------- Refresh Knowledge Base function ----------
+def refresh_kb():
+    return rag_engine.refresh_knowledge_base()
 
 # --------- Gradio UI ----------
 css_style="""
@@ -288,11 +372,12 @@ css_style="""
 .dark .gradio-container {
     background-color: #1e1e1e !important;
 }
+/* Dark theme for all components */
 .block {
     background-color: #2d2d2d !important;
     border: 1px solid #444 !important;
 }
-.gradio-textbox, .gradio-dropdown {
+.gradio-textbox {
     background-color: #3d3d3d !important;
     color: #ffffff !important;
     border: 1px solid #666 !important;
@@ -304,65 +389,98 @@ css_style="""
 }
 .gradio-button:hover {
     background-color: #5a5a5a !important;
+    color: #ffffff !important;
+    border: 1px solid #666 !important;
 }
 """
-
-with gr.Blocks() as demo:
-    gr.Markdown("# 🛡️ Advanced Phishing Detector System")
+with gr.Blocks(css=css_style) as demo:
+    gr.Markdown("# 🛡️ Phishing Detector (DeBERTa + LSTM + RAG)")
     
     with gr.Tabs():
-        # --- Tab 1: DeBERTa Analysis ---
-        with gr.TabItem("DeBERTa + LSTM Detection"):
+        # --- Tab 1: Standard Detection ---
+        with gr.TabItem("🔍 Standard Detection"):
             gr.Markdown("""
-            ### Hybrid Detection Model
-            Enter a URL or text. The system uses **DeBERTa** combined with **LSTM** to analyze semantic meaning and structure.
-            For URLs, it fetches the HTML content and performs a weighted analysis (30% URL structure + 70% Content).
+            Enter a URL or text for analysis using the DeBERTa + LSTM model.
+            
+            **Features:**
+            - **URL Analysis**: For URLs, the system will fetch HTML content and combine both URL and content analysis
+            - **Combined Prediction**: Uses weighted combination of URL structure and webpage content analysis
+            - **Visual Analysis**: Predict phishing/benign probability with visual charts
+            - **Token Importance**: Display the most important tokens in classification
+            - **Detailed Insights**: Comprehensive analysis of the impact of each token
+            
+            **How it works for URLs:**
+            1. Analyze the URL structure itself
+            2. Fetch the webpage HTML content
+            3. Analyze the webpage content
+            4. Combine both results for final prediction (30% URL + 70% content)
             """)
             
             with gr.Row():
                 with gr.Column(scale=2):
                     input_box = gr.Textbox(
-                        label="Input URL or Text", 
-                        placeholder="Enter URL (e.g., http://example.com) or paste suspicious text...",
+                        label="URL or text", 
+                        placeholder="Example: http://suspicious-site.example or paste any text",
                         lines=3
                     )
-                    btn_submit = gr.Button("Analyze Security", variant="primary")
+                    btn_submit = gr.Button("🔍 Analyze", variant="primary")
                     
                     gr.Examples(
                         examples=[
                             ["http://rendmoiunserviceeee.com"],
                             ["https://www.google.com"],
-                            ["Dear customer, your account has been suspended. Click here to verify."],
+                            ["Dear customer, your account has been suspended. Click here to verify your identity immediately."],
+                            ["https://mail-secure-login-verify.example/path?token=suspicious"],
                             ["http://paypaI-security-update.net/login"],
+                            ["Your package has been delivered successfully. Thank you for using our service."],
+                            ["https://github.com/user/repo"],
+                            ["Dear customer, your account has been suspended. Click here to verify."],
                         ],
                         inputs=input_box
                     )
                     
                 with gr.Column(scale=3):
                     output_html = gr.HTML(label="Analysis Result")
-            
+
             btn_submit.click(fn=predict_fn, inputs=input_box, outputs=output_html)
 
         # --- Tab 2: LLM + RAG Analysis ---
-        with gr.TabItem("LLM + RAG Assistant"):
+        with gr.TabItem("🤖 AI Assistant (RAG)"):
             gr.Markdown("""
-            ### AI Security Analyst (Qwen2.5 + RAG)
-            This agent retrieves context from a security knowledge base and uses a Large Language Model to explain *why* something is suspicious.
+            **AI Assistant** uses **Qwen2.5-3B** + **LangChain** to explain *why* a message is suspicious.
+            
+            **Features:**
+            - 🌐 Multilingual support (English + Vietnamese)
+            - 📚 Knowledge Base retrieval (Auto-sync)
+            - 🚀 No rate limits (self-hosted)
             """)
             
             with gr.Row():
                 with gr.Column(scale=1):
                     rag_input = gr.Textbox(
-                        label="Suspicious Content / Question", 
-                        placeholder="Paste email content, headers, or ask: 'Why is this url suspicious?'...",
+                        label="Suspicious Text/URL",
+                        placeholder="Paste the email content or URL here...",
                         lines=5
                     )
-                    btn_rag = gr.Button("Ask AI Assistant", variant="primary")
+                    with gr.Row():
+                        btn_rag = gr.Button("🤖 Ask AI Assistant", variant="primary")
+                        btn_refresh = gr.Button("♻️ Refresh Knowledge Base")
                     
+                    gr.Examples(
+                        examples=[
+                            ["Your PayPal account has been suspended. Click http://paypal-verify.com to unlock."],
+                            ["Tài khoản ngân hàng của bạn bị khóa. Nhấn vào đây để mở khóa ngay."],
+                            ["Your package is ready for delivery. Track here: https://fedex-track.com"],
+                        ],
+                        inputs=rag_input
+                    )
+                
                 with gr.Column(scale=1):
                     rag_output = gr.Markdown(label="AI Analysis")
+                    refresh_output = gr.Markdown(label="Status")
             
             btn_rag.click(fn=rag_predict_fn, inputs=[rag_input], outputs=rag_output)
+            btn_refresh.click(fn=refresh_kb, inputs=[], outputs=refresh_output)
 
 if __name__ == "__main__":
     demo.launch()