CyberSecurityChatBot/ProjectCyberAssistant

1import os2import gradio as gr3import time4from langchain.chains import RetrievalQA5from langchain_community.vectorstores import Chroma6from langchain_community.document_loaders import PyPDFLoader7from langchain_huggingface import HuggingFaceEmbeddings  # ✅ Fixed Import8from huggingface_hub import InferenceClient  # ✅ New method for querying Hugging Face LLM9 10# Install required dependencies (ensure latest versions)11os.system("pip install -U huggingface_hub langchain_huggingface langchain_core langchain gradio")12 13# Define paths for cybersecurity training PDFs14PDF_FILES = [15    "ISOIEC 27001_2ef522.pdf",16    "ISO-IEC-27005-2022.pdf",17    "MITRE ATLAS Overview Combined_v1.pdf",18    "NIST_CSWP_04162018.pdf"19]20 21# Fetch Hugging Face API token securely from environment variables22HUGGINGFACE_API_KEY = os.getenv("HUGGINGFACEHUB_API_TOKEN")23if HUGGINGFACE_API_KEY is None:24    raise ValueError("❌ Hugging Face API token is missing! Set it in Hugging Face Spaces Secrets.")25 26# Load PDFs into ChromaDB27def load_data():28    """Loads multiple PDFs and stores embeddings in ChromaDB"""29    all_docs = []30    for pdf in PDF_FILES:31        if os.path.exists(pdf):  # Ensure the PDF exists32            loader = PyPDFLoader(pdf)33            all_docs.extend(loader.load())34 35    # Use updated embedding model36    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")37 38    return Chroma.from_documents(all_docs, embeddings)39 40# Load the knowledge base41vector_db = load_data()42 43# Initialize Hugging Face Inference Client (new recommended method)44client = InferenceClient(45    "https://api-inference.huggingface.co/models/google/flan-t5-large",46    token=HUGGINGFACE_API_KEY47)48 49# Function to interact with the Hugging Face model50def query_llm(prompt):51    """Send query to Hugging Face API and return response"""52    response = client.chat_completion(messages=[{"role": "user", "content": prompt}])53    return response["choices"][0]["message"]["content"]54 55# Create Retrieval QA chain56qa_chain = RetrievalQA.from_chain_type(llm=query_llm, retriever=vector_db.as_retriever())57 58# Function to simulate futuristic typing effect59def chatbot_response(question):60    """Handles chatbot queries with a typing effect"""61    response = qa_chain.invoke(question)  # ✅ Use `invoke` instead of deprecated `run`62    displayed_response = ""63    for char in response:64        displayed_response += char65        time.sleep(0.02)  # Simulate typing delay66        yield displayed_response67 68# Custom futuristic CSS styling69custom_css = """70body {background-color: #0f172a; color: #0ff; font-family: 'Orbitron', sans-serif;}71.gradio-container {background: linear-gradient(to bottom, #020c1b, #001f3f);}72textarea {background: #011627; color: #0ff; font-size: 18px;}73button {background: #0088ff; color: white; font-size: 20px; border-radius: 5px; border: none; padding: 10px;}74button:hover {background: #00ffff; color: #000;}75"""76 77# Create Gradio Chatbot Interface78iface = gr.Interface(79    fn=chatbot_response,80    inputs="text",81    outputs="text",82    title="🤖 Cybersecurity AI Assistant",83    description="Ask me about NIST, ISO/IEC 27001, MITRE ATLAS, and ISO/IEC 27005. Powered by AI.",84    theme="default",85    css=custom_css,86    live=True,  # Enables real-time updates for typing effect87)88 89# Launch chatbot with public link90iface.launch(share=True)  # ✅ Now launches with a public link91