update

app.py

@@ -169,6 +169,8 @@ def run_full_pipeline(conversation: str, wingman_prompt: str = "", gemini_model_
         }
 
         # Model 5 – Google Gemini API
+        gemini_reply = ""
+        gemini_error = ""
         try:
             gemini_service = get_gemini_service(model_name=gemini_model_name)
             # Format conversation for Gemini: "Male: ... ||| Female: ..."
@@ -180,8 +182,25 @@ def run_full_pipeline(conversation: str, wingman_prompt: str = "", gemini_model_
             )
             gemini_error = ""
         except Exception as exc:
-            gemini_reply = ""
-            gemini_error = str(exc)
+            error_msg = str(exc)
+            # Try fallback to gemini-2.0-flash if current model fails (especially for MAX_TOKENS with 0 parts)
+            fallback_model = "gemini-2.0-flash"
+            if gemini_model_name != fallback_model and ("MAX_TOKENS" in error_msg or "quota" in error_msg.lower() or "429" in error_msg or "no text" in error_msg.lower()):
+                try:
+                    gemini_service = get_gemini_service(model_name=fallback_model)
+                    formatted_conversation = f"Male: {male} ||| Female: {female}"
+                    gemini_reply = gemini_service.generate_reply(
+                        conversation=formatted_conversation,
+                        trigger=trigger,
+                        move=move,
+                    )
+                    gemini_error = f"⚠️ Original model ({gemini_model_name}) failed. Used fallback: {fallback_model}"
+                except Exception as fallback_exc:
+                    gemini_reply = ""
+                    gemini_error = f"Model {gemini_model_name} failed: {error_msg[:150]}. Fallback ({fallback_model}) also failed: {str(fallback_exc)[:150]}"
+            else:
+                gemini_reply = ""
+                gemini_error = error_msg
 
         models_output["gemini"] = {
             "label": f"Model 5 – Gemini API ({gemini_model_name})",
@@ -279,22 +298,39 @@ with gr.Blocks(title=title) as demo:
                 )
                 gr.Markdown("Leave as-is for default behavior. Edits apply to Model 3 when its LoRA is used.")
             
-            # Model 5 – Gemini Model Selection
+            # Model 5 – Gemini Model Selection (using whitelist - 15 tested models)
             try:
-                gemini_models = get_available_gemini_models()
-                gemini_model_choices = [model["name"] for model in gemini_models]
-                gemini_model_display = [f"{model['displayName']} ({model['name']})" for model in gemini_models]
-                default_gemini_model = gemini_model_choices[0] if gemini_model_choices else "gemini-2.5-flash"
+                gemini_models = get_available_gemini_models(use_whitelist=True)
+                # Create choices with (label, value) format for better display
+                gemini_dropdown_choices = []
+                gemini_model_choices = []  # Keep for value matching
+                
+                for model in gemini_models:
+                    model_name = model["name"]
+                    display_name = model.get("displayName", model_name)
+                    # Remove "models/" prefix for cleaner display
+                    clean_name = model_name.replace("models/", "")
+                    # Format: (display_label, actual_value)
+                    label = f"{display_name} ({clean_name})"
+                    gemini_dropdown_choices.append((label, model_name))
+                    gemini_model_choices.append(model_name)
+                
+                # Default to gemini-2.0-flash (first in whitelist) or first available
+                default_gemini_model = "models/gemini-2.0-flash" if "models/gemini-2.0-flash" in gemini_model_choices else (gemini_model_choices[0] if gemini_model_choices else "models/gemini-2.0-flash")
+                
+                print(f"✓ Loaded {len(gemini_model_choices)} whitelisted Gemini models for dropdown")
             except Exception as e:
