app.py

import os
import gradio as gr
import google.generativeai as genai
from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from datetime import datetime
import pytz
import time
import shutil

# Get API key from Hugging Face Spaces secrets
google_api_key = os.environ.get("GOOGLE_API_KEY")

if not google_api_key:
    raise ValueError("GOOGLE_API_KEY not found in environment variables. Please add it to Hugging Face Space secrets.")

# Configure Google Generative AI
genai.configure(api_key=google_api_key)

# Function to get current date and time
def get_current_datetime():
    # Using UTC as default, but you can change to any timezone
    utc_now = datetime.now(pytz.UTC)
    
    # Convert to IST (Indian Standard Time) - modify as needed
    ist_timezone = pytz.timezone('Asia/Kolkata')
    ist_now = utc_now.astimezone(ist_timezone)
    
    # Format the datetime
    formatted_date = ist_now.strftime("%B %d, %Y")
    formatted_time = ist_now.strftime("%I:%M:%S %p")
    
    return formatted_date, formatted_time

# Load PDF and create vector store
def initialize_retriever():
    try:
        # Get current directory
        current_dir = os.getcwd()
        print(f"Current working directory: {current_dir}")
        
        # List files in current directory for debugging
        print(f"Files in directory: {os.listdir(current_dir)}")
        
        # Use absolute path for the PDF
        pdf_path = os.path.join(current_dir, "Team1.pdf")
        print(f"Attempting to load PDF from: {pdf_path}")
        
        # Check if file exists
        if not os.path.exists(pdf_path):
            raise FileNotFoundError(f"The file {pdf_path} does not exist")
            
        # Load PDF
        loader = PyPDFLoader(pdf_path)
        documents = loader.load()
        
        print(f"Successfully loaded {len(documents)} pages from the PDF")

        # Split text into chunks
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
        text_chunks = text_splitter.split_documents(documents)
        print(f"Split into {len(text_chunks)} text chunks")

        # Generate embeddings
        embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")

        # Store embeddings in FAISS index
        vectorstore = FAISS.from_documents(text_chunks, embeddings)
        print("Successfully created vector store")
        return vectorstore.as_retriever(search_kwargs={"k": 10})
    
    except Exception as e:
        print(f"Error in initialize_retriever: {str(e)}")
        # Return a dummy retriever for graceful failure
        class DummyRetriever:
            def get_relevant_documents(self, query):
                return []
        
        print("Returning dummy retriever due to error")
        return DummyRetriever()

# Initialize LLM
def get_llm():
    try:
        return ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
    except Exception as e:
        print(f"Error initializing LLM: {str(e)}")
        return None

llm = get_llm()

# RAG query function
def rag_query(query, retriever):
    if retriever is None:
        return "Error: Could not initialize document retriever. Please check if Team1.pdf exists."
    
    # Get current date and time for context
    current_date, current_time = get_current_datetime()
    
    try:
        # Retrieve relevant documents
        docs = retriever.get_relevant_documents(query)
        
        if not docs:
            return "No relevant information found in the document. Try a general query instead."
        
        # Create context from retrieved documents
        context = "\n".join([doc.page_content for doc in docs])
        prompt = f"""Context:\n{context}
        Current Date: {current_date}
        Current Time: {current_time}
        Question: {query}
        Answer directly and concisely, using the current date and time information if relevant:"""

        response = llm.invoke(prompt)
        return response.content
    except Exception as e:
        return f"Error in RAG processing: {str(e)}"

# General query function
def general_query(query):
    if llm is None:
        return "Error: Could not initialize language model. Please check your API key."
    
    # Get current date and time for context
    current_date, current_time = get_current_datetime()
    
    try:
        # Define the prompt with date and time context
        prompt_template = """Current Date: {date}
        Current Time: {time}
        Answer the following query, using the current date and time information if relevant: {query}"""
            
        prompt = PromptTemplate.from_template(prompt_template)
        
        # Create an LLM Chain
        chain = LLMChain(llm=llm, prompt=prompt)
        
        # Run chatbot and get response
        response = chain.run(date=current_date, time=current_time, query=query)
        return response
    
    except Exception as e:
        return f"Error in general query: {str(e)}"

