app.py

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()