Update main.py

main.py

@@ -2,53 +2,63 @@ import os
 import logging
 import asyncio
 import nest_asyncio
-
-# Import the Hugging Face Inference Client and login function
-from huggingface_hub import InferenceClient, login
-import langid
+import requests
+from hugchat import hugchat
 
 # Import Telegram bot components
 from telegram import Update, Bot
 from telegram.ext import Application, CommandHandler, MessageHandler, filters, CallbackContext
 
+# -------------------------
+# Load environment variables
+# -------------------------
+# Expected environment variables:
+# TELEGRAM_BOT_TOKEN - Telegram bot token
+# WEBHOOK_DOMAIN - example: your-space.hf.space
+# HF_EMAIL and HF_PASSWORD 
+
+TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
+if not TELEGRAM_BOT_TOKEN:
+    raise ValueError("Missing TELEGRAM_BOT_TOKEN environment variable.")
+
+WEBHOOK_DOMAIN = os.getenv("WEBHOOK_DOMAIN")
+if not WEBHOOK_DOMAIN:
+    raise ValueError("Missing WEBHOOK_DOMAIN environment variable.")
+
 # -------------------------
 # Configure logging
 # -------------------------
-logging.basicConfig(
-    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
-)
+logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # -------------------------
-# Environment variables and login to Hugging Face
+# Set up Hugging Chat integration
 # -------------------------
-HF_HUB_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
-if not HF_HUB_TOKEN:
-    raise ValueError("Missing Hugging Face API token. Please set HUGGINGFACEHUB_API_TOKEN.")
+# Logs into Hugging Face or uses a fallback to get cookies.
+if os.getenv("HF_EMAIL") and os.getenv("HF_PASSWORD"):
+    from hugchat.login import Login
+    logger.info("Logging into Hugging Face using provided credentials.")
+    sign = Login(os.getenv("HF_EMAIL"), os.getenv("HF_PASSWORD"))
+    cookies = sign.login()
+else:
+    logger.info("No HF_EMAIL/HF_PASSWORD provided. Using fallback method to get cookies.")
+    cookies = requests.get("https://huggingface.co/chat/").cookies
 
-login(token=HF_HUB_TOKEN)
-client = InferenceClient(api_key=HF_HUB_TOKEN)
-
-# Get Telegram Bot Token
-TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
-if not TOKEN:
-    raise ValueError("Missing Telegram Bot Token. Please set TELEGRAM_BOT_TOKEN.")
-
-# Get Webhook Domain (e.g., your-space.hf.space)
-WEBHOOK_DOMAIN = os.getenv("WEBHOOK_DOMAIN")
-if not WEBHOOK_DOMAIN:
-    raise ValueError("Missing Webhook Domain. Please set WEBHOOK_DOMAIN.")
-
-# Construct the full webhook URL.
-# We use the Telegram bot token as a unique path.
-WEBHOOK_URL = f"https://{WEBHOOK_DOMAIN}/{TOKEN}"
 
+# Create a ChatBot instance from the hugchat library.
+chatbot = hugchat.ChatBot(cookies=cookies.get_dict())
 
 # -------------------------
-# Function to detect language
+# Helper Functions
 # -------------------------
-def detect_language(user_input: str):
+
+def detect_language(user_input: str) -> str:
+    """
+    Detect language using langid.
+    Returns 'hebrew' if Hebrew, 'english' if English, otherwise 'unsupported'.
+    """
     try:
+        import langid
         lang, _ = langid.classify(user_input)
         if lang == "he":
             return "hebrew"
@@ -60,94 +70,103 @@ def detect_language(user_input: str):
         logger.error(f"Language detection error: {e}")
         return "unsupported"
 
-
-# -------------------------
-# Function to generate response using Hugging Face Chat models
-# -------------------------
-def generate_response(text: str) -> str:
-    language = detect_language(text)
-    
+        
+def generate_response_sync(message: str) -> str:
+    """
+    Generate a response using Hugging Chat in a synchronous manner.
+    This function is blocking; we'll run it in an executor.
+    """
+    language = detect_language(message)
     if language == "hebrew":
-        # For Hebrew, prompt the model accordingly.
-        content = "תענה בקצרה אבל תשתף את תהליך קבלת ההחלטות שלך, " + text
-        model = "microsoft/Phi-3.5-mini-instruct"
+        prompt = "תענה בקצרה אבל תשתף את תהליך קבלת ההחלטות שלך, " + message
     elif language == "english":
