main.py

import os
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from hugchat import hugchat
from hugchat.login import Login
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()

# HuggingFace credentials from environment variables
# HuggingFace credentials from environment variables
EMAIL = "[email protected]"
PASSWORD = "Allahisgreatest17"
ASSISTANT_ID = "682e0c1f5f0c3d952a27498e"

# Cookie directory path
COOKIE_PATH_DIR = "./cookies/"

# Global chatbot instance
chatbot = None

# Pydantic model for request body
class PromptRequest(BaseModel):
    prompt: str

# 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)

# 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")

        # 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)}")