# Function to handle the case when no PDF is found
def file_not_found_message():
    return ("The Team1.pdf file could not be found. Team Query mode will not work properly. "
            "Please ensure the PDF is correctly uploaded to the Hugging Face Space.")

# Query router function
def query_router(query, method, retriever):
    if method == "Team Query":
        if isinstance(retriever, type) or retriever is None:
            return file_not_found_message()
        return rag_query(query, retriever)
    elif method == "General Query":
        return general_query(query)
    return "Invalid selection!"

# Function to reset input and output
def reset_query_field():
    return "", ""  # Reset only the query input

# Function to update the clock
def update_datetime():
    date, time = get_current_datetime()
    return date, time


# Main function to create and launch the Gradio interface
def main():
    # Initialize retriever
    print("Initializing retriever...")
    retriever = initialize_retriever()
    
    # Define custom CSS to style the logo and header
    # Define local image paths
    logo_path = "Equinix-LOGO.jpeg"  # Ensure this file exists

        
    # Custom CSS for background styling
    # Custom CSS for background styling
    custom_css = """
    .gradio-container {
        background-color: #f0f0f0;
        text-align: center;
    }
    #logo img {
        display: block;
        margin: 0 auto;
        max-width: 200px; /* Adjust size */
    }
    /* Hide download buttons and controls */
    .download-button {
        display: none !important;
    }
    /* Hide other download options */
    .file-preview .download {
        display: none !important;
    }
    /* Hide the three dots menu that might contain download options */
    .icon-button.secondary {
        display: none !important;
    }
    """

    
    
    # Create the Gradio interface using Blocks
    with gr.Blocks(css=custom_css) as demo:
        gr.Image(logo_path, elem_id="logo", show_label=False, height=100, width=400,show_download_button=False)  # Display Logo
        

        #gr.Markdown(image_html)
        # Title & Description
        gr.Markdown("<h1 style='text-align: center; color: black;'>Equinix Chatbot for Automation Team</h1>")
            
        # Date and Time Display
        with gr.Row(elem_classes="datetime-display"):
            date_display = gr.Textbox(label="Date", interactive=False)
            time_display = gr.Textbox(label="Time", interactive=False)
        
        # Add refresh button for time
        refresh_btn = gr.Button("Update Date & Time")
        refresh_btn.click(fn=update_datetime, inputs=[], outputs=[date_display, time_display])
        
        gr.Markdown("<p style='text-align: center; color: black;'>Ask me anything!</p>")

        # Input & Dropdown Section
        with gr.Row():
            query_input = gr.Textbox(label="Enter your query")
            query_method = gr.Dropdown(["Team Query", "General Query"], label="Select Query Type", value="Team Query")

        # Output Textbox
        output_box = gr.Textbox(label="Response", interactive=False)
        
        # Buttons Section
        with gr.Row():
            submit_button = gr.Button("Submit")
            reset_button = gr.Button("Reset Query")
        
        # Button Click Events
        submit_button.click(
            lambda query, method: query_router(query, method, retriever), 
            inputs=[query_input, query_method], 
            outputs=output_box
        )
            
        # Reset only the query input
        reset_button.click(reset_query_field, inputs=[], outputs=[query_input, output_box])
       
        # Update date and time on submission
        submit_button.click(
            fn=update_datetime,
            inputs=[],
            outputs=[date_display, time_display]
        )
        
        # Initialize date and time values
        date_val, time_val = get_current_datetime()
        date_display.value = date_val
        time_display.value = time_val

    # Launch the interface
    demo.launch(share=True)

if __name__ == "__main__":
    main()