Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,112 +2,108 @@ import os
 import streamlit as st
 from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
 from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
-from io import StringIO
 
-# Set API tokens
+# API key setup
 os.environ['HUGGINGFACEHUB_API_TOKEN'] = os.getenv("key")
 os.environ['HF_TOKEN'] = os.getenv("key")
 
-# Page setup
-st.set_page_config(page_title="🧠 Multi-Mentor Chat", page_icon="🧠")
-st.title("👨‍🏫 Multi-Topic Mentor Chatbot")
-
-st.markdown(":red[Select a topic area to get personalized mentoring:]")
-
-# Define mentors and models
-MENTOR_MODELS = {
-    "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 & Power BI": {
-        "repo_id": "meta-llama/Meta-Llama-3-70B-Instruct",
-        "provider": "hyperbolic"
-    }
+# Streamlit config
+st.set_page_config(page_title="🧠 Multi-Mentor Chat", 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")
 }
 
-# Tabs for each mentor
-tabs = st.tabs(list(MENTOR_MODELS.keys()))
-
-# Initialize conversation storage
-if "conversations" not in st.session_state:
-    st.session_state.conversations = {key: [] for key in MENTOR_MODELS.keys()}
-
-# Iterate through each mentor tab
-for tab, (mentor_name, config) in zip(tabs, MENTOR_MODELS.items()):
-    with tab:
-        st.subheader(f"🧠 {mentor_name} Mentor")
-        experience = st.slider(f"Your experience in {mentor_name} (years):", 0, 20, 1, key=mentor_name)
-        question = st.text_input(f"Ask your {mentor_name} question:", key=f"q_{mentor_name}")
-        output = st.empty()
-
-        # Load model
-        model = HuggingFaceEndpoint(
-            repo_id=config["repo_id"],
-            provider=config["provider"],
-            temperature=0.5,
-            max_new_tokens=150,
-            task="conversational"
-        )
-
-        chat_model = ChatHuggingFace(
-            llm=model,
-            repo_id=config["repo_id"],
-            provider=config["provider"],
-            temperature=0.5,
-            max_new_tokens=150,
-            task="conversational"
-        )
-
-        # Handle question and response
-        if st.button("Ask", key=f"ask_{mentor_name}") and question:
+# Display custom styled buttons in columns
+st.markdown("### 🧑‍🏫 Choose a Mentor Domain")
+cols = st.columns(len(MENTORS))
+mentor_keys = list(MENTORS.keys())
+selected = None
+
+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
+
+# Use selection from session state if exists
+if "selected_mentor" in st.session_state:
+    selected = st.session_state["selected_mentor"]
+
+if selected:
+    mentor_name, repo_id, provider = MENTORS[selected]
+    st.subheader(f"{mentor_name} Mentor 💬")
+
+    # User experience slider and question input
+    experience = st.slider("📅 Your Experience (Years):", 0, 20, 1)
+    user_input = st.text_input("📝 Ask your question:")
+
+    output_container = st.container()
+
+    # Load the model
+    model = HuggingFaceEndpoint(
+        repo_id=repo_id,
+        provider=provider,
+        temperature=0.5,
+        max_new_tokens=200,
+        task="conversational"
+    )
+    chat_model = ChatHuggingFace(
+        llm=model,
+        repo_id=repo_id,
+        provider=provider,
+        temperature=0.5,
+        max_new_tokens=200,
+        task="conversational"
+    )
+
+    if "history" not in st.session_state:
+        st.session_state["history"] = []
+
+    # Ask button
+    if st.button("🧠 Ask Mentor"):
+        if user_input:
             prompt = ChatPromptTemplate.from_messages([
                 SystemMessagePromptTemplate.from_template(
-                    f"You are a helpful and experienced {mentor_name} mentor. The user has {experience} years of experience. Answer the question appropriately."
+                    f"You are a helpful and experienced {mentor_name} mentor assisting a learner with {experience} years of experience."
                 ),
                 HumanMessagePromptTemplate.from_template("{question}")
             ])
-            formatted_prompt = prompt.format_messages(question=question)
+            formatted_prompt = prompt.format_messages(question=user_input)
 
-            with st.spinner("Mentor is thinking..."):
+            with st.spinner("Thinking..."):
                 response = chat_model.invoke(formatted_prompt)
 
-            answer = response.content
-            output.markdown(f"👤 You: {question}")
-            output.markdown(f"🧠 Mentor: {answer}")
-
-            # Save conversation
-            st.session_state.conversations[mentor_name].append(f"You: {question}")
-            st.session_state.conversations[mentor_name].append(f"Mentor: {answer}")
+            # Save and show conversation
+            st.session_state["history"].append(f"👤 You: {user_input}")
+            st.session_state["history"].append(f"🧠 Mentor: {response.content}")
 
-        # Clear output
-        if st.button("Clear", key=f"clear_{mentor_name}"):
-            output.empty()
-            st.session_state.conversations[mentor_name] = []
+    # Display conversation history
+    if st.session_state["history"]:
+        with output_container:
+            st.markdown("### 🗨️ Conversation")
+            for msg in st.session_state["history"]:
+                st.markdown(msg)
 
-        # Download conversation
-        if st.session_state.conversations[mentor_name]:
-            convo_text = "\n".join(st.session_state.conversations[mentor_name])
+    # 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 Conversation",
-                data=convo_text,  # ← Fixed: Use string directly
-                file_name=f"{mentor_name.lower().replace(' ', '_')}_conversation.txt",
-                mime="text/plain",
-                key=f"download_{mentor_name}"
+                label="⬇️ Download Chat",
+                data="\n".join(st.session_state["history"]),
+                file_name=f"{selected}_mentor_chat.txt",
+                mime="text/plain"
             )