Update app.py

app.py

@@ -1,13 +1,13 @@
-# ------------------- MUST BE FIRST -------------------
 import streamlit as st
 import os
 import tempfile
 from pathlib import Path
+from knowledge_engine import KnowledgeManager  # Your class
 
-# Check if running in HuggingFace Spaces
+# Determine if running on HF Spaces or locally
 IS_HF_SPACES = os.getenv("SPACE_ID") is not None
 
-# Use temp directory for HF Spaces, local directory otherwise
+# Set knowledge base directory depending on environment
 if IS_HF_SPACES:
     KNOWLEDGE_DIR = Path(tempfile.gettempdir()) / "knowledge_base"
 else:
@@ -16,349 +16,59 @@ else:
 KNOWLEDGE_DIR.mkdir(parents=True, exist_ok=True)
 
 st.set_page_config(
-    page_title="Sirraya xBrain - Intelligent Assistant", 
-    layout="centered",
+    page_title="Simple Knowledge Assistant",
     page_icon="🧠",
-    initial_sidebar_state="collapsed"
+    layout="centered",
+    initial_sidebar_state="collapsed",
 )
-# -----------------------------------------------------
-
-from knowledge_engine import KnowledgeManager
 
 def check_hf_token():
-    """Check if HuggingFace token is available"""
-    token = os.getenv("HUGGINGFACEHUB_API_TOKEN") or os.getenv("HF_TOKEN")
-    return token is not None
-
-def initialize_lisa():
-    """Initialize LISA knowledge manager"""
-    if "lisa" not in st.session_state:
-        # Check for HF token first
-        if not check_hf_token():
-            st.error("❌ HuggingFace API token not found. Please set HUGGINGFACEHUB_API_TOKEN or HF_TOKEN environment variable.")
-            if not IS_HF_SPACES:
-                st.info("💡 For local development, you can set the token in your environment variables.")
-            else:
-                st.info("💡 For HuggingFace Spaces, add your token in the Spaces settings > Repository secrets.")
-            st.session_state.lisa = None
-            return
-        
-        with st.spinner("🔄 Initializing knowledge engine..."):
+    """Check if HuggingFace API token is set in environment variables."""
+    return os.getenv("HUGGINGFACEHUB_API_TOKEN") or os.getenv("HF_TOKEN")
+
+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()
+
+        with st.spinner("🔄 Initializing Knowledge Manager..."):
             try:
