Added multi model feature

src/streamlit_app.py

@@ -7,11 +7,30 @@ os.environ["STREAMLIT_HOME"] = "/tmp"
 os.environ["STREAMLIT_DISABLE_LOGGING"] = "1"
 os.environ["STREAMLIT_TELEMETRY_ENABLED"] = "0"
 
-
 # --- 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 assistant created to support data professionals, engineers, analysts, and businesses with their data and AI needs.
@@ -42,7 +61,8 @@ The standard consulting/freelance rate is $25/hour.
 Your goal is to ensure every user feels supported, empowered, and confident in their data and AI journey.
 """
 
-
+# --- 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:
@@ -71,10 +91,9 @@ if not st.session_state.chat_started:
     for i, suggestion in enumerate(suggestions):
         if cols[i].button(suggestion):
             st.session_state.pending_prompt = suggestion
-            st.session_state.chat_started = True  # <-- Fix: hide suggestions immediately
+            st.session_state.chat_started = True
             st.rerun()
 
-
 # --- Input box ---
 prompt = st.chat_input("Type your message here...")
 
@@ -86,34 +105,51 @@ if st.session_state.pending_prompt:
 elif prompt:
     st.session_state.chat_started = True
 
-# --- Process prompt (either from suggestion or typed input) ---
+# --- Process prompt ---
 if prompt:
     st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user"):
         st.markdown(prompt)
 
-    # --- Build chat prompt for model ---
-    chat_prompt = SYSTEM_PROMPT.strip() + "\n\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_info = MODEL_OPTIONS[st.session_state.selected_model]
+    HF_API_TOKEN = os.getenv("bk_token")
 
-    # --- Call Hugging Face model ---
     try:
-        # HF_API_TOKEN = os.getenv("hf_general_token")
-        HF_API_TOKEN = os.getenv("bk_token")
-        client = InferenceClient(api_key=HF_API_TOKEN, provider="hf-inference")
-        response = client.text_generation(
-            model="HuggingFaceH4/zephyr-7b-beta",
-            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()
+        client = InferenceClient(api_key=HF_API_TOKEN, provider=model_info["provider"])
+
+        if model_info["provider"] == "hf-inference":
+            # Format prompt as plain text conversation
+            chat_prompt = SYSTEM_PROMPT.strip() + "\n\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:"
+
+            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 Nebius
+            response = client.chat.completions.create(
+                model=model_info["model"],
+                messages=[
+                    {"role": "system", "content": SYSTEM_PROMPT.strip()},
+                    *[
+                        {"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)