Update main.py

main.py

@@ -1,30 +1,70 @@
+from fastapi import FastAPI, Query
 from hugchat import hugchat
 from hugchat.login import Login
 import os
 
-# Credentials
-EMAIL = os.getenv("EMAIL")
-PASSWD = os.getenv("PASSWD")
+app = FastAPI()
 
-# Path to store cookies
-cookie_path_dir = "./cookies/"  # Make sure this directory exists
+# Global variable to store the chatbot instance
+chatbot = None
 
-# Log in and get cookies
-sign = Login(EMAIL, PASSWD)
-cookies = sign.login(cookie_dir_path=cookie_path_dir, save_cookies=True)
-
-# Create ChatBot with cookies
-chatbot = hugchat.ChatBot(cookies=cookies.get_dict())
-
-# 🔁 Create a new conversation with a custom assistant ID
-ASSISTANT_ID = "66017fca58d60bd7d5c5c26c"  # Replace with your actual assistant ID
-chatbot.new_conversation(assistant=ASSISTANT_ID, switch_to=True)
+def setup_chatbot(email, password, cookie_path, assistant_id):
+    """
+    Sets up the Hugging Face chatbot with login and conversation.
+    
+    Args:
+        email (str): User email for login
+        password (str): User password for login
+        cookie_path (str): Directory to store cookies
+        assistant_id (str): ID of the assistant to use
+    
+    Returns:
+        hugchat.ChatBot: Configured chatbot instance
+    """
+    try:
+        # Create cookie directory if it doesn't exist
+        os.makedirs(cookie_path, exist_ok=True)
+        
+        sign = Login(email, password)
+        cookies = sign.login(cookie_dir_path=cookie_path, save_cookies=True)
+        chatbot = hugchat.ChatBot(cookies=cookies.get_dict())
+        chatbot.new_conversation(assistant=assistant_id, switch_to=True)
+        return chatbot
+    except Exception as e:
+        raise Exception(f"Failed to set up chatbot: {e}")
 
-# 🔁 Stream response character-by-character as it generates
-print("Assistant:", end=" ", flush=True)
-for token in chatbot.chat("Hello, how can I help you today?", stream=True):
+# Initialize chatbot at startup
+@app.on_event("startup")
+async def startup_event():
+    global chatbot
+    # Credentials and configuration
+    EMAIL = os.getenv("EMAIL")
+    PASSWD = os.getenv("PASSWD")
+    COOKIE_PATH_DIR = "./cookies/"
+    ASSISTANT_ID = "682e0c1f5f0c3d952a27498e"  # Replace with your actual assistant ID
     
+    chatbot = setup_chatbot(EMAIL, PASSWD, COOKIE_PATH_DIR, ASSISTANT_ID)
 
-# Optionally: web search example
-# response = chatbot.chat("How many models stored in huggingface?", web_search=True)
-# print("\n\nWeb Search Result:", response.wait_until_done())
+@app.post("/generate")
+async def generate_response(
+    prompt: str = Query(..., description="The prompt for the AI")
+):
+    """
+    Generates a response from the AI based on the provided prompt.
+    
+    Args:
+        prompt (str): The user's input prompt for the AI.
+    
+    Returns:
+        dict: A dictionary containing the AI's response or an error message.
+    """
+    global chatbot
+    if chatbot is None:
+        return {"error": "Chatbot not initialized. Please try again later."}
+    
+    try:
+        # Generate response (non-streaming for simplicity in API context)
+        response = chatbot.chat(prompt, stream=False)
+        return {"response": response}
+    except Exception as e:
+        return {"error": f"Failed to generate response: {str(e)}"}