-                gemini_model_choices = ["gemini-2.5-flash"]
-                gemini_model_display = ["Gemini 2.5 Flash (default - API key may be missing)"]
-                default_gemini_model = "gemini-2.5-flash"
+                print(f"⚠ Error loading Gemini whitelist: {e}")
+                gemini_dropdown_choices = [("Gemini 2.0 Flash (default - API key may be missing)", "models/gemini-2.0-flash")]
+                gemini_model_choices = ["models/gemini-2.0-flash"]
+                default_gemini_model = "models/gemini-2.0-flash"
             
             gemini_model_dropdown = gr.Dropdown(
-                choices=gemini_model_choices,
+                choices=gemini_dropdown_choices,
                 value=default_gemini_model,
                 label="Model 5 – Select Gemini Model",
-                info="Choose which Gemini model to use for reply generation",
+                info=f"Choose from {len(gemini_model_choices)} tested and working Gemini models (whitelist)",
+                interactive=True,
             )
             
             reply_btn = gr.Button("Generate Reply Suggestion", variant="primary", size="lg")

gemini_models_whitelist.json

@@ -0,0 +1,117 @@
+{
+  "whitelist": [
+    {
+      "name": "models/gemini-2.0-flash",
+      "displayName": "Gemini 2.0 Flash",
+      "description": "Gemini 2.0 Flash",
+      "version": "2.0",
+      "test_reply": "Mai anh rảnh đó, hay mình đổi gió đi đâu vui vui em nhỉ."
+    },
+    {
+      "name": "models/gemini-2.0-flash-001",
+      "displayName": "Gemini 2.0 Flash 001",
+      "description": "Stable version of Gemini 2.0 Flash, our fast and versatile multimodal model for scaling across diverse tasks, released in January of 2025.",
+      "version": "2.0",
+      "test_reply": "Mai anh rảnh đó, hay mình đổi gió đi đâu vui vui em nhỉ."
+    },
+    {
+      "name": "models/gemini-2.0-flash-lite-001",
+      "displayName": "Gemini 2.0 Flash-Lite 001",
+      "description": "Stable version of Gemini 2.0 Flash-Lite",
+      "version": "2.0",
+      "test_reply": "Mai anh rảnh, mình đi xem phim được không em."
+    },
+    {
+      "name": "models/gemini-2.0-flash-lite",
+      "displayName": "Gemini 2.0 Flash-Lite",
+      "description": "Gemini 2.0 Flash-Lite",
+      "version": "2.0",
+      "test_reply": "Mai anh rảnh, mình đi xem phim được không em."
+    },
+    {
+      "name": "models/gemini-2.0-flash-lite-preview-02-05",
+      "displayName": "Gemini 2.0 Flash-Lite Preview 02-05",
+      "description": "Preview release (February 5th, 2025) of Gemini 2.0 Flash-Lite",
+      "version": "preview-02-05",
+      "test_reply": "Mai anh rảnh, mình đi xem phim được không em."
+    },
+    {
+      "name": "models/gemini-2.0-flash-lite-preview",
+      "displayName": "Gemini 2.0 Flash-Lite Preview",
+      "description": "Preview release (February 5th, 2025) of Gemini 2.0 Flash-Lite",
+      "version": "preview-02-05",
+      "test_reply": "Mai anh rảnh, mình đi xem phim được không em."
+    },
+    {
+      "name": "models/gemma-3-1b-it",
+      "displayName": "Gemma 3 1B",
+      "description": "",
+      "version": "001",
+      "test_reply": "Anh cũng vậy, mong được gặp mặt em nhé."
+    },
+    {
+      "name": "models/gemma-3-4b-it",
+      "displayName": "Gemma 3 4B",
+      "description": "",
+      "version": "001",
+      "test_reply": "Mai thì em cứ hẹn anh nha, anh rảnh lắm đó ạ."
+    },
+    {
+      "name": "models/gemma-3-12b-it",
+      "displayName": "Gemma 3 12B",
+      "description": "",
+      "version": "001",
+      "test_reply": "Mai anh rảnh, em bảo anh đưa em đi đâu chơi nhé."
+    },
+    {
+      "name": "models/gemma-3-27b-it",
+      "displayName": "Gemma 3 27B",
+      "description": "",
+      "version": "001",
+      "test_reply": "Mai em rảnh thật ạ."
+    },
+    {
+      "name": "models/gemma-3n-e4b-it",
+      "displayName": "Gemma 3n E4B",
+      "description": "",
+      "version": "001",
+      "test_reply": "Dạ được em, mai anh rảnh nhé."
+    },
+    {
+      "name": "models/gemma-3n-e2b-it",
+      "displayName": "Gemma 3n E2B",
+      "description": "",
+      "version": "001",
+      "test_reply": "Mai em nhé, anh sẽ cố gắng sắp xếp ạ."
+    },
+    {
+      "name": "models/gemini-flash-lite-latest",
+      "displayName": "Gemini Flash-Lite Latest",
+      "description": "Latest release of Gemini Flash-Lite",
+      "version": "Gemini Flash-Lite Latest",
+      "test_reply": "Mai anh rảnh, mình đi đâu em thích nhé."
+    },
+    {
+      "name": "models/gemini-2.5-flash-lite",
+      "displayName": "Gemini 2.5 Flash-Lite",
+      "description": "Stable version of Gemini 2.5 Flash-Lite, released in July of 2025",
+      "version": "001",
+      "test_reply": "Mai anh rảnh, mình đi cà phê nha em."
+    },
+    {
+      "name": "models/gemini-2.5-flash-lite-preview-09-2025",
+      "displayName": "Gemini 2.5 Flash-Lite Preview Sep 2025",
+      "description": "Preview release (Septempber 25th, 2025) of Gemini 2.5 Flash-Lite",
+      "version": "2.5-preview-09-25",
+      "test_reply": "Mai anh rảnh, mình đi cà phê nhé em."
+    }
+  ],
+  "failed": 26,
+  "test_data": {
+    "conversation": "Male: Tối nay anh có lịch đột xuất. ||| Female: Thế mai được không?",
+    "trigger": "neutral",
+    "move": "escalate"
+  },
+  "total_tested": 41,
+  "passed": 15
+}

