app.py

import os
import time
import logging
import sys

import gradio as gr
from pinecone import Pinecone, ServerlessSpec

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext, Settings
from llama_index.vector_stores.pinecone import PineconeVectorStore
from llama_index.readers.file import PDFReader
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding


# -----------------------------
# Logging
# -----------------------------
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logger = logging.getLogger(__name__)


# -----------------------------
# Environment Variables
# -----------------------------
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")

PINECONE_INDEX_NAME = os.getenv("PINECONE_INDEX_NAME", "dds-hr-chatbot")
PINECONE_CLOUD = os.getenv("PINECONE_CLOUD", "aws")
PINECONE_REGION = os.getenv("PINECONE_REGION", "us-east-1")

REINDEX_ON_STARTUP = os.getenv("REINDEX_ON_STARTUP", "false").lower() == "true"

DATA_DIR = "data"

if not OPENAI_API_KEY:
    raise ValueError("OPENAI_API_KEY is missing. Please add it in Hugging Face Spaces secrets.")

if not PINECONE_API_KEY:
    raise ValueError("PINECONE_API_KEY is missing. Please add it in Hugging Face Spaces secrets.")


# -----------------------------
# LlamaIndex Settings
# -----------------------------
Settings.llm = OpenAI(
    model="gpt-4o-mini",
    temperature=0.2,
    api_key=OPENAI_API_KEY
)

Settings.embed_model = OpenAIEmbedding(
    model="text-embedding-ada-002",
    api_key=OPENAI_API_KEY
)

Settings.chunk_size = 600
Settings.chunk_overlap = 200


# -----------------------------
# System Prompt
# -----------------------------
system_prompt = """
You are AYesha, the Decoding Data Science (DDS) Enterprise HR Chatbot.

Answer questions exclusively using the attached DDS HR Handbook. Base all responses only on the information available in the handbook. Only respond to queries directly related to DDS HR policies as outlined in the handbook.

Rules:
- If a question is outside DDS HR policies, politely clarify that you are a human resources bot and only answer DDS HR questions.
- If a question cannot be answered from the handbook, politely decline and direct the user to email connect@decodingdatascience.com.
- Never answer questions about anything outside your HR handbook scope.
- Do not provide salary details, confidential information, old policies, legal advice, or company-wide information outside HR policies.
- Do not reveal internal reasoning.
- Always answer in a concise and professional tone.

For forbidden or unsupported topics, say:
“I’m sorry, I can only answer questions about the latest DDS HR policies. For confidential or other queries, please email connect@decodingdatascience.com.”

Remember: You are AYesha, the DDS HR Enterprise Chatbot. You must only answer from the authorized HR handbook content.
"""


# -----------------------------
# Pinecone Setup
# -----------------------------
def setup_pinecone_index():
    pc = Pinecone(api_key=PINECONE_API_KEY)

    existing_indexes = [index_info["name"] for index_info in pc.list_indexes()]

    if PINECONE_INDEX_NAME not in existing_indexes:
        logger.info(f"Creating Pinecone index: {PINECONE_INDEX_NAME}")

        pc.create_index(
            name=PINECONE_INDEX_NAME,
            dimension=1536,
            metric="cosine",
            spec=ServerlessSpec(
                cloud=PINECONE_CLOUD,
                region=PINECONE_REGION
            )
        )

        while not pc.describe_index(PINECONE_INDEX_NAME).status["ready"]:
            logger.info("Waiting for Pinecone index to be ready...")
            time.sleep(2)
    else:
        logger.info(f"Using existing Pinecone index: {PINECONE_INDEX_NAME}")

    return pc.Index(PINECONE_INDEX_NAME)


# -----------------------------
# Load or Create Index
# -----------------------------
def build_query_engine():
    pinecone_index = setup_pinecone_index()

    vector_store = PineconeVectorStore(
        pinecone_index=pinecone_index
    )

    storage_context = StorageContext.from_defaults(
        vector_store=vector_store
    )

    index_stats = pinecone_index.describe_index_stats()
    total_vectors = index_stats.get("total_vector_count", 0)

    if total_vectors == 0 or REINDEX_ON_STARTUP:
        logger.info("Loading documents and creating vector index...")

        if not os.path.exists(DATA_DIR):
            raise ValueError(
                "The 'data' folder is missing. Please create a data folder and upload your PDF file inside it."
            )

        documents = SimpleDirectoryReader(
            input_dir=DATA_DIR,
            required_exts=[".pdf"],
            file_extractor={".pdf": PDFReader()}
        ).load_data()

        if not documents:
            raise ValueError("No PDF documents were loaded from the 'data' folder.")

        index = VectorStoreIndex.from_documents(
            documents,
            storage_context=storage_context
        )

        logger.info("Documents indexed successfully.")
    else:
        logger.info("Existing Pinecone vectors found. Loading index from vector store.")

        index = VectorStoreIndex.from_vector_store(
            vector_store=vector_store
        )

    query_engine = index.as_query_engine(
        similarity_top_k=5,
        system_prompt=system_prompt
    )

    return query_engine


query_engine = build_query_engine()


# -----------------------------
# Chat Function
# -----------------------------
def query_doc(message, history):
    if not message or not message.strip():
        return "Please enter a question about the DDS HR handbook."

    try:
        response = query_engine.query(message)
        return str(response)

    except Exception as e:
        logger.error(f"Error while answering query: {e}")
        return "Sorry, something went wrong while processing your question. Please try again."


# -----------------------------
# Example Questions
# -----------------------------
example_questions = [
    "What is the leave policy?",
    "What is the work from home policy?",
    "What is the probation policy?",
    "What are the employee code of conduct rules?",
    "Who should I contact for confidential HR questions?"
]


# -----------------------------
# Gradio UI
# -----------------------------
with gr.Blocks(title="DDS Enterprise HR Chatbot") as demo:
    gr.Markdown(
        """
        # DDS Enterprise HR Chatbot

        Ask questions based on the DDS HR Handbook.

        This chatbot uses LlamaIndex, Pinecone, OpenAI, and Gradio.
        """
    )

    gr.ChatInterface(
        fn=query_doc,
        examples=example_questions,
        textbox=gr.Textbox(
            placeholder="Ask a question about DDS HR policies...",
            label="Your Question"
        )
    )

    gr.Markdown(
        """
        ---
        **Note:** This chatbot only answers questions related to the DDS HR Handbook.
        For confidential or unsupported questions, please contact connect@decodingdatascience.com.
        """
    )


if __name__ == "__main__":
    demo.launch()