-                st.session_state.lisa = KnowledgeManager()
-                
-                # Check if initialization was successful
-                if hasattr(st.session_state.lisa, 'embeddings') and st.session_state.lisa.embeddings:
-                    st.success("✅ Knowledge engine initialized successfully!")
-                    
-                    # Show initialization status
-                    if st.session_state.lisa.llm:
-                        st.info("🤖 LLM model loaded successfully")
-                    else:
-                        st.warning("⚠️ LLM model failed to load - using fallback responses")
-                        
-                    # Show knowledge base summary
-                    summary = st.session_state.lisa.get_knowledge_summary()
-                    st.info(f"📚 {summary}")
-                else:
-                    st.error("❌ Failed to initialize embeddings. Check your internet connection.")
-                    st.session_state.lisa = None
-                    
+                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 system: {str(e)}")
-                st.session_state.lisa = None
-                
-                # Show debugging info
-                with st.expander("🐛 Debug Information"):
-                    st.code(f"Error: {e}")
-                    st.code(f"HF Spaces: {IS_HF_SPACES}")
-                    st.code(f"Token available: {check_hf_token()}")
-
-def render_sidebar():
-    """Render the sidebar for knowledge management"""
-    with st.sidebar:
-        st.header("📚 Knowledge Management")
-        
-        # Show environment info
-        if IS_HF_SPACES:
-            st.info("🚀 Running on HuggingFace Spaces")
-        else:
-            st.info("💻 Running locally")
-        
-        # Token status
-        if check_hf_token():
-            st.success("✅ HF Token detected")
-        else:
-            st.error("❌ HF Token missing")
-        
-        st.divider()
-        
-        # File upload section
-        st.subheader("📄 Add Knowledge")
-        
-        # Text input method (more reliable for HF Spaces)
-        st.markdown("**Method 1: Paste Text Content**")
-        filename = st.text_input("Filename", placeholder="my_document.txt")
-        content = st.text_area("Text Content", height=150, placeholder="Paste your text content here...")
-        
-        if st.button("➕ Add Text Content", type="primary"):
-            if filename and content and st.session_state.lisa:
-                try:
-                    success = st.session_state.lisa.add_text_content(filename, content)
-                    if success:
-                        st.success(f"✅ Added {filename} to knowledge base!")
-                        st.balloons()
-                    else:
-                        st.error("❌ Failed to add content")
-                except Exception as e:
-                    st.error(f"❌ Error adding content: {e}")
-            elif not filename or not content:
-                st.warning("⚠️ Please provide both filename and content")
-            else:
-                st.error("❌ Knowledge engine not initialized")
-        
-        st.divider()
-        
-        # File upload method (backup)
-        st.markdown("**Method 2: Upload File**")
-        uploaded_file = st.file_uploader(
-            "Upload text file", 
-            type=["txt", "md", "csv"],
-            help="Upload text files to expand LISA's knowledge base"
-        )
-        
-        if uploaded_file:
-            if st.session_state.lisa:
-                try:
-                    # Read file content
-                    file_content = uploaded_file.read()
-                    
-                    # Add to knowledge base using the new method
-                    success = st.session_state.lisa.add_uploaded_file(file_content, uploaded_file.name)
-                    
-                    if success:
-                        st.success(f"✅ Added {uploaded_file.name} to knowledge base!")
-                        st.balloons()
-                    else:
-                        st.error("❌ Failed to add uploaded file")
-                        
-                except Exception as e:
-                    st.error(f"❌ Error processing upload: {e}")
-            else:
-                st.error("❌ Knowledge engine not initialized")
-        
-        st.divider()
-        
-        # Rebuild button (for troubleshooting)
-        if st.button("🔄 Restart System", help="Reinitialize the entire system"):
-            # Clear session state
-            if "lisa" in st.session_state:
-                if hasattr(st.session_state.lisa, 'cleanup'):
-                    st.session_state.lisa.cleanup()
-                del st.session_state.lisa
-            
-            if "messages" in st.session_state:
-                del st.session_state.messages
-                
-            st.success("🔄 System restarted. Page will reload...")
-            st.rerun()
-        
-        st.divider()
-        
-        # System info section
-        st.subheader("🔧 System Status")
-        
-        if st.session_state.lisa:
-            # Show system information
-            file_count = st.session_state.lisa.get_knowledge_files_count()
-            st.metric("📄 Knowledge Files", file_count)
-            
-            # Show detailed status
-            with st.expander("📊 Detailed Status"):
-                st.text(f"Embeddings: {'✅' if st.session_state.lisa.embeddings else '❌'}")
-                st.text(f"Vector DB: {'✅' if st.session_state.lisa.vector_db else '❌'}")
-                st.text(f"BM25: {'✅' if st.session_state.lisa.bm25_retriever else '❌'}")
-                st.text(f"LLM: {'✅' if st.session_state.lisa.llm else '❌'}")
-                st.text(f"QA Chain: {'✅' if st.session_state.lisa.qa_chain else '❌'}")
-        else:
-            st.error("❌ System not initialized")
-        
-        # Model info
-        with st.expander("🤖 Model Information"):
-            st.markdown("""
-            **Embeddings:** sentence-transformers/all-mpnet-base-v2  
-            **LLM:** mistralai/Mistral-7B-Instruct-v0.1  
-            **Fallbacks:** Multiple models available  
-            **Retrieval:** Hybrid (Vector + BM25)  
-            **Storage:** In-memory + temp files
-            """)
-
-def render_chat_interface():
-    """Render the main chat interface"""
-    # Initialize chat history
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
-    
-    # Show welcome message if no messages
-    if not st.session_state.messages:
-        with st.chat_message("assistant"):
-            welcome_msg = """
-            👋 Hello! I'm LISA, your intelligent assistant powered by Sirraya xBrain.
-            
-            I can help you with questions about any knowledge you've added to my database. 
-            
-            **To get started:**
-            1. Add some knowledge using the sidebar
-            2. Ask me questions about that content
-            3. I'll provide accurate answers with source citations
-            
-            What would you like to know?
-            """
-            st.markdown(welcome_msg)
-    
-    # Display chat history
-    for msg in st.session_state.messages:
-        with st.chat_message(msg["role"]):
-            st.write(msg["content"])
-            
-            # Show sources if available
-            if msg["role"] == "assistant" and msg.get("sources"):
-                with st.expander(f"📖 View Sources ({len(msg['sources'])} found)", expanded=False):
-                    for i, source in enumerate(msg["sources"]):
-                        st.markdown(f"**📄 Source {i+1}:**")
-                        
-                        # Truncate long content
-                        content = source.page_content
-                        if len(content) > 400:
-                            content = content[:400] + "..."
-                        st.text(content)
-                        
-                        # Show metadata if available
-                        if hasattr(source, 'metadata') and source.metadata:
-                            source_file = source.metadata.get('source', 'Unknown')
-                            st.caption(f"📁 From: {source_file}")
-                        
-                        if i < len(msg["sources"]) - 1:
-                            st.divider()
-    
-    # Handle new user query
-    if prompt := st.chat_input("Ask LISA about anything in your knowledge base..."):
-        # Add user message
-        st.session_state.messages.append({"role": "user", "content": prompt})
-        
-        with st.chat_message("user"):
-            st.write(prompt)
-        
-        # Generate response
-        with st.chat_message("assistant"):
-            if st.session_state.lisa:
-                with st.spinner("🤔 Thinking..."):
-                    try:
-                        result = st.session_state.lisa.query(prompt)
-                        
-                        # Display answer
-                        st.write(result["answer"])
-                        
-                        # Show processing time and stats
-                        col1, col2 = st.columns(2)
-                        with col1:
-                            if result["processing_time"] > 0:
-                                st.caption(f"⚡ {result['processing_time']:.0f}ms")
-                        with col2:
-                            source_count = len(result["source_chunks"]) if result["source_chunks"] else 0
-                            st.caption(f"📚 {source_count} sources")
-                        
-                        # Store message with sources
-                        st.session_state.messages.append({
-                            "role": "assistant",
-                            "content": result["answer"],
-                            "sources": result["source_chunks"] if result["source_chunks"] else None
-                        })
-                        
-                    except Exception as e:
-                        error_msg = f"❌ Error processing query: {str(e)}"
-                        st.error(error_msg)
-                        st.session_state.messages.append({
-                            "role": "assistant",
-                            "content": error_msg
-                        })
-            else:
-                error_msg = "❌ LISA is not properly initialized. Please check the sidebar for system status and try restarting."
-                st.error(error_msg)
-                st.session_state.messages.append({
-                    "role": "assistant",
-                    "content": error_msg
-                })
-
-def render_sample_content():
-    """Render sample content section"""
-    with st.expander("💡 Need sample content to test? Click here!", expanded=False):
-        st.markdown("**Here's some sample content you can add to test the system:**")
-        
-        sample_content = """Artificial Intelligence (AI) is a rapidly evolving field that focuses on creating intelligent machines capable of performing tasks that typically require human intelligence. These tasks include learning, reasoning, problem-solving, perception, and language understanding.
-
-Machine Learning is a subset of AI that enables computers to learn and improve from experience without being explicitly programmed. It uses algorithms and statistical models to analyze and draw inferences from patterns in data.
+                st.error(f"❌ Error initializing Knowledge Manager:\n{e}")
+                st.stop()
 
-Deep Learning is a specialized subset of machine learning that uses artificial neural networks with multiple layers to model and understand complex patterns in data. It has been particularly successful in areas like image recognition, natural language processing, and speech recognition.
+# Initialize KnowledgeManager once
+initialize_knowledge_manager()
 
-Natural Language Processing (NLP) is a branch of AI that helps computers understand, interpret, and generate human language in a valuable way. It combines computational linguistics with statistical, machine learning, and deep learning models."""
+km = st.session_state.get("knowledge_manager")
 