-        # For English, use a different prompt.
-        content = "keep it short but tell your decision making process, " + text
-        model = "mistralai/Mistral-Nemo-Instruct-2407"
+        prompt = "keep it short but tell your decision making process, " + message
     else:
         return "Sorry, I only support Hebrew and English."
     
-    messages = [{"role": "user", "content": content}]
-    
+    response_queue = ""
+    full_response = ""
     try:
-        # Call the chat completion API
-        completion = client.chat.completions.create(
-            model=model,
-            messages=messages,
-            max_tokens=2048,
-            temperature=0.5,
-            top_p=0.7
-        )
-        return completion.choices[0].message.content
+        # The huggingchat.ChatBot.chat() method returns a generator (streaming tokens)
+        for resp in chatbot.chat(prompt, _stream_yield_all=True):
+            if resp and "token" in resp:
+                response_queue += resp["token"]
+            # Flush the response if it's long (optional, here we accumulate)
+        full_response = response_queue
+        return full_response
     except Exception as e:
-        logger.error(f"Error generating response: {e}")
+        logger.error(f"Error generating response via Hugging Chat: {e}")
         return "Error: Could not generate response."
 
 
+async def generate_response(message: str) -> str:
+    """
+    Asynchronous wrapper around generate_response_sync using run_in_executor.
+    """
+    loop = asyncio.get_event_loop()
+    response = await loop.run_in_executor(None, generate_response_sync, message)
+    return response
+
+
 # -------------------------
 # Telegram Bot Handlers
 # -------------------------
 async def start(update: Update, context: CallbackContext):
-    # Respond to the /start command.
+    """
+    Handler for the /start command.
+    """
     await update.message.reply_text("Hello! Tell me your decision-making issue, and I'll try to help.")
-    logger.info("Start command received.")
+    logger.info("Received /start command.")
 
 
 async def handle_message(update: Update, context: CallbackContext):
+    """
+    Handler for incoming text messages.
+    It generates a response using the Hugging Chat integration.
+    """
     user_text = update.message.text
-    logger.info(f"User message: {user_text}")
-    
-    # Generate a response using our Hugging Face chat logic.
-    response_text = generate_response(user_text)
+    logger.info(f"Received message: {user_text}")
+    response_text = await generate_response(user_text)
     logger.info(f"Generated response: {response_text}")
-    
     await update.message.reply_text(response_text)
 
 
+# -------------------------
+# Webhook Configuration for Telegram
+# -------------------------
+# Construct the full webhook URL.
+# We use the bot token as the URL path for uniqueness.
+WEBHOOK_URL = f"https://{WEBHOOK_DOMAIN}/{TELEGRAM_BOT_TOKEN}"
+
 # -------------------------
 # Main function to run the Telegram Bot using Webhook mode
 # -------------------------
 async def main():
-    # Build the Application with the Telegram Bot Token.
-    application = Application.builder().token(TOKEN).build()
+    # Build the Telegram Application.
+    application = Application.builder().token(TELEGRAM_BOT_TOKEN).build()
 
-    # Add command and message handlers.
+    # Add handlers.
     application.add_handler(CommandHandler("start", start))
     application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))
 
     logger.info("Starting bot in webhook mode...")
     print("Starting bot in webhook mode...")
 
-    # Run the bot in webhook mode.
-    # The bot will listen on all interfaces (0.0.0.0) at port 7860.
-    # Telegram will send updates to WEBHOOK_URL.
+    # Run the bot using webhook mode.
     await application.run_webhook(
-        listen="0.0.0.0",         # Listen on all interfaces.
-        port=7860,                # Port to listen on.
-        url_path=TOKEN,           # Use the bot token as the URL path.
-        webhook_url=WEBHOOK_URL   # Full webhook URL.
+        listen="0.0.0.0",           # Listen on all interfaces.
+        port=7860,                  # Port to listen on – must match the exposed port in Dockerfile.
+        url_path=TELEGRAM_BOT_TOKEN,  # Use the bot token as the URL path.
+        webhook_url=WEBHOOK_URL     # Full webhook URL that Telegram will call.
     )
 
-
 # -------------------------
-# Run the main function without closing a running event loop
+# Run the main function (handling event loop issues)
 # -------------------------
 if __name__ == "__main__":
-    # Apply nest_asyncio to allow nested event loops (needed in some HF Space environments).
+    # Apply nest_asyncio to support nested event loops (useful in HF Spaces).
     nest_asyncio.apply()
     loop = asyncio.get_event_loop()
-    
     try:
         loop.run_until_complete(main())
     except Exception as e: