src/chatbot.py

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import streamlit as st
import requests
import json
from huggingface_hub import InferenceClient

class PersianChatbot:
    def __init__(self):
        """
        کلاس چت‌بات فارسی با استفاده از مدل‌های Hugging Face
        """
        self.model_name = "alibidaran/Gemma2_Farsi"
        self.tokenizer = None
        self.model = None
        self.inference_client = None
        self.is_loaded = False
        
    def load_model_local(self):
        """
        بارگذاری مدل به صورت محلی (برای استفاده در محیط‌های قدرتمند)
        """
        try:
            st.info("در حال بارگذاری مدل چت‌بات...")
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                device_map="auto" if torch.cuda.is_available() else None
            )
            self.is_loaded = True
            st.success("مدل چت‌بات با موفقیت بارگذاری شد!")
            return True
        except Exception as e:
            st.error(f"خطا در بارگذاری مدل: {str(e)}")
            return False
    
    def load_model_inference(self):
        """
        استفاده از Inference API برای مدل‌های Hugging Face (مناسب برای Spaces)
        """
        try:
            self.inference_client = InferenceClient(model=self.model_name)
            self.is_loaded = True
            st.success("اتصال به مدل چت‌بات برقرار شد!")
            return True
        except Exception as e:
            st.error(f"خطا در اتصال به مدل: {str(e)}")
            return False
    
    def generate_response_local(self, user_input, max_length=512):
        """
        تولید پاسخ با استفاده از مدل محلی
        """
        if not self.is_loaded or self.model is None:
            return "مدل چت‌بات بارگذاری نشده است."
        
        try:
            # آماده‌سازی ورودی
            prompt = f"کاربر: {user_input}\nدستیار: "
            inputs = self.tokenizer.encode(prompt, return_tensors="pt")
            
            # تولید پاسخ
            with torch.no_grad():
                outputs = self.model.generate(
                    inputs,
                    max_length=max_length,
                    num_return_sequences=1,
                    temperature=0.7,
                    do_sample=True,
                    pad_token_id=self.tokenizer.eos_token_id
                )
            
            # استخراج پاسخ
            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            response = response.split("دستیار: ")[-1].strip()
            
            return response
        
        except Exception as e:
            return f"خطا در تولید پاسخ: {str(e)}"
    
    def generate_response_inference(self, user_input, max_tokens=512):
        """
        تولید پاسخ با استفاده از Inference API
        """
        if not self.is_loaded or self.inference_client is None:
            return "اتصال به مدل چت‌بات برقرار نیست."
        
        try:
            # آماده‌سازی پرامپت
            prompt = f"سوال: {user_input}\nپاسخ:"
            
            # ارسال درخواست به API
            response = self.inference_client.text_generation(
                prompt,
                max_new_tokens=max_tokens,
                temperature=0.7,
                do_sample=True,
                return_full_text=False
            )
            
            return response.strip()
        
        except Exception as e:
            return f"خطا در دریافت پاسخ: {str(e)}"
    
    def generate_response_fallback(self, user_input):
        """
        پاسخ‌دهی پیش‌فرض در صورت عدم دسترسی به مدل
        """
        # تحلیل کلمات کلیدی و ارائه پاسخ‌های از پیش تعریف شده
        user_input_lower = user_input.lower()
        
        # پاسخ‌های مرتبط با اتوماسیون وب
        if any(word in user_input_lower for word in ['وب‌سایت', 'سایت', 'مرورگر', 'کلیک', 'فرم']):
            return """برای اتوماسیون وب می‌توانید از دستورات زیر استفاده کنید:
            
🌐 **دستورات مرورگر:**
- "وب‌سایت [آدرس] را باز کن"
- "روی دکمه [نام دکمه] کلیک کن"
- "فرم را با اطلاعات [اطلاعات] پر کن"
- "عکس از صفحه بگیر"
- "صفحه را اسکرول کن"

💡 **نکته:** لطفاً دستورات خود را واضح و مشخص بیان کنید تا بتوانم بهترین کمک را به شما ارائه دهم."""

        # پاسخ‌های مرتبط با TradingView
        elif any(word in user_input_lower for word in ['تریدینگ', 'trading', 'اندیکاتور', 'سیگنال', 'بورس']):
            return """برای نظارت بر TradingView می‌توانید از قابلیت‌های زیر استفاده کنید:

📊 **قابلیت‌های TradingView:**
- نظارت بر 5 اندیکاتور همزمان (RSI, MACD, EMA, Bollinger, Stochastic)
- ارسال اعلان هنگام همگرایی سیگنال‌ها
- پشتیبانی از ایمیل و وب‌هوک
- نظارت 24 ساعته

🔧 **تنظیمات:**
- به بخش "TradingView" بروید
- نماد مورد نظر را انتخاب کنید
- اندیکاتورهای مورد نظر را فعال کنید
- نظارت را شروع کنید"""

        # پاسخ‌های مرتبط با تنظیمات
        elif any(word in user_input_lower for word in ['تنظیمات', 'ایمیل', 'وب‌هوک', 'اعلان']):
            return """برای تنظیم اعلان‌ها مراحل زیر را دنبال کنید:

📧 **تنظیمات ایمیل:**
- سرور SMTP (مثل smtp.gmail.com)
- پورت (معمولاً 587)
- نام کاربری و رمز عبور
- ایمیل گیرنده

🔗 **تنظیمات وب‌هوک:**
- آدرس وب‌هوک
- روش ارسال (POST/GET)
- محتوای JSON

🧪 **تست:**
- از دکمه‌های تست استفاده کنید
- اتصال خود را بررسی کنید"""

        # پاسخ‌های عمومی
        else:
            responses = [
                f"متشکرم از سوال شما درباره '{user_input}'. چطور می‌توانم بهتر کمکتان کنم؟",
                f"سوال جالبی درباره '{user_input}' پرسیدید. لطفاً بیشتر توضیح دهید.",
                f"برای پاسخ بهتر به سوال '{user_input}' نیاز به اطلاعات بیشتری دارم.",
                "من اینجا هستم تا در زمینه اتوماسیون وب، TradingView و تنظیمات سیستم کمکتان کنم. چه کاری می‌توانم انجام دهم؟"
            ]
            import random
            return random.choice(responses)
    
    def chat(self, user_input):
        """
        تابع اصلی چت که بهترین روش موجود را انتخاب می‌کند
        """
        if not user_input.strip():
            return "لطفاً پیام خود را وارد کنید."
        
        # تلاش برای استفاده از مدل اصلی
        if self.is_loaded:
            if self.model is not None:
                return self.generate_response_local(user_input)
            elif self.inference_client is not None:
                return self.generate_response_inference(user_input)
        
        # در صورت عدم دسترسی، استفاده از پاسخ‌های پیش‌فرض
        return self.generate_response_fallback(user_input)
    
    def initialize(self, use_inference=True):
        """
        راه‌اندازی چت‌بات
        """
        if use_inference:
            return self.load_model_inference()
        else:
            return self.load_model_local()

