Update main.py

main.py

@@ -1,31 +1,93 @@
+import os
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 from hugchat import hugchat
 from hugchat.login import Login
-import fastapi
-from fastapi import FastAPI, Query
+from dotenv import load_dotenv
+from contextlib import asynccontextmanager
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Load environment variables
+load_dotenv()
+
+# Initialize FastAPI app
 app = FastAPI()
-# Your HuggingFace credentials
+
+# HuggingFace credentials from environment variables
+# HuggingFace credentials from environment variables
 EMAIL = "[email protected]"
 PASSWORD = "Allahisgreatest17"
-@app.post("/ask")
-def main(prompt: str = Query(..., description="The prompt for the AI")):
-    # Login and save cookies (optional cookie path omitted for simplicity)
-    sign = Login(EMAIL, PASSWORD)
-    cookies = sign.login(save_cookies=True)
+ASSISTANT_ID = "682e0c1f5f0c3d952a27498e"
+
+# Cookie directory path
+COOKIE_PATH_DIR = "./cookies/"
+
+# Global chatbot instance
+chatbot = None
 
-    # Initialize chatbot with cookies
-    chatbot = hugchat.ChatBot(cookies=cookies.get_dict())
+# Pydantic model for request body
+class PromptRequest(BaseModel):
+    prompt: str
 
-    # Start a new conversation with your assistant ID
-    ASSISTANT_ID = "682e0c1f5f0c3d952a27498e"  # Replace with your assistant ID
-    chatbot.new_conversation(assistant=ASSISTANT_ID, switch_to=True)
+# Ensure cookie directory exists
+def ensure_cookie_directory():
+    if not os.path.exists(COOKIE_PATH_DIR):
+        os.makedirs(COOKIE_PATH_DIR)
+        logger.info("Created cookie directory: %s", COOKIE_PATH_DIR)
 
-    # Send message and get the full response (no streaming)
-    response = chatbot.chat(f"{prompt}", stream=False)
+# Lifespan handler for startup and shutdown
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Startup: Initialize chatbot
+    global chatbot
+    try:
+        # Validate credentials
+        if not EMAIL or not PASSWORD or not ASSISTANT_ID:
+            raise ValueError("HuggingFace credentials or assistant ID not provided in environment variables")
 
-    # Print full response
-    print("Assistant:", response)
-    return response
+        # Ensure cookie directory exists
+        ensure_cookie_directory()
 
+        # Log in to HuggingFace
+        sign = Login(EMAIL, PASSWORD)
+        cookies = sign.login(cookie_path_dir=COOKIE_PATH_DIR, save_cookies=True)
+        logger.info("Successfully logged in to HuggingFace")
+
+        # Initialize chatbot
+        chatbot = hugchat.ChatBot(cookies=cookies.get_dict())
+        # Start a new conversation
+        chatbot.new_conversation(assistant=ASSISTANT_ID, switch_to=True)
+        logger.info("Chatbot initialized with assistant ID: %s", ASSISTANT_ID)
+    except Exception as e:
+        logger.error("Failed to initialize chatbot: %s", str(e))
+        raise HTTPException(status_code=500, detail="Failed to initialize chatbot")
+
+    yield  # Application runs here
+
+    # Shutdown: Clean up (optional)
+    logger.info("Shutting down application")
+
+# Attach lifespan handler to FastAPI app
+app.lifespan = lifespan
+
+# Endpoint to handle chat requests
+@app.post("/ask")
+async def ask(prompt_request: PromptRequest):
+    try:
+        if not chatbot:
+            raise HTTPException(status_code=500, detail="Chatbot not initialized")
 
+        # Send prompt to chatbot
+        response = chatbot.chat(prompt_request.prompt, stream=False)
+        logger.info("Received response for prompt: %s", prompt_request.prompt)
 
+        # Return structured response
+        return {"status": "success", "response": response}
+    except Exception as e:
+        logger.error("Error processing request: %s", str(e))
+        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")