gemini_service.py

@@ -2,7 +2,9 @@
 Service for generating replies using Google Gemini API.
 """
 import os
+import json
 import requests
+from pathlib import Path
 from typing import Optional, List, Dict, Any
 import google.generativeai as genai
 
@@ -111,7 +113,7 @@ def fetch_gemini_models(api_key: Optional[str] = None) -> List[Dict[str, Any]]:
 class GeminiReplyService:
     """Service for generating replies using Google Gemini API."""
     
-    def __init__(self, api_key: Optional[str] = None, model_name: str = "gemini-2.5-flash"):
+    def __init__(self, api_key: Optional[str] = None, model_name: str = "gemini-2.0-flash"):
         """
         Initialize Gemini service.
         
@@ -145,7 +147,7 @@ class GeminiReplyService:
         trigger: str,
         move: str,
         temperature: float = 0.2,
-        max_output_tokens: int = 80,
+        max_output_tokens: int = 500,  # Increased significantly to avoid truncation issues
     ) -> str:
         """
         Generate reply using Google Gemini API.
@@ -177,9 +179,65 @@ MOVE: "{move}"
                     temperature=temperature,
                     max_output_tokens=max_output_tokens,
                 ),
+                safety_settings=[
+                    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
+                    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
+                    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
+                    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
+                ],
             )
             
-            raw = response.text.strip() if response.text else ""
+            # Check if response has text
+            try:
+                raw = response.text.strip()
+            except Exception as text_error:
+                # Check finish reason if text access fails
+                if response.candidates and len(response.candidates) > 0:
+                    candidate = response.candidates[0]
+                    finish_reason = candidate.finish_reason
+                    
+                    # Map finish_reason codes to names
+                    finish_reason_map = {
+                        0: "FINISH_REASON_UNSPECIFIED",
+                        1: "STOP",
+                        2: "MAX_TOKENS",
+                        3: "SAFETY",
+                        4: "RECITATION",
+                        5: "OTHER"
+                    }
+                    finish_reason_name = finish_reason_map.get(finish_reason, f"UNKNOWN({finish_reason})")
+                    
+                    if finish_reason == 2:  # MAX_TOKENS - response was truncated
+                        # Try to get partial text if available
+                        try:
+                            if hasattr(candidate, 'content') and candidate.content:
+                                if hasattr(candidate.content, 'parts') and candidate.content.parts:
+                                    parts_text = []
+                                    for part in candidate.content.parts:
+                                        if hasattr(part, 'text') and part.text:
+                                            parts_text.append(part.text)
+                                    if parts_text:
+                                        raw = " ".join(parts_text).strip()
+                                        # Continue with processing below
+                                    else:
+                                        raise Exception(f"Response truncated due to max_tokens limit (finish_reason: {finish_reason_name}). No text parts found. Try increasing max_output_tokens.")
+                                else:
+                                    raise Exception(f"Response truncated due to max_tokens limit (finish_reason: {finish_reason_name}). No content parts. Try increasing max_output_tokens.")
+                            else:
+                                raise Exception(f"Response truncated due to max_tokens limit (finish_reason: {finish_reason_name}). No candidate content. Try increasing max_output_tokens.")
+                        except Exception as parts_error:
+                            # If we couldn't extract partial text, raise the original error
+                            raise Exception(f"Response truncated due to max_tokens limit (finish_reason: {finish_reason_name}). Could not extract partial text: {str(parts_error)}. Try increasing max_output_tokens.")
+                        else:
+                            raise Exception(f"Response truncated due to max_tokens limit (finish_reason: {finish_reason_name}). Try increasing max_output_tokens.")
+                    elif finish_reason == 3:  # SAFETY
+                        raise Exception(f"Response blocked by safety filter (finish_reason: {finish_reason_name}). The content may have triggered safety filters. Try adjusting the prompt or using a different model.")
+                    elif finish_reason == 4:  # RECITATION
+                        raise Exception(f"Response blocked due to recitation policy (finish_reason: {finish_reason_name}).")
+                    else:
+                        raise Exception(f"Response has no text. Finish reason: {finish_reason_name} ({finish_reason}). Error: {str(text_error)}")
+                else:
+                    raise Exception(f"Response has no candidates. Error: {str(text_error)}")
             
             # Hậu xử lý: lấy câu đầu, giới hạn 25 từ
             import re
@@ -198,7 +256,14 @@ MOVE: "{move}"
             return limited
             
         except Exception as e:
-            raise Exception(f"Gemini API error: {str(e)}")
+            error_msg = str(e)
+            # Check for quota errors and provide clearer message
+            if "429" in error_msg or "quota" in error_msg.lower() or "Quota exceeded" in error_msg:
+                raise Exception(
+                    f"Gemini API quota exceeded. Please check your billing/plan at https://ai.google.dev/gemini-api/docs/rate-limits. "
+                    f"Original error: {error_msg[:200]}"
+                )
+            raise Exception(f"Gemini API error: {error_msg}")
 
 
 # Global singleton instance
@@ -208,7 +273,7 @@ _cached_models = None
 
 def get_gemini_service(
     api_key: Optional[str] = None,
-    model_name: str = "gemini-2.5-flash",
+    model_name: str = "gemini-2.0-flash",
 ) -> GeminiReplyService:
     """Get or create the global Gemini service instance."""
     global _gemini_service
@@ -219,18 +284,41 @@ def get_gemini_service(
     return _gemini_service
 
 
-def get_available_gemini_models(api_key: Optional[str] = None, use_cache: bool = True) -> List[Dict[str, Any]]:
+def get_available_gemini_models(api_key: Optional[str] = None, use_cache: bool = True, use_whitelist: bool = True) -> List[Dict[str, Any]]:
     """
     Get list of available Gemini models that support generateContent.
