Update app.py

app.py

@@ -1,26 +1,13 @@
 import os
 import gradio as gr
 from langchain.chains import RetrievalQA
-from langchain_community.vectorstores import Chroma  # Use langchain_community for vector stores
-from langchain.llms import OpenAI, HuggingFaceHub
-from langchain.embeddings import OpenAIEmbeddings, HuggingFaceEmbeddings
-from langchain.document_loaders import PyPDFLoader
+from langchain_community.vectorstores import Chroma  # ✅ Fixed Import
+from langchain.llms import OpenAI
+from langchain_huggingface import HuggingFaceEndpoint  # ✅ Corrected Import
+from langchain.embeddings import OpenAIEmbeddings
+from langchain_community.embeddings import HuggingFaceEmbeddings  # ✅ Corrected Import
+from langchain_community.document_loaders import PyPDFLoader  # ✅ Corrected Import
 import time
-from langchain_huggingface import HuggingFaceEndpoint
-
-
-# Fetch API token securely from environment variables
-HUGGINGFACE_API_KEY = os.getenv("HUGGINGFACEHUB_API_TOKEN")
-
-if HUGGINGFACE_API_KEY is None:
-    raise ValueError("Hugging Face API token is missing! Set it as a secret in Hugging Face Spaces.")
-
-# Load LLM securely
-llm = HuggingFaceEndpoint(
-    repo_id="google/flan-t5-large",
-    model_kwargs={"temperature": 0.5, "max_length": 512},
-    huggingfacehub_api_token=HUGGINGFACE_API_KEY  # ✅ Secure API key usage
-)
 
 # Define paths for cybersecurity training PDFs
 PDF_FILES = [
@@ -30,9 +17,12 @@ PDF_FILES = [
     "NIST_CSWP_04162018.pdf"
 ]
 
-# Choose whether to use OpenAI API (Online) or Hugging Face (Offline)
-USE_OPENAI = False  # Set to True if using OpenAI API for better responses
+# Fetch Hugging Face API token securely from environment variables
+HUGGINGFACE_API_KEY = os.getenv("HUGGINGFACEHUB_API_TOKEN")
+if HUGGINGFACE_API_KEY is None:
+    raise ValueError("❌ Hugging Face API token is missing! Set it in Hugging Face Spaces Secrets.")
 
+# Load PDFs into ChromaDB
 def load_data():
     """Loads multiple PDFs and stores embeddings in ChromaDB"""
     all_docs = []
@@ -41,27 +31,25 @@ def load_data():
             loader = PyPDFLoader(pdf)
             all_docs.extend(loader.load())
 
-    # Use OpenAI embeddings (Online) or Hugging Face embeddings (Offline)
-    if USE_OPENAI:
-        embeddings = OpenAIEmbeddings()
-    else:
-        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")  # ✅ Use updated embeddings
 
     return Chroma.from_documents(all_docs, embeddings)
 
-# Load the knowledge base from the uploaded PDFs
+# Load the knowledge base
 vector_db = load_data()
 
-# Select the LLM model (Online: OpenAI | Offline: Hugging Face)
-if USE_OPENAI:
-    llm = OpenAI()
-else:
-    llm = HuggingFaceHub(repo_id="google/flan-t5-large", model_kwargs={"temperature": 0.5, "max_length": 512})
+# Load LLM from Hugging Face securely
+llm = HuggingFaceEndpoint(
+    repo_id="google/flan-t5-large",
+    temperature=0.5,  # ✅ Explicitly set temperature
+    max_length=512,  # ✅ Explicitly set max_length
+    huggingfacehub_api_token=HUGGINGFACE_API_KEY
+)
 
-# Create Retrieval QA chain for document-based responses
+# Create Retrieval QA chain
 qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=vector_db.as_retriever())
 
-# Function to generate chatbot responses with a typing effect
+# Function to simulate futuristic typing effect
 def chatbot_response(question):
     """Handles chatbot queries with a typing effect"""
     response = qa_chain.run(question)
@@ -74,59 +62,23 @@ def chatbot_response(question):
 # Custom futuristic CSS styling
 custom_css = """
 body {background-color: #0f172a; color: #0ff; font-family: 'Orbitron', sans-serif;}
-#chatbot-container {border: 2px solid #00ffff; background: rgba(0, 0, 0, 0.8); padding: 20px; border-radius: 15px;}
 .gradio-container {background: linear-gradient(to bottom, #020c1b, #001f3f);}
 textarea {background: #011627; color: #0ff; font-size: 18px;}
 button {background: #0088ff; color: white; font-size: 20px; border-radius: 5px; border: none; padding: 10px;}
 button:hover {background: #00ffff; color: #000;}
 """
 
-# 3D Avatar using Three.js
-three_js_html = """
-<div id="avatar-container">
-  <script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script>
-  <script>
-    function create3DAvatar() {
-        var scene = new THREE.Scene();
-        var camera = new THREE.PerspectiveCamera(75, 1, 0.1, 1000);
-        var renderer = new THREE.WebGLRenderer({ alpha: true });
-        renderer.setSize(300, 300);
-        document.getElementById('avatar-container').appendChild(renderer.domElement);
-
-        var geometry = new THREE.SphereGeometry(1, 32, 32);
-        var material = new THREE.MeshStandardMaterial({ color: 0x00ffff, wireframe: true });
-        var avatar = new THREE.Mesh(geometry, material);
-        scene.add(avatar);
-
-        var light = new THREE.PointLight(0x00ffff, 1, 100);
-        light.position.set(2, 2, 5);
-        scene.add(light);
-
-        camera.position.z = 3;
-
-        function animate() {
-            requestAnimationFrame(animate);
-            avatar.rotation.y += 0.01;
-            renderer.render(scene, camera);
-        }
-        animate();
-    }
-    window.onload = create3DAvatar;
-  </script>
-</div>
-"""
-
-# Create Gradio Chatbot Interface with Custom UI & 3D Avatar
+# Create Gradio Chatbot Interface
 iface = gr.Interface(
     fn=chatbot_response,
     inputs="text",
     outputs="text",
     title="🤖 Cybersecurity AI Assistant",
-    description="Ask me about NIST, ISO/IEC 27001, MITRE ATLAS, and ISO/IEC 27005. Powered by AI and real-time 3D visualization.",
+    description="Ask me about NIST, ISO/IEC 27001, MITRE ATLAS, and ISO/IEC 27005. Powered by AI.",
     theme="default",
     css=custom_css,
     live=True,  # Enables real-time updates for typing effect
 )
 
-# Embed 3D Avatar into the interface
-iface.launch(share=True, custom_js=three_js_html)
+# Launch chatbot
+iface.launch()