update

app.py

@@ -4,6 +4,8 @@ from typing import Dict, Any, Tuple
 
 from reply_service import get_reply_service
 from trigger_move_identifier import get_trigger_move_identifier
+from perplexity_service import get_perplexity_service
+from gemini_service import get_gemini_service, get_available_gemini_models
 
 
 TRIGGER_MODEL_DIR = "./models/trigger_detector"
@@ -75,8 +77,8 @@ def parse_conversation(text: str) -> Tuple[str, str]:
     return male, female
 
 
-def run_full_pipeline(conversation: str, wingman_prompt: str = "") -> Dict[str, Any]:
-    """Run trigger detector and generate replies from three prompt styles (1 backend)."""
+def run_full_pipeline(conversation: str, wingman_prompt: str = "", gemini_model_name: str = "gemini-2.5-flash") -> Dict[str, Any]:
+    """Run trigger detector and generate replies from 5 models (3 prompt styles + Perplexity + Gemini)."""
     try:
         male, female = parse_conversation(conversation)
         identifier = get_trigger_move_identifier(
@@ -127,6 +129,21 @@ def run_full_pipeline(conversation: str, wingman_prompt: str = "") -> Dict[str,
             wingman_reply = ""
             wingman_error = str(exc)
 
+        # Model 4 – Perplexity API
+        try:
+            perplexity_service = get_perplexity_service()
+            # Format conversation for Perplexity: "Male: ... ||| Female: ..."
+            formatted_conversation = f"Male: {male} ||| Female: {female}"
+            perplexity_reply = perplexity_service.generate_reply(
+                conversation=formatted_conversation,
+                trigger=trigger,
+                move=move,
+            )
+            perplexity_error = ""
+        except Exception as exc:
+            perplexity_reply = ""
+            perplexity_error = str(exc)
+
         models_output["llama"] = {
             "label": "Model 1 – Prompt style: an toàn / nhẹ nhàng",
             "reply": llama_reply,
@@ -144,6 +161,33 @@ def run_full_pipeline(conversation: str, wingman_prompt: str = "") -> Dict[str,
             "reply": wingman_reply,
             "error": wingman_error,
         }
+
+        models_output["perplexity"] = {
+            "label": "Model 4 – Perplexity API",
+            "reply": perplexity_reply,
+            "error": perplexity_error,
+        }
+
+        # Model 5 – Google Gemini API
+        try:
+            gemini_service = get_gemini_service(model_name=gemini_model_name)
+            # Format conversation for Gemini: "Male: ... ||| Female: ..."
+            formatted_conversation = f"Male: {male} ||| Female: {female}"
+            gemini_reply = gemini_service.generate_reply(
+                conversation=formatted_conversation,
+                trigger=trigger,
+                move=move,
+            )
+            gemini_error = ""
+        except Exception as exc:
+            gemini_reply = ""
+            gemini_error = str(exc)
+
+        models_output["gemini"] = {
+            "label": f"Model 5 – Gemini API ({gemini_model_name})",
+            "reply": gemini_reply,
+            "error": gemini_error,
+        }
         
         return {
             "trigger": trigger,
@@ -201,8 +245,8 @@ with gr.Blocks(title=title) as demo:
     gr.Markdown("---")
     
     # Main Reply Suggestion Tab
-    gr.Markdown("### 🎯 Generate AI Reply Suggestions (3 Styles, 1 Backend)")
-    gr.Markdown("Nhập hội thoại và hệ thống sẽ chạy pipeline Trigger → Move → 3 style prompt (an toàn, flirt tinh tế, wingman).")
+    gr.Markdown("### 🎯 Generate AI Reply Suggestions (5 Models)")
+    gr.Markdown("Nhập hội thoại và hệ thống sẽ chạy pipeline Trigger → Move → 5 models (3 prompt styles + Perplexity API + Gemini API).")
     
     with gr.Row():
         with gr.Column(scale=2):
@@ -234,6 +278,25 @@ with gr.Blocks(title=title) as demo:
                     info="You can reference {conversation}, {trigger}, {move} inside your prompt.",
                 )
                 gr.Markdown("Leave as-is for default behavior. Edits apply to Model 3 when its LoRA is used.")
+            
+            # Model 5 – Gemini Model Selection
+            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"
+            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"
+            
+            gemini_model_dropdown = gr.Dropdown(
+                choices=gemini_model_choices,
+                value=default_gemini_model,
+                label="Model 5 – Select Gemini Model",
+                info="Choose which Gemini model to use for reply generation",
+            )
+            
             reply_btn = gr.Button("Generate Reply Suggestion", variant="primary", size="lg")
     
     reply_out = gr.JSON(
@@ -262,12 +325,26 @@ with gr.Blocks(title=title) as demo:
         placeholder="Reply từ mô hình Wingman LoRA (hoặc fallback prompt) sẽ xuất hiện tại đây.",
     )
 
-    def generate_reply_with_extraction(conversation: str, wingman_prompt: str) -> Tuple[Dict[str, Any], str, str, str]:
-        """Generate replies from three models."""
-        result = run_full_pipeline(conversation, wingman_prompt)
+    perplexity_box = gr.Textbox(
+        lines=3,
+        label="Model 4 – Perplexity API",
+        interactive=False,
+        placeholder="Reply từ Perplexity API sẽ xuất hiện tại đây.",
+    )
+
+    gemini_box = gr.Textbox(
+        lines=3,
+        label="Model 5 – Gemini API",
+        interactive=False,
+        placeholder="Reply từ Gemini API sẽ xuất hiện tại đây.",
+    )
+
+    def generate_reply_with_extraction(conversation: str, wingman_prompt: str, gemini_model_name: str) -> Tuple[Dict[str, Any], str, str, str, str, str]:
+        """Generate replies from five models."""
+        result = run_full_pipeline(conversation, wingman_prompt, gemini_model_name)
         if "error" in result:
             error_msg = f"❌ {result['error']}"
-            return result, error_msg, error_msg, error_msg
+            return result, error_msg, error_msg, error_msg, error_msg, error_msg
 
         models = result.get("models", {})
 
@@ -284,12 +361,14 @@ with gr.Blocks(title=title) as demo:
             extract_text("llama"),
             extract_text("pho"),
             extract_text("wingman"),
+            extract_text("perplexity"),
+            extract_text("gemini"),
         )
     
     reply_btn.click(
         generate_reply_with_extraction,
-        inputs=[reply_in, wingman_prompt_in],
-        outputs=[reply_out, llama_box, pho_box, wingman_box],
+        inputs=[reply_in, wingman_prompt_in, gemini_model_dropdown],
+        outputs=[reply_out, llama_box, pho_box, wingman_box, perplexity_box, gemini_box],
         api_name="reply"
     )
     

gemini_service.py

@@ -0,0 +1,239 @@
+"""
+Service for generating replies using Google Gemini API.
+"""
+import os
+import requests
+from typing import Optional, List, Dict, Any
+import google.generativeai as genai
+
+# Reuse the same system prompt from Perplexity service
+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 fetch_gemini_models(api_key: Optional[str] = None) -> List[Dict[str, Any]]:
+    """
+    Fetch available Gemini models from Google API.
+    
+    Args:
+        api_key: Google API key. If None, will try to get from GOOGLE_API_KEY env var.
+    
+    Returns:
+        List of model dictionaries with name, displayName, and description
+    """
+    api_key = api_key or os.getenv("GOOGLE_API_KEY")
+    if not api_key:
+        raise ValueError(
+            "Google API key is required.\n\n"
+            "Set environment variable:\n"
+            "  export GOOGLE_API_KEY=AIzaSy...\n\n"
+            "Or pass api_key parameter."
+        )
+    
+    url = f"https://generativelanguage.googleapis.com/v1beta/models?key={api_key}"
+    
+    try:
+        response = requests.get(url, timeout=10)
+        response.raise_for_status()
+        data = response.json()
+        
+        models = data.get("models", [])
+        
+        # Filter models that support generateContent
+        generative_models = [
+            {
+                "name": model.get("name", ""),
+                "displayName": model.get("displayName", ""),
+                "description": model.get("description", ""),
+                "version": model.get("version", ""),
+            }
+            for model in models
+            if "generateContent" in model.get("supportedGenerationMethods", [])
+        ]
+        
+        return generative_models
+        
+    except requests.exceptions.RequestException as e:
+        raise Exception(f"Failed to fetch Gemini models: {str(e)}")
+    except Exception as e:
+        raise Exception(f"Error parsing Gemini models: {str(e)}")
+
+
+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"):
+        """
+        Initialize Gemini service.
+        
+        Args:
+            api_key: Google API key. If None, will try to get from GOOGLE_API_KEY env var.
+            model_name: Model name to use (e.g., "gemini-2.5-flash", "models/gemini-2.5-pro")
+                       Can be with or without "models/" prefix
+        """
+        self.api_key = api_key or os.getenv("GOOGLE_API_KEY")
+        if not self.api_key:
+            raise ValueError(
+                "Google API key is required.\n\n"
+                "Set environment variable:\n"
+                "  export GOOGLE_API_KEY=AIzaSy...\n\n"
+                "Or pass api_key parameter when initializing GeminiReplyService."
+            )
+        
+        # Configure the API
+        genai.configure(api_key=self.api_key)
+        
+        # Normalize model name (remove "models/" prefix if present)
+        if model_name.startswith("models/"):
+            model_name = model_name.replace("models/", "")
+        
+        self.model_name = model_name
+        self.model = genai.GenerativeModel(model_name)
+    
+    def generate_reply(
+        self,
+        conversation: str,
+        trigger: str,
+        move: str,
+        temperature: float = 0.2,
+        max_output_tokens: int = 80,
+    ) -> str:
+        """
+        Generate reply using Google Gemini API.
+        
+        Args:
+            conversation: Conversation text in format "Male: ... ||| Female: ..."
+            trigger: Trigger label (e.g., "rapport_bid", "flirt_charm")
+            move: Move label (e.g., "charm", "invite", "validate")
+            temperature: Sampling temperature
+            max_output_tokens: Maximum tokens to generate
+        
+        Returns:
+            Generated reply text (1 sentence, ≤25 words)
+        """
+        user_content = f"""
+HỘI THOẠI: "{conversation}"
+TRIGGER: "{trigger}"
+MOVE: "{move}"
+""".strip()
+        
+        # Combine system prompt and user content
+        full_prompt = f"{SYSTEM_PROMPT}\n\n{user_content}"
+        
+        try:
+            # Generate content
+            response = self.model.generate_content(
+                full_prompt,
+                generation_config=genai.types.GenerationConfig(
+                    temperature=temperature,
+                    max_output_tokens=max_output_tokens,
+                ),
+            )
+            
+            raw = response.text.strip() if response.text else ""
+            
+            # Hậu xử lý: lấy câu đầu, giới hạn 25 từ
+            import re
+            # Tách theo dấu câu, lấy câu đầu
+            sentences = re.split(r'[.!?]', raw)
+            one_sentence = sentences[0].strip() if sentences else raw.strip()
+            
+            # Giới hạn 25 từ
+            words = one_sentence.split()
+            limited = " ".join(words[:25])
+            
+            # Đảm bảo kết thúc bằng dấu câu nếu cần
+            if limited and not limited[-1] in ".!?":
+                limited = limited.rstrip(",;:") + "."
+            
+            return limited
+            
+        except Exception as e:
+            raise Exception(f"Gemini API error: {str(e)}")
+
+
+# Global singleton instance
+_gemini_service = None
+_cached_models = None
+
+
+def get_gemini_service(
+    api_key: Optional[str] = None,
+    model_name: str = "gemini-2.5-flash",
+) -> GeminiReplyService:
+    """Get or create the global Gemini service instance."""
+    global _gemini_service
+    # Normalize model name (remove "models/" prefix if present)
+    normalized_name = model_name.replace("models/", "") if model_name.startswith("models/") else model_name
+    if _gemini_service is None or _gemini_service.model_name != normalized_name:
+        _gemini_service = GeminiReplyService(api_key=api_key, model_name=normalized_name)
+    return _gemini_service
+
+
+def get_available_gemini_models(api_key: Optional[str] = None, use_cache: bool = True) -> List[Dict[str, Any]]:
+    """
+    Get list of available Gemini models that support generateContent.
+    
+    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)
+    
+    Returns:
+        List of model dictionaries
+    """
+    global _cached_models
+    if use_cache and _cached_models is not None:
+        return _cached_models
+    
+    _cached_models = fetch_gemini_models(api_key=api_key)
+    return _cached_models
+

perplexity_service.py

@@ -0,0 +1,167 @@
+"""
+Service for generating replies using Perplexity API.
+"""
+import os
+from typing import Optional
+from openai import OpenAI
+
+
+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()
+
+
+class PerplexityReplyService:
+    """Service for generating replies using Perplexity API."""
+    
+    def __init__(self, api_key: Optional[str] = None, model: str = "mistral-7b-instruct"):
+        """
+        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")
+        """
+        self.api_key = api_key or os.getenv("PERPLEXITY_API_KEY")
+        if not self.api_key:
+            raise ValueError(
+                "Perplexity API key is required.\n\n"
+                "Set environment variable:\n"
+                "  export PERPLEXITY_API_KEY=pplx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx\n\n"
+                "Or pass api_key parameter when initializing PerplexityReplyService."
+            )
+        
+        self.client = OpenAI(
+            api_key=self.api_key,
+            base_url="https://api.perplexity.ai",
+        )
+        self.model = model
+    
+    def generate_reply(
+        self,
+        conversation: str,
+        trigger: str,
+        move: str,
+        temperature: float = 0.2,
+        max_tokens: int = 80,
+    ) -> str:
+        """
+        Generate reply using Perplexity API.
+        
+        Args:
+            conversation: Conversation text in format "Male: ... ||| Female: ..."
+            trigger: Trigger label (e.g., "rapport_bid", "flirt_charm")
+            move: Move label (e.g., "charm", "invite", "validate")
+            temperature: Sampling temperature
+            max_tokens: Maximum tokens to generate
+        
+        Returns:
+            Generated reply text (1 sentence, ≤25 words)
+        """
+        user_content = f"""
+HỘI THOẠI: "{conversation}"
+TRIGGER: "{trigger}"
+MOVE: "{move}"
+""".strip()
+        
+        try:
+            completion = self.client.chat.completions.create(
+                model=self.model,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                messages=[
+                    {
+                        "role": "system",
+                        "content": SYSTEM_PROMPT,
+                    },
+                    {
+                        "role": "user",
+                        "content": user_content,
+                    },
+                ],
+            )
+            
+            raw = completion.choices[0].message.content.strip() if completion.choices[0].message.content else ""
+            
+            # Hậu xử lý: lấy câu đầu, giới hạn 25 từ
+            import re
+            # Tách theo dấu câu, lấy câu đầu
+            sentences = re.split(r'[.!?]', raw)
+            one_sentence = sentences[0].strip() if sentences else raw.strip()
+            
+            # Giới hạn 25 từ
+            words = one_sentence.split()
+            limited = " ".join(words[:25])
+            
+            # Đảm bảo kết thúc bằng dấu câu nếu cần
+            if limited and not limited[-1] in ".!?":
+                limited = limited.rstrip(",;:") + "."
+            
+            return limited
+            
+        except Exception as e:
+            raise Exception(f"Perplexity API error: {str(e)}")
+
+
+# Global singleton instance
+_perplexity_service = None
+
+
+def get_perplexity_service(
+    api_key: Optional[str] = None,
+    model: str = "mistral-7b-instruct",
+) -> PerplexityReplyService:
+    """Get or create the global Perplexity service instance."""
+    global _perplexity_service
+    if _perplexity_service is None:
+        _perplexity_service = PerplexityReplyService(api_key=api_key, model=model)
+    return _perplexity_service
+

requirements.txt

@@ -16,4 +16,6 @@ bitsandbytes
 datasets
 pandas
 einops
+openai>=1.0.0
+google-generativeai>=0.3.0