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 = """
"""
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('', unsafe_allow_html=True)
st.markdown('🕶️ MATRIX AI CHAT
', unsafe_allow_html=True)
st.markdown('ENTER THE MATRIX: Advanced AI with Retrieval-Augmented Generation
', 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('', 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('
', unsafe_allow_html=True)
# Token status
if HF_TOKEN:
st.markdown('✅ HUGGING FACE TOKEN: CONNECTED
', unsafe_allow_html=True)
# Test connection
client, config = get_model_client(selected_model)
if client:
st.markdown(f'✅ {selected_model}: READY
', unsafe_allow_html=True)
else:
st.markdown(f'❌ {selected_model}: CONNECTION FAILED
', unsafe_allow_html=True)
else:
st.markdown('❌ NO HUGGING FACE TOKEN FOUND
', 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()