import gradio as gr
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np

# Knowledge base
documents = [
    "Leave policy: Employees are entitled to 20 annual leaves per year.",
    "Health insurance: The company provides full coverage for employees and partial coverage for dependents.",
    "Working hours: Standard office hours are 9 AM to 6 PM, Monday to Friday.",
    "Remote work: Employees may work from home 2 days per week with manager approval.",
    "Promotions: Promotions are based on yearly performance reviews."
]

# Load embedding model
embedder = SentenceTransformer("all-MiniLM-L6-v2")
doc_embeddings = embedder.encode(documents, convert_to_numpy=True)

# Create FAISS index
dimension = doc_embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(doc_embeddings)

# RAG function
def rag_chat(message, history):
    query_embedding = embedder.encode([message], convert_to_numpy=True)
    D, I = index.search(query_embedding, k=1)
    best_doc = documents[I[0][0]]
    response = f"📘 Policy Reference: {best_doc}"
    return response

# Gradio UI
demo = gr.ChatInterface(
    fn=rag_chat,
    title="💡 WellCare AI - HR Assistant",
    description="Ask me anything about leave, health insurance, or company policy.",
    theme="soft"
)

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