+    By default, only returns models from whitelist (tested and working models).
     
     Args:
         api_key: Google API key. If None, will try to get from GOOGLE_API_KEY env var.
         use_cache: Whether to use cached models list (default: True)
+        use_whitelist: Whether to filter by whitelist (default: True)
     
     Returns:
-        List of model dictionaries
+        List of model dictionaries (whitelisted models only if use_whitelist=True)
     """
     global _cached_models
+    
+    # Load whitelist if requested
+    if use_whitelist:
+        whitelist_path = Path(__file__).parent / "gemini_models_whitelist.json"
+        if whitelist_path.exists():
+            try:
+                import json
+                with open(whitelist_path, "r", encoding="utf-8") as f:
+                    whitelist_data = json.load(f)
+                    whitelist = whitelist_data.get("whitelist", [])
+                    if whitelist:
+                        if use_cache and _cached_models is not None:
+                            return _cached_models
+                        _cached_models = whitelist
+                        return whitelist
+            except Exception as e:
+                print(f"Warning: Could not load whitelist: {e}. Falling back to fetching all models.")
+        else:
+            print(f"Warning: Whitelist file not found at {whitelist_path}. Falling back to fetching all models.")
+    
+    # Fallback to fetching all models
     if use_cache and _cached_models is not None:
         return _cached_models
     

perplexity_service.py

@@ -62,13 +62,16 @@ Nếu vì bất kỳ lý do gì bạn không thể tuân thủ tất cả quy t
 class PerplexityReplyService:
     """Service for generating replies using Perplexity API."""
     
-    def __init__(self, api_key: Optional[str] = None, model: str = "mistral-7b-instruct"):
+    def __init__(self, api_key: Optional[str] = None, model: str = "sonar"):
         """
         Initialize Perplexity service.
         
         Args:
             api_key: Perplexity API key. If None, will try to get from PERPLEXITY_API_KEY env var.
