src/streamlit_app.py

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)