Update tribunal_2.py

tribunal_2.py

@@ -85,6 +85,17 @@ three_column_style = """
     </style>
 """
 
+#-----
+def get_active_model():
+    return genparam.SELECTED_MODEL_1 if genparam.ACTIVE_MODEL == 0 else genparam.SELECTED_MODEL_2
+
+def get_active_prompt_template():
+    return genparam.PROMPT_TEMPLATE_1 if genparam.ACTIVE_MODEL == 0 else genparam.PROMPT_TEMPLATE_2
+
+def get_active_vector_index():
+    return st.secrets["vector_index_id_1"] if genparam.ACTIVE_INDEX == 0 else st.secrets["vector_index_id_2"]
+#-----
+
 def setup_client(project_id):
     credentials = Credentials(
         url=st.secrets["url"],
@@ -197,7 +208,7 @@ def fetch_response(user_input, milvus_client, emb, vector_index_properties, vect
     prompt_data = apply_prompt_syntax(
         prompt,
         system_prompt,
-        genparam.PROMPT_TEMPLATE,
+        get_active_prompt_template(),
         genparam.BAKE_IN_PROMPT_SYNTAX
     )
 
@@ -205,7 +216,7 @@ def fetch_response(user_input, milvus_client, emb, vector_index_properties, vect
     
     watsonx_llm = ModelInference(
         api_client=client, 
-        model_id=genparam.SELECTED_MODEL,
+        model_id=get_active_model(),
         verify=genparam.VERIFY
     )
 
@@ -217,7 +228,19 @@ def fetch_response(user_input, milvus_client, emb, vector_index_properties, vect
         GenParams.STOP_SEQUENCES: genparam.STOP_SEQUENCES
     }
 
-    with st.chat_message("Tribunal", avatar="🥸"):
+    bot_name = None
+    bot_avatar = None
+    if chat_history == st.session_state.chat_history_1:
+        bot_name = genparam.BOT_1_NAME
+        bot_avatar = genparam.BOT_1_AVATAR
+    elif chat_history == st.session_state.chat_history_2:
+        bot_name = genparam.BOT_2_NAME
+        bot_avatar = genparam.BOT_2_AVATAR
+    else:
+        bot_name = genparam.BOT_3_NAME
+        bot_avatar = genparam.BOT_3_AVATAR
+
+    with st.chat_message(bot_name, avatar=bot_avatar):
         if genparam.TOKEN_CAPTURE_ENABLED:
             st.code(prompt_data, line_numbers=True, wrap_lines=True)
         stream = generate_response(watsonx_llm, prompt_data, params)
@@ -278,17 +301,18 @@ def main():
             st.subheader(genparam.BOT_1_NAME)
             # Display chat history for bot 1
             for message in st.session_state.chat_history_1:
-                with st.chat_message(message["role"], avatar="👤" if message["role"] == "user" else "🥸"):
+                with st.chat_message(message["role"], avatar=genparam.USER_AVATAR if message["role"] == "user" else genparam.BOT_1_AVATAR):
                     #st.markdown(f"<span style='color: #1565C0;'>{message['content']}</span>", unsafe_allow_html=True)
                     st.markdown(message['content'])
             
             # Add user message and get bot 1 response
-            st.session_state.chat_history_1.append({"role": "user", "content": user_input, "avatar":"👤"})
+            st.session_state.chat_history_1.append({"role": "user", "content": user_input, "avatar": genparam.USER_AVATAR})
             milvus_client, emb, vector_index_properties, vector_store_schema = setup_vector_index(
-                    client, 
-                    wml_credentials, 
-                    st.secrets["vector_index_id"]
-                )
+                client, 
+                wml_credentials,
+                get_active_vector_index()
+               #st.secrets["vector_index_id"]
+            )
             system_prompt = genparam.BOT_1_PROMPT
             
             response = fetch_response(
@@ -300,7 +324,7 @@ def main():
                 system_prompt,
                 st.session_state.chat_history_1
             )
-            st.session_state.chat_history_1.append({"role": "Tribunal", "content": response, "avatar":"🥸"})
+            st.session_state.chat_history_1.append({"role": genparam.BOT_1_NAME, "content": response, "avatar": genparam.BOT_1_AVATAR})
             st.markdown("</div>", unsafe_allow_html=True)
             
         with col2:
@@ -308,18 +332,20 @@ def main():
             st.subheader(genparam.BOT_2_NAME)
             # Display chat history for bot 2
             for message in st.session_state.chat_history_2:
-                with st.chat_message(message["role"], avatar="👤" if message["role"] == "user" else "🥸"):
+                with st.chat_message(message["role"], avatar=genparam.USER_AVATAR if message["role"] == "user" else genparam.BOT_2_AVATAR):
                     #st.markdown(f"<span style='color: #2E7D32;'>{message['content']}</span>", unsafe_allow_html=True)
                     st.markdown(message['content'])
             
             
             # Add user message and get bot 2 response
-            st.session_state.chat_history_2.append({"role": "user", "content": user_input, "avatar":"👤"})
+            st.session_state.chat_history_2.append({"role": "user", "content": user_input, "avatar": genparam.USER_AVATAR})
             milvus_client, emb, vector_index_properties, vector_store_schema = setup_vector_index(
                 client, 
-                wml_credentials, 
-                st.secrets["vector_index_id"]
+                wml_credentials,
+                get_active_vector_index()
+               #st.secrets["vector_index_id"]
             )
+            system_prompt = genparam.BOT_2_PROMPT
             
             response = fetch_response(
                 user_input, 
@@ -327,10 +353,10 @@ def main():
                 emb, 
                 vector_index_properties, 
                 vector_store_schema,
-                genparam.BOT_2_PROMPT,
+                system_prompt,
                 st.session_state.chat_history_2
             )
-            st.session_state.chat_history_2.append({"role": "Tribunal", "content": response, "avatar":"🥸"})
+            st.session_state.chat_history_2.append({"role": genparam.BOT_2_NAME, "content": response, "avatar": genparam.BOT_2_AVATAR})
             st.markdown("</div>", unsafe_allow_html=True)
             
         with col3:
@@ -338,18 +364,20 @@ def main():
             st.subheader(genparam.BOT_3_NAME)
             # Display chat history for bot 3
             for message in st.session_state.chat_history_3:
-                with st.chat_message(message["role"], avatar="👤" if message["role"] == "user" else "🥸"):
+                with st.chat_message(message["role"], avatar=genparam.USER_AVATAR if message["role"] == "user" else genparam.BOT_3_AVATAR):
                     #st.markdown(f"<span style='color: #6A1B9A;'>{message['content']}</span>", unsafe_allow_html=True)
                     st.markdown(message['content'])
             
             
             # Add user message and get bot 3 response
-            st.session_state.chat_history_3.append({"role": "user", "content": user_input, "avatar":"👤"})
+            st.session_state.chat_history_3.append({"role": "user", "content": user_input, "avatar": genparam.USER_AVATAR})
             milvus_client, emb, vector_index_properties, vector_store_schema = setup_vector_index(
                 client, 
-                wml_credentials, 
-                st.secrets["vector_index_id"]
+                wml_credentials,
+                get_active_vector_index()
+               #st.secrets["vector_index_id"]
             )
+            system_prompt = genparam.BOT_3_PROMPT
             
             response = fetch_response(
                 user_input, 
@@ -357,10 +385,10 @@ def main():
                 emb, 
                 vector_index_properties, 
                 vector_store_schema,
-                genparam.BOT_3_PROMPT,
+                system_prompt,
                 st.session_state.chat_history_3
             )
-            st.session_state.chat_history_3.append({"role": "Tribunal", "content": response, "avatar":"🥸"})
+            st.session_state.chat_history_3.append({"role": genparam.BOT_3_NAME, "content": response, "avatar": genparam.BOT_3_AVATAR})
             st.markdown("</div>", unsafe_allow_html=True)
 
 if __name__ == "__main__":