Create app.py

app.py

@@ -0,0 +1,807 @@
+import streamlit as st
+import numpy as np
+import uuid
+import json
+import os
+import time
+from datetime import datetime
+from huggingface_hub import InferenceClient
+from sentence_transformers import SentenceTransformer
+from sklearn.metrics.pairwise import cosine_similarity
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from openai import OpenAI
+
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
+
+# Create necessary directories if they don't exist
+os.makedirs("data/sessions", exist_ok=True)
+os.makedirs("data/documents", exist_ok=True)
+os.makedirs("data/embeddings", exist_ok=True)
+
+# Configure page settings
+st.set_page_config(
+    page_title="Matrix AI Chat with RAG",
+    page_icon="🕶️",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
+
+# Matrix-style CSS
+def load_css():
+    matrix_css = """
+    <style>
+    @import url('https://fonts.googleapis.com/css2?family=Courier+New:wght@400;700&display=swap');
+    
+    /* Global Matrix styling */
+    .stApp {
+        background-color: #000000 !important;
+        color: #00ff00 !important;
+        font-family: 'Courier New', monospace !important;
+    }
+    
+    /* Main content area */
+    .main .block-container {
+        background-color: #000000 !important;
+        color: #00ff00 !important;
+    }
+    
+    /* Sidebar */
+    .css-1d391kg {
+        background-color: #000000 !important;
+        border-right: 2px solid #00ff00 !important;
+    }
+    
+    /* Chat messages */
+    .stChatMessage {
+        background-color: #001100 !important;
+        border: 1px solid #00ff00 !important;
+        border-radius: 5px !important;
+        padding: 15px !important;
+        margin: 10px 0 !important;
+        color: #00ff00 !important;
+        font-family: 'Courier New', monospace !important;
+        box-shadow: 0 0 10px rgba(0, 255, 0, 0.3) !important;
+    }
+    
+    /* Input containers */
+    .stTextInput > div > div > input,
+    .stTextArea > div > div > textarea {
+        background-color: #000000 !important;
+        color: #00ff00 !important;
+        border: 1px solid #00ff00 !important;
+        font-family: 'Courier New', monospace !important;
+    }
+    
+    /* Selectbox */
+    .stSelectbox > div > div > div {
+        background-color: #000000 !important;
+        color: #00ff00 !important;
+        border: 1px solid #00ff00 !important;
+        font-family: 'Courier New', monospace !important;
+    }
+    
+    /* Buttons */
+    .stButton > button {
+        background-color: #000000 !important;
+        color: #00ff00 !important;
+        border: 1px solid #00ff00 !important;
+        font-family: 'Courier New', monospace !important;
+        font-weight: bold !important;
+        transition: all 0.3s ease !important;
+    }
+    
+    .stButton > button:hover {
+        background-color: #00ff00 !important;
+        color: #000000 !important;
+        box-shadow: 0 0 15px rgba(0, 255, 0, 0.7) !important;
+    }
+    
+    /* Headers */
+    h1, h2, h3, h4, h5, h6 {
+        color: #00ff00 !important;
+        font-family: 'Courier New', monospace !important;
+        text-shadow: 0 0 10px rgba(0, 255, 0, 0.8) !important;
+    }
+    
+    /* Main header */
+    .main-header {
+        text-align: center;
+        color: #00ff00 !important;
+        margin-bottom: 2rem;
+        font-size: 3rem !important;
+        text-shadow: 0 0 20px rgba(0, 255, 0, 1) !important;
+        animation: matrix-glow 2s ease-in-out infinite alternate;
+    }
+    
+    @keyframes matrix-glow {
+        from { text-shadow: 0 0 20px rgba(0, 255, 0, 0.8); }
+        to { text-shadow: 0 0 30px rgba(0, 255, 0, 1), 0 0 40px rgba(0, 255, 0, 0.8); }
+    }
+    
+    /* Status indicators */
+    .status-success {
+        color: #00ff00 !important;
+        font-weight: bold !important;
+        text-shadow: 0 0 5px rgba(0, 255, 0, 0.8) !important;
+    }
+    
+    .status-error {
+        color: #ff0000 !important;
+        font-weight: bold !important;
+        text-shadow: 0 0 5px rgba(255, 0, 0, 0.8) !important;
+    }
+    
+    /* Chat input */
+    .stChatInputContainer {
+        background-color: #000000 !important;
+        border-top: 1px solid #00ff00 !important;
+    }
+    
+    /* Expander */
+    .streamlit-expanderHeader {
+        background-color: #000000 !important;
+        color: #00ff00 !important;
+        border: 1px solid #00ff00 !important;
+    }
+    
+    /* Info boxes */
+    .stInfo {
+        background-color: #001100 !important;
+        color: #00ff00 !important;
+        border: 1px solid #00ff00 !important;
+    }
+    
+    /* Warning boxes */
+    .stWarning {
+        background-color: #110100 !important;
+        color: #ffff00 !important;
+        border: 1px solid #ffff00 !important;
+    }
+    
+    /* Error boxes */
+    .stError {
+        background-color: #110000 !important;
+        color: #ff0000 !important;
+        border: 1px solid #ff0000 !important;
+    }
+    
+    /* Success boxes */
+    .stSuccess {
+        background-color: #001100 !important;
+        color: #00ff00 !important;
+        border: 1px solid #00ff00 !important;
+    }
+    
+    /* Spinner */
+    .stSpinner {
+        color: #00ff00 !important;
+    }
+    
+    /* Caption */
+    .caption {
+        color: #00aa00 !important;
+        font-family: 'Courier New', monospace !important;
+        text-align: center;
+        font-style: italic;
+    }
+    
+    /* Matrix rain effect */
+    .matrix-bg::before {
+        content: "";
+        position: fixed;
+        top: 0;
+        left: 0;
+        width: 100%;
+        height: 100%;
+        background: repeating-linear-gradient(
+            90deg,
+            transparent,
+            transparent 98px,
+            rgba(0, 255, 0, 0.03) 100px
+        );
+        pointer-events: none;
+        z-index: -1;
+    }
+    
+    /* Model selection highlight */
+    .model-selector {
+        border: 2px solid #00ff00 !important;
+        border-radius: 5px !important;
+        padding: 10px !important;
+        background-color: #001100 !important;
+        margin: 10px 0 !important;
+    }
+    
+    /* Scrollbar */
+    ::-webkit-scrollbar {
+        width: 12px;
+    }
+    
+    ::-webkit-scrollbar-track {
+        background: #000000;
+    }
+    
+    ::-webkit-scrollbar-thumb {
+        background: #00ff00;
+        border-radius: 6px;
+    }
+    
+    ::-webkit-scrollbar-thumb:hover {
+        background: #00aa00;
+    }
+    </style>
+    """
+    st.markdown(matrix_css, unsafe_allow_html=True)
+
+# Model configurations
+MODEL_CONFIGS = {
+    "DeepSeek-R1": {
+        "provider": "together",
+        "model_name": "deepseek-ai/DeepSeek-R1-0528",
+        "type": "api"
+    },
+    "Llama-3.2-3B": {
+        "provider": "huggingface",
+        "model_name": "meta-llama/Llama-3.2-3B",
+        "type": "local"
+    },
+    "Qwen2.5-VL-7B-Instruct": {
+        "provider": "hyperbolic",
+        "model_name": "Qwen/Qwen2.5-VL-7B-Instruct",
+        "type": "api"
+    }
+}
+
+# Initialize clients based on selected model
+@st.cache_resource
+def get_model_client(model_name):
+    try:
+        if not HF_TOKEN:
+            st.error("❌ Hugging Face token is required!")
+            return None, None
+        
+        config = MODEL_CONFIGS[model_name]
+        
+        if config["type"] == "api":
+            if config["provider"] == "together":
+                client = InferenceClient(
+                    provider="together",
+                    api_key=HF_TOKEN,
+                )
+                return client, config
+            elif config["provider"] == "hyperbolic":
+                client = OpenAI(
+                    base_url="https://router.huggingface.co/hyperbolic/v1",
+                    api_key=HF_TOKEN,
+                )
+                return client, config
+        elif config["type"] == "local":
+            # For local models, we'll load tokenizer and model
+            tokenizer = AutoTokenizer.from_pretrained(config["model_name"])
+            model = AutoModelForCausalLM.from_pretrained(config["model_name"])
+            return (tokenizer, model), config
+        
+        return None, None
+    except Exception as e:
+        st.error(f"❌ Error initializing {model_name} client: {e}")
+        return None, None
+
+# Initialize session management
+def get_session_id():
+    if "session_id" not in st.session_state:
+        st.session_state.session_id = str(uuid.uuid4())
+        save_session_metadata(st.session_state.session_id)
+    return st.session_state.session_id
+
+# Save session metadata
+def save_session_metadata(session_id):
+    try:
+        session_file = f"data/sessions/{session_id}_metadata.json"
+        metadata = {
+            "session_id": session_id,
+            "created_at": datetime.now().isoformat(),
+            "last_updated": datetime.now().isoformat()
+        }
+        with open(session_file, "w") as f:
+            json.dump(metadata, f, indent=2)
+    except Exception as e:
+        st.warning(f"Could not save session metadata: {e}")
+
+# Update session timestamp
+def update_session_timestamp(session_id):
+    try:
+        session_file = f"data/sessions/{session_id}_metadata.json"
+        if os.path.exists(session_file):
+            with open(session_file, "r") as f:
+                metadata = json.load(f)
+            metadata["last_updated"] = datetime.now().isoformat()
+            with open(session_file, "w") as f:
+                json.dump(metadata, f, indent=2)
+    except Exception as e:
+        st.warning(f"Could not update session timestamp: {e}")
+
+# Save chat history
+def save_chat_history(prompt, response, embedding=None, context=""):
+    try:
+        session_id = get_session_id()
+        history_file = f"data/sessions/{session_id}_history.json"
+        
+        # Load existing history or create new
+        if os.path.exists(history_file):
+            with open(history_file, "r") as f:
+                history = json.load(f)
+        else:
+            history = []
+        
+        # Get message order
+        message_order = len(history) + 1
+        
+        # Create history entry
+        entry = {
+            "message_id": message_order,
+            "prompt": prompt,
+            "response": response,
+            "context": context,
+            "timestamp": datetime.now().isoformat()
+        }
+        
+        # Save embedding if available
+        if embedding is not None:
+            embedding_file = f"data/embeddings/{session_id}_{message_order}.npy"
+            np.save(embedding_file, np.array(embedding))
+            entry["embedding_path"] = embedding_file
+        
+        # Append and save history
+        history.append(entry)
+        with open(history_file, "w") as f:
+            json.dump(history, f, indent=2)
+        
+        # Update session timestamp
+        update_session_timestamp(session_id)
+    except Exception as e:
+        st.warning(f"Could not save chat history: {e}")
+
+# Add a document to the RAG system
+def add_document(title, content, embedding=None):
+    try:
+        # Generate document ID
+        doc_id = str(uuid.uuid4())
+        
+        # Save document
+        document_file = f"data/documents/{doc_id}.json"
+        document = {
+            "id": doc_id,
+            "title": title,
+            "content": content,
+            "created_at": datetime.now().isoformat()
+        }
+        
+        with open(document_file, "w") as f:
+            json.dump(document, f, indent=2)
+        
+        # Save embedding if available
+        if embedding is not None:
+            embedding_file = f"data/embeddings/doc_{doc_id}.npy"
+            np.save(embedding_file, np.array(embedding))
+            
+            # Save embedding reference
+            document["embedding_path"] = embedding_file
+            with open(document_file, "w") as f:
+                json.dump(document, f, indent=2)
+        
+        return doc_id
+    except Exception as e:
+        st.error(f"Error adding document: {e}")
+        return None
+
+# Function to get embedding vector
+@st.cache_resource
+def load_embedding_model():
+    try:
+        model = SentenceTransformer('all-MiniLM-L6-v2')
+        return model
+    except Exception as e:
+        st.error(f"Error loading embedding model: {e}")
+        return None
+
+def get_embedding(text):
+    model = load_embedding_model()
+    if model:
+        try:
+            return model.encode(text)
+        except Exception as e:
+            st.warning(f"Embedding error: {e}")
+    return None
+
+# Generate conversation context
+def generate_context(user_query, max_turns=3):
+    try:
+        session_id = get_session_id()
+        history_file = f"data/sessions/{session_id}_history.json"
+        
+        if not os.path.exists(history_file):
+            return ""
+        
+        with open(history_file, "r") as f:
+            history = json.load(f)
+        
+        # Get last N conversation turns
+        recent_history = history[-max_turns:] if len(history) >= max_turns else history
+        
+        # Create context string
+        context = ""
+        for entry in recent_history:
+            context += f"User: {entry['prompt']}\nAssistant: {entry['response']}\n\n"
+        
+        return context.strip()
+    except Exception as e:
+        st.warning(f"Error generating context: {e}")
+        return ""
+
+# Fetch stored embeddings
+def fetch_embeddings():
+    try:
+        session_id = get_session_id()
+        history_file = f"data/sessions/{session_id}_history.json"
+        
+        if not os.path.exists(history_file):
+            return [], np.array([])
+        
+        with open(history_file, "r") as f:
+            history = json.load(f)
+        
+        prompts, responses, embeddings, contexts = [], [], [], []
+        
+        for entry in history:
+            if "embedding_path" in entry and os.path.exists(entry["embedding_path"]):
+                try:
+                    embedding = np.load(entry["embedding_path"])
+                    embeddings.append(embedding)
+                    prompts.append(entry["prompt"])
+                    responses.append(entry["response"])
+                    contexts.append(entry.get("context", ""))
+                except Exception:
+                    continue  # Skip corrupted embeddings
+        
+        return list(zip(prompts, responses, contexts)), np.array(embeddings) if embeddings else np.array([])
+    except Exception as e:
+        st.warning(f"Error fetching embeddings: {e}")
+        return [], np.array([])
+
+# Search for similar documents in the RAG system
+def search_rag_documents(query_embedding, top_k=3, threshold=0.7):
+    try:
+        if not os.path.exists("data/documents"):
+            return []
+        
+        results = []
+        document_files = [f for f in os.listdir("data/documents") if f.endswith(".json")]
+        
+        for doc_file in document_files:
+            try:
+                with open(f"data/documents/{doc_file}", "r") as f:
+                    document = json.load(f)
+                
+                # Check if embedding exists
+                if "embedding_path" in document and os.path.exists(document["embedding_path"]):
+                    doc_embedding = np.load(document["embedding_path"])
+                    
+                    # Calculate similarity
+                    similarity = cosine_similarity([query_embedding], [doc_embedding])[0][0]
+                    
+                    # Add if above threshold
+                    if similarity >= threshold:
+                        results.append((
+                            document["id"], 
+                            document["title"], 
+                            document["content"], 
+                            similarity
+                        ))
+            except Exception:
+                continue  # Skip corrupted documents
+        
+        # Sort by similarity score (descending)
+        results.sort(key=lambda x: x[3], reverse=True)
+        return results[:top_k]
+    except Exception as e:
+        st.warning(f"Error searching RAG documents: {e}")
+        return []
+
+# Similarity Search Function
+def find_similar_response(user_query, user_embedding, threshold=0.85):
+    try:
+        # First check for similar responses in conversation history
+        data, embeddings = fetch_embeddings()
+        
+        if embeddings.size > 0:
+            similarities = cosine_similarity([user_embedding], embeddings)[0]
+            best_match_index = np.argmax(similarities)
+            
+            if similarities[best_match_index] >= threshold:
+                matched_prompt, matched_response, matched_context = data[best_match_index]
+                return matched_response, ""
+        
+        # If no match in history, search RAG documents
+        rag_results = search_rag_documents(user_embedding)
+        if rag_results:
+            context_docs = "\n\n".join([
+                f"**{title}**\n{content}" 
+                for _, title, content, _ in rag_results
+            ])
+            return None, context_docs
+        
+        return None, ""
+    except Exception as e:
+        st.warning(f"Error in similarity search: {e}")
+        return None, ""
+
+# Generate response using selected model
+def generate_response(prompt, system_prompt="", rag_context="", selected_model="DeepSeek-R1"):
+    client, config = get_model_client(selected_model)
+    
+    if not client:
+        return f"❌ {selected_model} client not available. Please check your configuration."
+    
+    try:
+        # Construct user content
+        user_content = prompt
+        if rag_context:
+            user_content = f"Context information:\n{rag_context}\n\nQuestion: {prompt}"
+        
+        if config["type"] == "api":
+            # Handle API-based models
+            messages = []
+            
+            if system_prompt:
+                messages.append({
+                    "role": "system",
+                    "content": system_prompt
+                })
+            
+            messages.append({
+                "role": "user",
+                "content": user_content
+            })
+            
+            # Generate response based on provider
+            if config["provider"] == "together":
+                completion = client.chat.completions.create(
+                    model=config["model_name"],
+                    messages=messages,
+                    max_tokens=1000,
+                    temperature=0.7,
+                    top_p=0.9,
+                )
+                return completion.choices[0].message.content
+            
+            elif config["provider"] == "hyperbolic":
+                completion = client.chat.completions.create(
+                    model=config["model_name"],
+                    messages=messages,
+                    max_tokens=1000,
+                    temperature=0.7,
+                )
+                return completion.choices[0].message.content
+        
+        elif config["type"] == "local":
+            # Handle local models
+            tokenizer, model = client
+            
+            # Prepare input
+            full_prompt = f"{system_prompt}\n\nUser: {user_content}\nAssistant:"
+            inputs = tokenizer(full_prompt, return_tensors="pt")
+            
+            # Generate response
+            with torch.no_grad():
+                outputs = model.generate(
+                    inputs.input_ids,
+                    max_length=inputs.input_ids.shape[1] + 500,
+                    temperature=0.7,
+                    do_sample=True,
+                    pad_token_id=tokenizer.eos_token_id
+                )
+            
+            # Decode response
+            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            # Extract only the assistant's response
+            response = response.split("Assistant:")[-1].strip()
+            return response
+        
+    except Exception as e:
+        st.error(f"Error generating response: {e}")
+        return f"I apologize, but I encountered an error while processing your request: {str(e)}"
+
+# Main UI
+def main():
+    # Load CSS
+    load_css()
+    
+    # Matrix background div
+    st.markdown('<div class="matrix-bg"></div>', unsafe_allow_html=True)
+    
+    st.markdown('<h1 class="main-header">🕶️ MATRIX AI CHAT</h1>', unsafe_allow_html=True)
+    st.markdown('<p class="caption">ENTER THE MATRIX: Advanced AI with Retrieval-Augmented Generation</p>', unsafe_allow_html=True)
+    
+    # Get session ID
+    session_id = get_session_id()
+    
+    # Sidebar Configuration
+    with st.sidebar:
+        st.markdown("## ⚙️ MATRIX CONTROL PANEL")
+        
+        # Model Selection
+        st.markdown('<div class="model-selector">', unsafe_allow_html=True)
+        st.markdown("### 🤖 AI MODEL SELECTION")
+        selected_model = st.selectbox(
+            "Choose your AI:",
+            options=list(MODEL_CONFIGS.keys()),
+            index=0,
+            help="Select the AI model to power your conversations"
+        )
+        st.markdown('</div>', unsafe_allow_html=True)
+        
+        # Token status
+        if HF_TOKEN:
+            st.markdown('<p class="status-success">✅ HUGGING FACE TOKEN: CONNECTED</p>', unsafe_allow_html=True)
+            
+            # Test connection
+            client, config = get_model_client(selected_model)
+            if client:
+                st.markdown(f'<p class="status-success">✅ {selected_model}: READY</p>', unsafe_allow_html=True)
+            else:
+                st.markdown(f'<p class="status-error">❌ {selected_model}: CONNECTION FAILED</p>', unsafe_allow_html=True)
+        else:
+            st.markdown('<p class="status-error">❌ NO HUGGING FACE TOKEN FOUND</p>', unsafe_allow_html=True)
+            st.info("Please set your HF_TOKEN environment variable to enter the Matrix.")
+        
+        st.divider()
+        
+        # System prompt configuration
+        system_prompt = st.text_area(
+            "SYSTEM PROMPT",
+            value=f"You are {selected_model}, an advanced AI assistant operating within the Matrix. Provide accurate, detailed, and helpful responses. If given context information, use it to enhance your answers. Embrace the digital realm.",
+            height=120,
+            help="Define how the AI should behave in the Matrix"
+        )
+        
+        st.divider()
+        
+        # Session controls
+        st.markdown("### 🔄 SESSION CONTROLS")
+        col1, col2 = st.columns(2)
+        
+        with col1:
+            if st.button("NEW JACK IN", use_container_width=True):
+                # Reset session
+                for key in ["session_id", "message_log"]:
+                    if key in st.session_state:
+                        del st.session_state[key]
+                st.rerun()
+        
+        with col2:
+            if st.button("PURGE ALL", use_container_width=True):
+                # Clear all data (with confirmation)
+                if st.session_state.get("confirm_clear", False):
+                    try:
+                        import shutil
+                        if os.path.exists("data"):
+                            shutil.rmtree("data")
+                        os.makedirs("data/sessions", exist_ok=True)
+                        os.makedirs("data/documents", exist_ok=True)
+                        os.makedirs("data/embeddings", exist_ok=True)
+                        st.success("Matrix data purged!")
+                        st.session_state.confirm_clear = False
+                        st.rerun()
+                    except Exception as e:
+                        st.error(f"Error purging Matrix: {e}")
+                else:
+                    st.session_state.confirm_clear = True
+                    st.warning("Click again to confirm Matrix purge")
+        
+        st.divider()
+        
+        # RAG Document Upload
+        st.markdown("### 📚 KNOWLEDGE MATRIX")
+        
+        with st.expander("UPLOAD DATA"):
+            doc_title = st.text_input("DATA TITLE", placeholder="Enter data identifier...")
+            doc_content = st.text_area(
+                "DATA CONTENT", 
+                placeholder="Upload your knowledge to the Matrix...", 
+                height=200
+            )
+            
+            if st.button("📝 INJECT DATA", use_container_width=True):
+                if doc_title and doc_content:
+                    with st.spinner("Integrating into Matrix..."):
+                        doc_embedding = get_embedding(doc_content)
+                        doc_id = add_document(doc_title, doc_content, doc_embedding)
+                        if doc_id:
+                            st.success(f"✅ Data '{doc_title}' integrated into Matrix!")
+                        else:
+                            st.error("❌ Failed to integrate data")
+                else:
+                    st.warning("Please provide both title and content")
+        
+        # Display document count
+        try:
+            doc_count = len([f for f in os.listdir("data/documents") if f.endswith(".json")])
+            st.info(f"📄 {doc_count} data nodes in Matrix")
+        except:
+            st.info("📄 0 data nodes in Matrix")
+        
+        st.divider()
+        st.markdown(f"**SESSION ID:** `{session_id[:8]}...`")
+    
+    # Initialize chat history
+    if "message_log" not in st.session_state:
+        st.session_state.message_log = [{
+            "role": "assistant", 
+            "content": f"🕶️ Welcome to the Matrix. I am {selected_model}, your guide through the digital realm. The red pill or the blue pill - what will you choose to explore today?"
+        }]
+    
+    # Display chat history
+    for message in st.session_state.message_log:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+    
+    # Chat input
+    user_query = st.chat_input("Enter your query into the Matrix...")
+    
+    # Process user query
+    if user_query and HF_TOKEN:
+        # Add user message to chat
+        st.session_state.message_log.append({"role": "user", "content": user_query})
+        
+        # Display user message
+        with st.chat_message("user"):
+            st.markdown(user_query)
+        
+        # Generate response
+        with st.chat_message("assistant"):
+            with st.spinner(f"🧠 {selected_model} is processing in the Matrix..."):
+                # Get embedding for similarity search
+                user_embedding = get_embedding(user_query)
+                
+                # Check for similar responses or RAG context
+                cached_response = None
+                rag_context = ""
+                
+                if user_embedding is not None:
+                    cached_response, rag_context = find_similar_response(user_query, user_embedding)
+                
+                if cached_response:
+                    # Use cached response
+                    st.info("🔍 Found similar data in Matrix")
+                    response_text = cached_response
+                else:
+                    # Generate new response
+                    response_text = generate_response(user_query, system_prompt, rag_context, selected_model)
+                
+                # Display response with Matrix-style streaming effect
+                response_placeholder = st.empty()
+                displayed_response = ""
+                
+                # Simulate Matrix-style streaming
+                for char in response_text:
+                    displayed_response += char
+                    response_placeholder.markdown(displayed_response + "█")
+                    time.sleep(0.02)  # Slightly slower for Matrix effect
+                
+                # Final response
+                response_placeholder.markdown(response_text)
+        
+        # Add response to chat history
+        st.session_state.message_log.append({"role": "assistant", "content": response_text})
+        
+        # Save to persistent storage
+        save_chat_history(user_query, response_text, user_embedding, generate_context(user_query))
+        
+        # Rerun to update UI
+        st.rerun()
+    
+    elif user_query and not HF_TOKEN:
+        st.error("❌ Please set your Hugging Face token to enter the Matrix.")
+
+if __name__ == "__main__":
+    main()