-            model: Model name to use (default: "mistral-7b-instruct")
+            model: Model name to use (default: "sonar")
+                    Valid models: sonar, sonar-pro, llama-3.1-sonar-small-32k-online, etc.
+                    See: https://docs.perplexity.ai/getting-started/models
+                    Note: Model names may vary. Check documentation for current valid models.
         """
         self.api_key = api_key or os.getenv("PERPLEXITY_API_KEY")
         if not self.api_key:
@@ -157,7 +160,7 @@ _perplexity_service = None
 
 def get_perplexity_service(
     api_key: Optional[str] = None,
-    model: str = "mistral-7b-instruct",
+    model: str = "sonar",
 ) -> PerplexityReplyService:
     """Get or create the global Perplexity service instance."""
     global _perplexity_service

test_all_gemini_models.py

@@ -0,0 +1,196 @@
+"""
+Test all Gemini models and create whitelist of working models.
+"""
+import os
+import json
+from pathlib import Path
+from dotenv import load_dotenv
+from gemini_service import fetch_gemini_models, get_gemini_service
+
+# Load .env file
+env_path = Path(__file__).parent / '.env'
+if env_path.exists():
+    load_dotenv(env_path)
+
+# Test data
+TEST_CONVERSATION = "Male: Tối nay anh có lịch đột xuất. ||| Female: Thế mai được không?"
+TEST_TRIGGER = "neutral"
+TEST_MOVE = "escalate"
+
+SYSTEM_PROMPT = """
+Bạn là một wingman AI tinh tế, chuyên giúp Nam soạn 1 tin nhắn trả lời duy nhất trong hội thoại hẹn hò tiếng Việt. Bạn luôn nhìn từ góc nhìn của Nam, xưng "anh" và gọi đối phương là "em".
+
+Bạn được cung cấp:
+
+- HỘI THOẠI: đoạn hội thoại gần nhất giữa Nam (Male) và Nữ (Female), phân tách các tin bằng ký hiệu "|||".
+
+- TRIGGER: intent hiện tại (ví dụ: neutral, positive, negative, confused...).
+
+- MOVE: chiến lược hiện tại (ví dụ: escalate, hold, de-escalate, tease, comfort...).
+
+Nhiệm vụ của bạn:
+
+- Dựa trên HỘI THOẠI + TRIGGER + MOVE, hãy chọn một hướng phản hồi tự nhiên, duyên dáng, đúng chiến lược (không quá đẩy hay quá lùi so với MOVE).
+
+- Ưu tiên giữ mạch cảm xúc nhất quán với hội thoại, tránh tạo thông tin fact mới về thế giới bên ngoài hoặc về hai người.
+
+QUY TẮC CỨNG:
+
+- Chỉ trả về đúng 1 câu duy nhất.
+
+- Tối đa 25 từ tiếng Việt.
+
+- Lịch sự, ấm áp, thân thiện; không phán xét, không thô lỗ.
+
+- Không giải thích meta (không nói về "prompt", "AI", "chiến lược", "MOVE", "TRIGGER"...).
+
+- Không lặp lại nguyên văn câu của đối phương.
+
+- Không thêm fact mới (chỉ dựa trên những gì có trong hội thoại, hoặc các câu nói chung chung, không cụ thể hóa thông tin chưa có).
+
+Khi TRIGGER hoặc MOVE có vẻ mâu thuẫn với HỘI THOẠI:
+
+- Hãy ưu tiên sự an toàn và mềm mại.
+
+- Có thể hỏi lại nhẹ nhàng để làm rõ, nhưng vẫn giữ frame chủ động, tự tin của Nam.
+
+PHONG CÁCH:
+
+- Ấm áp, tự tin nhưng không tự cao.
+
+- Có thể dùng từ đệm tự nhiên (nha, nhé, ạ, dạ) khi phù hợp với ngữ cảnh.
+
+- Phản chiếu cảm xúc của đối phương.
+
+- Giữ mạch trò chuyện mở để còn đất tăng tương tác về sau.
+
+Nếu vì bất kỳ lý do gì bạn không thể tuân thủ tất cả quy tắc trên:
+
+- Hãy ưu tiên vẫn trả về đúng 1 câu, ≤25 từ, không chứa meta, không chứa thông tin fact mới.
+""".strip()
+
+
+def test_model(model_name: str, max_retries: int = 2) -> tuple[bool, str, str]:
+    """
+    Test a single Gemini model.
+    
+    Returns:
+        (success: bool, reply: str, error: str)
+    """
+    for attempt in range(max_retries):
+        try:
+            service = get_gemini_service(model_name=model_name)
+            formatted_conversation = f"Male: {TEST_CONVERSATION.split('|||')[0].strip()} ||| Female: {TEST_CONVERSATION.split('|||')[1].strip()}"
+            
+            reply = service.generate_reply(
+                conversation=formatted_conversation,
+                trigger=TEST_TRIGGER,
+                move=TEST_MOVE,
+                max_output_tokens=200,
+            )
+            
+            if reply and len(reply.strip()) > 0:
+                return True, reply, ""
+            else:
+                return False, "", "Empty reply"
+                
+        except Exception as e:
+            error_msg = str(e)
+            if attempt < max_retries - 1:
+                continue  # Retry
+            return False, "", error_msg
+    
+    return False, "", "Max retries exceeded"
+
+
+def main():
+    """Test all Gemini models and create whitelist."""
+    print("=" * 60)
+    print("Testing All Gemini Models for Whitelist")
+    print("=" * 60)
+    
+    # Fetch all available models
+    try:
+        print("\nFetching available Gemini models...")
+        all_models = fetch_gemini_models()
+        print(f"✓ Found {len(all_models)} models with generateContent support")
+    except Exception as e:
+        print(f"✗ Error fetching models: {str(e)}")
+        return 1
+    
+    # Test each model
+    whitelist = []
+    failed_models = []
+    
+    print(f"\nTesting {len(all_models)} models...")
+    print("=" * 60)
+    
+    for idx, model_info in enumerate(all_models, 1):
+        model_name = model_info["name"]
+        display_name = model_info.get("displayName", model_name)
+        
+        print(f"\n[{idx}/{len(all_models)}] Testing: {display_name}")
+        print(f"  Model name: {model_name}")
+        
+        success, reply, error = test_model(model_name)
+        
+        if success:
+            print(f"  ✓ PASSED - Reply: {reply[:60]}...")
+            whitelist.append({
+                "name": model_name,
+                "displayName": display_name,
+                "description": model_info.get("description", ""),
+                "version": model_info.get("version", ""),
+                "test_reply": reply[:100],  # Store sample reply
+            })
+        else:
+            print(f"  ✗ FAILED - {error[:100]}")
+            failed_models.append({
+                "name": model_name,
+                "displayName": display_name,
+                "error": error[:200],
+            })
+    
+    # Save whitelist
+    whitelist_file = Path(__file__).parent / "gemini_models_whitelist.json"
+    with open(whitelist_file, "w", encoding="utf-8") as f:
+        json.dump({
+            "whitelist": whitelist,
+            "failed": failed_models,
+            "test_data": {
+                "conversation": TEST_CONVERSATION,
+                "trigger": TEST_TRIGGER,
+                "move": TEST_MOVE,
+            },
+            "total_tested": len(all_models),
+            "passed": len(whitelist),
+            "failed": len(failed_models),
+        }, f, indent=2, ensure_ascii=False)
+    
+    # Print summary
+    print("\n" + "=" * 60)
+    print("Test Summary")
+    print("=" * 60)
+    print(f"Total models tested: {len(all_models)}")
+    print(f"✓ Passed (whitelist): {len(whitelist)}")
+    print(f"✗ Failed: {len(failed_models)}")
+    print(f"\nWhitelist saved to: {whitelist_file}")
+    
+    if whitelist:
+        print("\n✓ Working models (whitelist):")
+        for model in whitelist:
+            print(f"  - {model['displayName']} ({model['name']})")
+    
+    if failed_models:
+        print("\n✗ Failed models:")
+        for model in failed_models[:10]:  # Show first 10
+            print(f"  - {model['displayName']} ({model['name']}): {model['error'][:50]}...")
+        if len(failed_models) > 10:
+            print(f"  ... and {len(failed_models) - 10} more")
+    
+    return 0 if whitelist else 1
+
+
+if __name__ == "__main__":
+    exit(main())
+

test_api_models.py

@@ -0,0 +1,136 @@
+"""
+Test script for Model 4 (Perplexity) and Model 5 (Gemini) APIs.
+"""
+import os
+from pathlib import Path
+from dotenv import load_dotenv
+from perplexity_service import get_perplexity_service
+from gemini_service import get_gemini_service
+
+# Load .env file
+env_path = Path(__file__).parent / '.env'
+if env_path.exists():
+    load_dotenv(env_path)
+    print(f"✓ Loaded .env file from {env_path}")
+else:
+    print(f"⚠ .env file not found at {env_path}, using environment variables only")
+
+# Test data
+TEST_CONVERSATION = "Tối nay anh có lịch đột xuất. ||| Thế mai được không?"
+TEST_TRIGGER = "neutral"
+TEST_MOVE = "escalate"
+
+def test_perplexity():
+    """Test Perplexity API (Model 4)."""
+    print("=" * 60)
+    print("Testing Model 4 - Perplexity API")
+    print("=" * 60)
+    
+    try:
+        service = get_perplexity_service()
+        print(f"✓ Perplexity service initialized")
+        print(f"  Model: {service.model}")
+        print(f"  API Key: {'Set' if service.api_key else 'Not set'}")
+        
+        print("\nGenerating reply...")
+        reply = service.generate_reply(
+            conversation=TEST_CONVERSATION,
+            trigger=TEST_TRIGGER,
+            move=TEST_MOVE,
+        )
+        
+        print(f"\n✓ Success!")
+        print(f"  Input conversation: {TEST_CONVERSATION}")
+        print(f"  Trigger: {TEST_TRIGGER}")
+        print(f"  Move: {TEST_MOVE}")
+        print(f"  Generated reply: {reply}")
+        print(f"  Reply length: {len(reply.split())} words")
+        
+        return True, reply
+        
+    except Exception as e:
+        print(f"\n✗ Error: {str(e)}")
+        return False, str(e)
+
+
+def test_gemini():
+    """Test Gemini API (Model 5)."""
+    print("\n" + "=" * 60)
+    print("Testing Model 5 - Gemini API")
+    print("=" * 60)
+    
+    try:
+        # Try default model first
+        model_name = "gemini-2.0-flash"
+        service = get_gemini_service(model_name=model_name)
+        print(f"✓ Gemini service initialized")
+        print(f"  Model: {service.model_name}")
+        print(f"  API Key: {'Set' if service.api_key else 'Not set'}")
+        
+        print("\nGenerating reply...")
+        formatted_conversation = f"Male: {TEST_CONVERSATION.split('|||')[0].strip()} ||| Female: {TEST_CONVERSATION.split('|||')[1].strip()}"
+        reply = service.generate_reply(
+            conversation=formatted_conversation,
+            trigger=TEST_TRIGGER,
+            move=TEST_MOVE,
+            max_output_tokens=500,  # Use higher limit for testing
+        )
+        
+        print(f"\n✓ Success!")
+        print(f"  Input conversation: {formatted_conversation}")
+        print(f"  Trigger: {TEST_TRIGGER}")
+        print(f"  Move: {TEST_MOVE}")
+        print(f"  Generated reply: {reply}")
+        print(f"  Reply length: {len(reply.split())} words")
+        
+        return True, reply
+        
+    except Exception as e:
+        print(f"\n✗ Error: {str(e)}")
+        return False, str(e)
+
+
+def main():
+    """Run all tests."""
+    print("\n" + "=" * 60)
+    print("API Endpoint Test Script")
+    print("=" * 60)
+    print(f"\nTest Data:")
+    print(f"  Conversation: {TEST_CONVERSATION}")
+    print(f"  Trigger: {TEST_TRIGGER}")
+    print(f"  Move: {TEST_MOVE}")
+    print()
+    
+    # Check environment variables
+    print("Environment Variables:")
+    perplexity_key = os.getenv("PERPLEXITY_API_KEY")
+    google_key = os.getenv("GOOGLE_API_KEY")
+    print(f"  PERPLEXITY_API_KEY: {'✓ Set' if perplexity_key else '✗ Not set'}")
+    print(f"  GOOGLE_API_KEY: {'✓ Set' if google_key else '✗ Not set'}")
+    print()
+    
+    # Test Perplexity
+    perplexity_success, perplexity_result = test_perplexity()
+    
+    # Test Gemini
+    gemini_success, gemini_result = test_gemini()
+    
+    # Summary
+    print("\n" + "=" * 60)
+    print("Test Summary")
+    print("=" * 60)
+    print(f"  Model 4 (Perplexity): {'✓ PASSED' if perplexity_success else '✗ FAILED'}")
+    print(f"  Model 5 (Gemini): {'✓ PASSED' if gemini_success else '✗ FAILED'}")
+    print()
+    
+    if perplexity_success and gemini_success:
+        print("✓ All tests passed!")
+        return 0
+    else:
+        print("✗ Some tests failed. Check errors above.")
+        return 1
+
+
+if __name__ == "__main__":
+    exit(main())
+