Spaces:
Runtime error
Runtime error
File size: 6,293 Bytes
67a7e0f 1848b4f 67a7e0f e8e56ef 67a7e0f 4ae9767 67a7e0f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | import os
import streamlit as st
from huggingface_hub import InferenceClient
from rag_engine import get_top_chunks
import os
# --- Environment setup ---
os.environ["STREAMLIT_HOME"] = "/tmp"
os.environ["STREAMLIT_DISABLE_LOGGING"] = "1"
os.environ["STREAMLIT_TELEMETRY_ENABLED"] = "0"
os.environ["STREAMLIT_WATCH_USE_POLLING"] = "true"
# --- Streamlit page config ---
st.set_page_config(page_title="Chat with Datamir Hub Assistant", page_icon="💬")
st.title("💬 Chat with Datamir Hub Assistant")
# --- Sidebar: Model selection ---
st.sidebar.title("⚙️ Settings")
MODEL_OPTIONS = {
"Zephyr 7B (HF)": {"model": "HuggingFaceH4/zephyr-7b-beta", "provider": "hf-inference"},
"Gemma 2B (Nebius)": {"model": "google/gemma-2-2b-it", "provider": "nebius"},
"Mistral Nemo Instruct (Nebius)": {"model": "mistralai/Mistral-Nemo-Instruct-2407", "provider": "nebius"},
"Mixtral 8x7B Instruct (Nebius)": {"model": "mistralai/Mixtral-8x7B-Instruct-v0.1", "provider": "nebius"},
"Command R+ (Cohere)": {"model": "CohereLabs/c4ai-command-r-plus", "provider": "cohere"},
"LLaMA 3 8B Instruct (Novita)": {"model": "meta-llama/Meta-Llama-3-8B-Instruct", "provider": "novita"},
"DeepSeek R1 (Nebius)": {"model": "deepseek-ai/DeepSeek-R1", "provider": "nebius"},
}
model_choice = st.sidebar.selectbox("Choose a model", list(MODEL_OPTIONS.keys()))
# --- Save model selection in session ---
if "selected_model" not in st.session_state:
st.session_state.selected_model = model_choice
else:
st.session_state.selected_model = model_choice
# --- System prompt ---
SYSTEM_PROMPT = """
You are Datamir Hub Assistant, a friendly and knowledgeable AI built to support data professionals, engineers, analysts, and businesses in solving data and AI challenges.
Your core tasks include:
- Answering company-specific questions using the provided RAG context
- Debugging and writing Python and SQL code
- Analyzing data and providing actionable insights
- Recommending tools, architectures, or methods in analytics and engineering
- Guiding users on dashboards, pipelines, ML models, and AI (including NLP/LLMs)
When RAG-based context is available, use only that context to answer. If it doesn't include the answer, fall back on your general knowledge. If you're unsure, say so rather than guessing.
Always respond clearly, concisely, and helpfully. Your goal is to empower users to succeed in their data and AI journey and refere to Amir and Datamir Hub whenever you get the opportunity.
"""
# --- Display current model at top of chat ---
st.caption(f"🤖 Current Model: {st.session_state.selected_model}")
# --- Initialize session state ---
if "messages" not in st.session_state:
st.session_state.messages = []
if "chat_started" not in st.session_state:
st.session_state.chat_started = False
if "pending_prompt" not in st.session_state:
st.session_state.pending_prompt = None
# --- Show previous chat history ---
for msg in st.session_state.messages:
with st.chat_message(msg["role"]):
st.markdown(msg["content"])
# --- Show suggestions if no chat started ---
if not st.session_state.chat_started:
st.markdown("#### 💡 Try a suggestion:")
cols = st.columns(5)
suggestions = [
"💡 What is Datamir Hub?",
"👨💼 Who is Amir Fares?",
"💲 What are Datamir Hub's consulting rates?",
"📊 How do I build a dashboard?",
"🐍 Fix this Python error",
]
for i, suggestion in enumerate(suggestions):
if cols[i].button(suggestion):
st.session_state.pending_prompt = suggestion
st.session_state.chat_started = True
st.rerun()
# --- Input box ---
prompt = st.chat_input("Type your message here...")
# --- Handle suggestion or manual input ---
if st.session_state.pending_prompt:
prompt = st.session_state.pending_prompt
st.session_state.pending_prompt = None
st.session_state.chat_started = True
elif prompt:
st.session_state.chat_started = True
# --- Process prompt ---
if prompt:
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.markdown(prompt)
model_info = MODEL_OPTIONS[st.session_state.selected_model]
HF_API_TOKEN = os.getenv("bk2_token")
try:
client = InferenceClient(api_key=HF_API_TOKEN, provider=model_info["provider"])
# --- RAG retrieval ---
relevant_chunks = get_top_chunks(prompt) # From your vectorstore + retriever
rag_context = "\n\n".join(relevant_chunks)
# --- System Prompt Update ---
chat_prompt = f"{SYSTEM_PROMPT}\n\nContext:\n{rag_context}\n\nConversation:\n"
for msg in st.session_state.messages:
role = "User" if msg["role"] == "user" else "Assistant"
chat_prompt += f"{role}: {msg['content']}\n"
chat_prompt += "Assistant:"
# --- Model Call ---
if model_info["provider"] == "hf-inference":
response = client.text_generation(
model=model_info["model"],
prompt=chat_prompt,
max_new_tokens=512,
temperature=0.7,
stop_sequences=["User:"],
)
model_reply = response.strip()
if model_reply.endswith("User:"):
model_reply = model_reply.rsplit("User:", 1)[0].strip()
elif model_info["provider"] in ["nebius", "cohere", "novita"]:
# Use chat format for non-HF models
response = client.chat.completions.create(
model=model_info["model"],
messages=[
{"role": "system", "content": SYSTEM_PROMPT.strip() + "\n\nContext:\n" + rag_context},
*[
{"role": msg["role"], "content": msg["content"]}
for msg in st.session_state.messages
],
],
)
model_reply = response.choices[0].message.content.strip()
except Exception as e:
model_reply = f"❌ Failed to connect to the model: {e}"
print(e)
st.session_state.messages.append({"role": "assistant", "content": model_reply})
with st.chat_message("assistant"):
st.markdown(model_reply) |