# کلاس مدیریت گفتگو
class ConversationManager:
    def __init__(self):
        self.conversation_history = []
        self.max_history = 10
    
    def add_message(self, role, content):
        """
        افزودن پیام به تاریخچه گفتگو
        """
        self.conversation_history.append({
            "role": role,
            "content": content,
            "timestamp": str(pd.Timestamp.now())
        })
        
        # محدود کردن تاریخچه
        if len(self.conversation_history) > self.max_history * 2:
            self.conversation_history = self.conversation_history[-self.max_history * 2:]
    
    def get_context(self):
        """
        دریافت متن گفتگو برای ارسال به مدل
        """
        context = ""
        for msg in self.conversation_history[-6:]:  # آخرین 6 پیام
            if msg["role"] == "user":
                context += f"کاربر: {msg['content']}\n"
            else:
                context += f"دستیار: {msg['content']}\n"
        return context
    
    def clear_history(self):
        """
        پاک کردن تاریخچه گفتگو
        """
        self.conversation_history = []

# تابع کمکی برای راه‌اندازی چت‌بات در Streamlit
@st.cache_resource
def get_chatbot():
    """
    دریافت نمونه چت‌بات (با کش برای بهینه‌سازی)
    """
    chatbot = PersianChatbot()
    # تلاش برای راه‌اندازی با Inference API
    chatbot.initialize(use_inference=True)
    return chatbot

# تابع کمکی برای مدیریت گفتگو
def get_conversation_manager():
    """
    دریافت مدیر گفتگو
    """
    if "conversation_manager" not in st.session_state:
        st.session_state.conversation_manager = ConversationManager()
    return st.session_state.conversation_manager