Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,70 +2,111 @@ 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
 os.environ['HUGGINGFACEHUB_API_TOKEN'] = os.getenv("key")
 os.environ['HF_TOKEN'] = os.getenv("key")
 
-st.set_page_config(page_title="👨‍🏫 Multi-Mentor Chat", page_icon="🧠")
-st.title("🧠 Multi-Topic Mentor")
+# Page setup
+st.set_page_config(page_title="🧠 Multi-Mentor Chat", page_icon="🧠")
+st.title("👨‍🏫 Multi-Topic Mentor Chatbot")
 
+# 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"
+    }
+}
 
-st.markdown("### Select your Mentor:")
-col1, col2, col3, col4, col5, col6 = st.columns(6)
-with col1: st.write("🟢 Python")
-with col2: st.write("🟠 ML")
-with col3: st.write("🔵 DL")
-with col4: st.write("🟣 Stats")
-with col5: st.write("🟡 Data_Anaylasis")
-with col6: st.write("🔴 sql and powerbi")
+# Tabs for each mentor
+tabs = st.tabs(list(MENTOR_MODELS.keys()))
 
-mentor_type = st.selectbox("Choose a mentor:", ["", "python", "machine_learning", "deep_learning", "stats", "data_anaylasis", "sql and powerbi"])
+# Initialize conversation storage
+if "conversations" not in st.session_state:
+    st.session_state.conversations = {key: [] for key in MENTOR_MODELS.keys()}
 
-if mentor_type:
-    st.subheader(f"🧠 {mentor_type.upper()} Mentor Chat")
-    experience = st.slider("Your experience (in years):", 0, 20, 1)
-    user_input = st.text_input("Ask your question:")
+# 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"
+        )
 
-    output_container = st.empty()
+        chat_model = ChatHuggingFace(
+            llm=model,
+            repo_id=config["repo_id"],
+            provider=config["provider"],
+            temperature=0.5,
+            max_new_tokens=150,
+            task="conversational"
+        )
 
-    if mentor_type == "python":
-        model = HuggingFaceEndpoint(repo_id="meta-llama/Llama-3.1-8B-Instruct", provider="nebius", temperature=0.5, max_new_tokens=150, task="conversational")
-    elif mentor_type == "machine_learning":
-        model = HuggingFaceEndpoint(repo_id="deepseek-ai/DeepSeek-R1", provider="nebius", temperature=0.5, max_new_tokens=150, task="conversational")
-    elif mentor_type == "deep_learning":
-        model = HuggingFaceEndpoint(repo_id="deepseek-ai/DeepSeek-R1", provider="nebius", temperature=0.5, max_new_tokens=150, task="conversational")
-    elif mentor_type == "stats":
-        model = HuggingFaceEndpoint(repo_id="meta-llama/Llama-3.2-1B-Instruct", provider="nebius", temperature=0.5, max_new_tokens=150, task="conversational")
-    elif mentor_type == "data_anaylasis":
-        model = HuggingFaceEndpoint(repo_id="meta-llama/Llama-3.3-70B-Instruct", provider="nebius", temperature=0.5, max_new_tokens=150, task="conversational")
-    elif mentor_type == "sql and powerbi":
-        model = HuggingFaceEndpoint(repo_id="meta-llama/Meta-Llama-3-70B-Instruct", provider="hyperbolic", temperature=0.5, max_new_tokens=150, task="conversational")
+        # Handle question and response
+        if st.button("Ask", key=f"ask_{mentor_name}") and question:
+            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."
+                ),
+                HumanMessagePromptTemplate.from_template("{question}")
+            ])
+            formatted_prompt = prompt.format_messages(question=question)
 
-    chat_model = ChatHuggingFace(
-        llm=model,
-        repo_id=model.repo_id,
-        provider="nebius",
-        temperature=0.5,
-        max_new_tokens=150,
-        task="conversational"
-    )
+            with st.spinner("Mentor is thinking..."):
+                response = chat_model.invoke(formatted_prompt)
 
+            answer = response.content
+            output.markdown(f"👤 You: {question}")
+            output.markdown(f"🧠 Mentor: {answer}")
 
-    if st.button("Ask") and user_input:
-        prompt = ChatPromptTemplate.from_messages([
-            SystemMessagePromptTemplate.from_template(
-                f"You are a helpful and experienced {mentor_type.upper()} mentor assisting a learner with {experience} years of experience."
-            ),
-            HumanMessagePromptTemplate.from_template("{question}")
-        ])
-        formatted_prompt = prompt.format_messages(question=user_input)
+            # Save conversation
+            st.session_state.conversations[mentor_name].append(f"You: {question}")
+            st.session_state.conversations[mentor_name].append(f"Mentor: {answer}")
 
-        with st.spinner("Mentor is thinking..."):
-            response = chat_model.invoke(formatted_prompt)
+        # Clear output
+        if st.button("Clear", key=f"clear_{mentor_name}"):
+            output.empty()
+            st.session_state.conversations[mentor_name] = []
 
-        output_container.markdown(f"👤 You: {user_input}")
-        output_container.markdown(f"🧠 Mentor: {response.content}")
-
-    if st.button("Clear Output"):
-        output_container.empty()
+        # Download conversation
+        if st.session_state.conversations[mentor_name]:
+            convo_text = "\n".join(st.session_state.conversations[mentor_name])
+            buffer = StringIO(convo_text)
+            st.download_button(
+                label="⬇️ Download Conversation",
+                data=buffer,
+                file_name=f"{mentor_name.lower().replace(' ', '_')}_conversation.txt",
+                mime="text/plain",
+                key=f"download_{mentor_name}"
+            )