-        if st.button("📝 Add Sample AI Content"):
-            if st.session_state.lisa:
-                success = st.session_state.lisa.add_text_content("ai_basics.txt", sample_content)
-                if success:
-                    st.success("✅ Sample AI content added! Try asking: 'What is machine learning?'")
-                else:
-                    st.error("❌ Failed to add sample content")
-            else:
-                st.error("❌ Please initialize the system first")
-
-def main():
-    """Main application function"""
-    # Header with better styling
-    st.title("🧠 Sirraya xBrain")
-    st.markdown("### *LISA - Intelligent Assistant powered by Advanced RAG Technology*")
-    st.markdown("*Created by Amir Hameed*")
-    
-    # Show environment status
-    col1, col2, col3 = st.columns(3)
-    with col1:
-        if IS_HF_SPACES:
-            st.info("🚀 HF Spaces")
-        else:
-            st.info("💻 Local")
-    with col2:
-        if check_hf_token():
-            st.success("🔑 Token OK")
-        else:
-            st.error("🔑 No Token")
-    with col3:
-        if "lisa" in st.session_state and st.session_state.lisa:
-            st.success("🤖 LISA Ready")
-        else:
-            st.warning("🤖 LISA Loading")
-    
-    st.divider()
-    
-    # Initialize LISA
-    initialize_lisa()
-    
-    # Show sample content helper
-    if "lisa" in st.session_state and st.session_state.lisa:
-        file_count = st.session_state.lisa.get_knowledge_files_count()
-        if file_count <= 2:  # Only default files
-            render_sample_content()
-    
-    # Render sidebar
-    render_sidebar()
+# If KnowledgeManager is initialized, show status and simple query input
+if km:
+    st.info("🤖 Model & knowledge loaded and ready.")
     
-    # Render chat interface
-    render_chat_interface()
-    
-    # Footer
-    st.divider()
-    st.markdown("---")
-    st.markdown("*Powered by Sirraya xBrain | Advanced Hybrid RAG System*")
-
-if __name__ == "__main__":
-    main()
+    # Show knowledge summary if available
+    if hasattr(km, "get_knowledge_summary"):
+        summary = km.get_knowledge_summary()
+        st.markdown(f"**Knowledge Summary:** {summary}")
+
+    # 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.")