Update app.py

app.py

@@ -1,74 +1,41 @@
-import streamlit as st
 import os
-import tempfile
-from pathlib import Path
-from knowledge_engine import KnowledgeManager  # Your class
-
-# Determine if running on HF Spaces or locally
-IS_HF_SPACES = os.getenv("SPACE_ID") is not None
-
-# Set knowledge base directory depending on environment
-if IS_HF_SPACES:
-    KNOWLEDGE_DIR = Path(tempfile.gettempdir()) / "knowledge_base"
-else:
-    KNOWLEDGE_DIR = Path("knowledge_base")
-
-KNOWLEDGE_DIR.mkdir(parents=True, exist_ok=True)
+import streamlit as st
+from knowledge_manager import KnowledgeManager
 
 st.set_page_config(
-    page_title="Simple Knowledge Assistant",
-    page_icon="🧠",
+    page_title="Sirraya xBrain - LangChain QA Assistant",
     layout="centered",
-    initial_sidebar_state="collapsed",
+    page_icon="🧠",
+    initial_sidebar_state="expanded",
 )
 
-def check_hf_token():
-    """Check if HuggingFace API token is set in environment variables."""
-    return os.getenv("HUGGINGFACEHUB_API_TOKEN") or os.getenv("HF_TOKEN")
+st.title("🧠 Sirraya xBrain — Intelligent QA Assistant")
 
-def initialize_knowledge_manager():
-    """Initialize KnowledgeManager and handle errors."""
-    if "knowledge_manager" not in st.session_state:
-        token = check_hf_token()
-        if not token:
-            st.error("❌ HuggingFace API token not found.\n\n"
-                     "Set HUGGINGFACEHUB_API_TOKEN or HF_TOKEN environment variable.")
-            st.stop()
+if "km" not in st.session_state:
+    with st.spinner("Initializing knowledge base and LLM..."):
+        try:
+            st.session_state.km = KnowledgeManager()
+            st.success("✅ Knowledge engine initialized successfully!")
+            st.info(st.session_state.km.get_knowledge_summary())
+        except Exception as e:
+            st.error(f"❌ Failed to initialize system: {e}")
+            st.session_state.km = None
 
-        with st.spinner("🔄 Initializing Knowledge Manager..."):
-            try:
-                km = KnowledgeManager()
-                if not hasattr(km, "embeddings") or not km.embeddings:
-                    st.error("❌ Embeddings failed to initialize.")
-                    st.stop()
-                st.session_state.knowledge_manager = km
-                st.success("✅ Knowledge Manager initialized successfully!")
-            except Exception as e:
-                st.error(f"❌ Error initializing Knowledge Manager:\n{e}")
-                st.stop()
+if st.session_state.km is None:
+    st.warning("Knowledge base not loaded or failed to initialize.")
+    st.stop()
 
-# Initialize KnowledgeManager once
-initialize_knowledge_manager()
+question = st.text_input("Ask a question about the knowledge base:", "")
 
-km = st.session_state.get("knowledge_manager")
+if question:
+    with st.spinner("Generating answer..."):
+        answer, sources = st.session_state.km.query(question)
 
-# If KnowledgeManager is initialized, show status and simple query input
-if km:
-    st.info("🤖 Model & knowledge loaded and ready.")
-    
-    # Show knowledge summary if available
-    if hasattr(km, "get_knowledge_summary"):
-        summary = km.get_knowledge_summary()
-        st.markdown(f"**Knowledge Summary:** {summary}")
+        st.markdown("### Answer:")
+        st.write(answer)
 
-    # Simple input to ask a question and get an answer (assuming km has an 'ask' method)
-    question = st.text_input("Ask something from the knowledge base:")
-    if question:
-        with st.spinner("🧠 Thinking..."):
-            try:
-                answer = km.ask(question)  # Replace with your actual query method
-                st.markdown(f"**Answer:** {answer}")
-            except Exception as e:
-                st.error(f"❌ Failed to get answer: {e}")
-else:
-    st.warning("Knowledge Manager is not initialized yet.")
+        if sources:
+            with st.expander("📚 Source documents"):
+                for i, src in enumerate(sources, 1):
+                    st.write(f"Source {i}:")
+                    st.write(src)