Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,101 +3,104 @@ import streamlit as st
 from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
 from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
 
-# API key setup
+# Load Hugging Face token
 os.environ['HUGGINGFACEHUB_API_TOKEN'] = os.getenv("key")
 os.environ['HF_TOKEN'] = os.getenv("key")
 
-# Streamlit config
-st.set_page_config(page_title="🧠 Multi-Domain Knowledge Assistant", page_icon="🎓", layout="wide")
-st.title("🎓 Multi-Topic Mentor Assistant")
-st.markdown("Get expert help in different data & tech domains. Choose a mentor below to get started:")
-
-# Mentor options with emojis and display names
-MENTORS = {
-    "python": ("🐍 Python", "meta-llama/Llama-3.1-8B-Instruct", "nebius"),
-    "machine_learning": ("🤖 ML", "deepseek-ai/DeepSeek-R1", "nebius"),
-    "deep_learning": ("🧠 DL", "deepseek-ai/DeepSeek-R1", "nebius"),
-    "stats": ("📊 Stats", "meta-llama/Llama-3.2-1B-Instruct", "nebius"),
-    "data_analysis": ("📈 Data Analysis", "meta-llama/Llama-3.3-70B-Instruct", "nebius"),
-    "sql_powerbi": ("🗄️ SQL + Power BI", "meta-llama/Meta-Llama-3-70B-Instruct", "hyperbolic")
+# Page config
+st.set_page_config(page_title="🧠 AI Mentor Hub", page_icon="🎓")
+st.title("🎓 AI Mentor Hub - Learn Smarter, Faster!")
+
+st.markdown("### 🔍 Select your mentor and ask any question:")
+
+# Emoji mentor labels
+mentor_labels = {
+    "python": "🐍 Python",
+    "machine_learning": "🤖 Machine Learning",
+    "deep_learning": "🧠 Deep Learning",
+    "stats": "📊 Statistics",
+    "data_analysis": "🧾 Data Analysis",
+    "sql_powerbi": "🗃️ SQL & Power BI"
 }
 
-# Display custom styled buttons in columns
-st.markdown("### 🧑‍🏫 Choose a Mentor Domain")
-cols = st.columns(len(MENTORS))
-mentor_keys = list(MENTORS.keys())
-selected = None
+mentor_configs = {
+    "python": {"repo_id": "meta-llama/Llama-3.1-8B-Instruct", "provider": "nebius"},
+    "machine_learning": {"repo_id": "deepseek-ai/DeepSeek-R1", "provider": "nebius"},
+    "deep_learning": {"repo_id": "deepseek-ai/DeepSeek-R1", "provider": "nebius"},
+    "stats": {"repo_id": "meta-llama/Llama-3.2-1B-Instruct", "provider": "nebius"},
+    "data_analysis": {"repo_id": "meta-llama/Llama-3.3-70B-Instruct", "provider": "nebius"},
+    "sql_powerbi": {"repo_id": "meta-llama/Meta-Llama-3-70B-Instruct", "provider": "hyperbolic"}
+}
 
-for i, col in enumerate(cols):
-    with col:
-        if st.button(MENTORS[mentor_keys[i]][0]):
-            selected = mentor_keys[i]
-            st.session_state["selected_mentor"] = selected
+# Show options
+mentor_choice = st.selectbox("Choose a mentor topic:", [""] + list(mentor_labels.keys()))
 
-# Use selection from session state if exists
-if "selected_mentor" in st.session_state:
-    selected = st.session_state["selected_mentor"]
+# Storage for conversation
+if "conversation" not in st.session_state:
+    st.session_state.conversation = []
 
-if selected:
-    mentor_name, repo_id, provider = MENTORS[selected]
-    st.subheader(f"{mentor_name} Mentor 💬")
+if mentor_choice:
+    label = mentor_labels[mentor_choice]
+    st.subheader(f"{label} Mentor Chat")
 
-    # User experience slider and question input
-    experience = st.slider("📅 Your Experience (Years):", 0, 20, 1)
-    user_input = st.text_input("📝 Ask your question:")
+    # Inputs
+    experience = st.slider("📅 Your experience (years):", 0, 20, 1)
+    question = st.text_input("💬 Ask your question:")
 
-    output_container = st.container()
+    # Model setup
+    config = mentor_configs[mentor_choice]
+    repo_id = config["repo_id"]
+    provider = config["provider"]
 
-    # Load the model
     model = HuggingFaceEndpoint(
         repo_id=repo_id,
         provider=provider,
         temperature=0.5,
-        max_new_tokens=200,
+        max_new_tokens=150,
         task="conversational"
     )
+
     chat_model = ChatHuggingFace(
         llm=model,
         repo_id=repo_id,
-        provider=provider,
+        provider=provider,  # FIXED: use correct provider dynamically
         temperature=0.5,
-        max_new_tokens=200,
+        max_new_tokens=150,
         task="conversational"
     )
 
-    if "history" not in st.session_state:
-        st.session_state["history"] = []
+    output_box = st.empty()
 
-    # Ask button
-    if st.button("🧠 Ask Mentor"):
-        if user_input:
+    if st.button("🧠 Get Answer"):
+        if not question.strip():
+            st.warning("❗ Please enter a question.")
+        else:
             prompt = ChatPromptTemplate.from_messages([
                 SystemMessagePromptTemplate.from_template(
-                    f"You are a helpful and experienced {mentor_name} mentor assisting a learner with {experience} years of experience."
+                    f"You are a helpful and expert {mentor_choice.replace('_', ' ').title()} mentor. The user has {experience} years of experience. Answer clearly."
                 ),
                 HumanMessagePromptTemplate.from_template("{question}")
             ])
-            formatted_prompt = prompt.format_messages(question=user_input)
+            messages = prompt.format_messages(question=question)
 
             with st.spinner("Thinking..."):
-                response = chat_model.invoke(formatted_prompt)
-
-            # Save and show conversation
-            st.session_state["history"].append(f"👤 You: {user_input}")
-            st.session_state["history"].append(f"🧠 Mentor: {response.content}")
-
-
-    # Clear and download buttons
-    col_clear, col_download = st.columns(2)
-    with col_clear:
-        if st.button("🧹 Clear Conversation"):
-            st.session_state["history"] = []
-            output_container.empty()
-    with col_download:
-        if st.session_state["history"]:
-            st.download_button(
-                label="⬇️ Download Chat",
-                data="\n".join(st.session_state["history"]),
-                file_name=f"{selected}_mentor_chat.txt",
-                mime="text/plain"
-            )
+                response = chat_model.invoke(messages)
+
+            answer = response.content
+            output_box.markdown(f"👤 **You:** {question}")
+            output_box.markdown(f"🧠 **Mentor:** {answer}")
+            st.session_state.conversation.append(f"You: {question}")
+            st.session_state.conversation.append(f"Mentor: {answer}")
+
+    if st.button("🗑️ Clear Chat"):
+        output_box.empty()
+        st.session_state.conversation = []
+
+    if st.session_state.conversation:
+        convo_text = "\n".join(st.session_state.conversation)
+        st.download_button(
+            "⬇️ Download Conversation",
+            data=convo_text,
+            file_name=f"{mentor_choice}_chat.txt",
+            mime="text/plain"
+        )