Update neo_sages.py

neo_sages.py

@@ -217,6 +217,7 @@ def generate_response(watsonx_llm, prompt_data, params):
         yield chunk
 
 def fetch_response(user_input, milvus_client, emb, vector_index_properties, vector_store_schema, system_prompt, chat_history):
+    # Get grounding documents
     grounding = proximity_search(
         question=user_input,
         milvus_client=milvus_client,
@@ -224,17 +225,54 @@ def fetch_response(user_input, milvus_client, emb, vector_index_properties, vect
         vector_index_properties=vector_index_properties,
         vector_store_schema=vector_store_schema
     )
-    prompt = prepare_prompt(user_input, chat_history)
 
-    prompt_data = apply_prompt_syntax(
-        prompt,
-        system_prompt,
-        get_active_prompt_template(),
-        genparam.BAKE_IN_PROMPT_SYNTAX
-    )
-
-    prompt_data = prompt_data.replace("__grounding__", grounding)
-    
+    # Special handling for PATH-er B. (first column)
+    if chat_history == st.session_state.chat_history_1:
+        # Display user question first
+        with st.chat_message("user", avatar=genparam.USER_AVATAR):
+            st.markdown(user_input)
+        
+        # Parse and display each document from the grounding
+        documents = grounding.split("\n\n")[2:]  # Skip the count line and first newline
+        for doc in documents:
+            if doc.strip():  # Only process non-empty strings
+                parts = doc.split("\n")
+                doc_name = parts[0].replace("Document: ", "")
+                content = parts[1].replace("Content: ", "")
+                
+                # Display document with delay
+                time.sleep(0.5)
+                st.markdown(f"**{doc_name}**")
+                st.code(content)
+
+        # Store in chat history
+        return grounding
+
+    # For MOD-ther S. (second column)
+    elif chat_history == st.session_state.chat_history_2:
+        prompt = prepare_prompt(user_input, chat_history)
+        prompt_data = apply_prompt_syntax(
+            prompt,
+            system_prompt,
+            get_active_prompt_template(),
+            genparam.BAKE_IN_PROMPT_SYNTAX
+        )
+        prompt_data = prompt_data.replace("__grounding__", grounding)
+
+    # For SYS-ter V. (third column)
+    else:
+        # Get chat history from MOD-ther S.
+        mod_ther_history = st.session_state.chat_history_2
+        prompt = prepare_prompt(user_input, mod_ther_history)
+        prompt_data = apply_prompt_syntax(
+            prompt,
+            system_prompt,
+            get_active_prompt_template(),
+            genparam.BAKE_IN_PROMPT_SYNTAX
+        )
+        prompt_data = prompt_data.replace("__grounding__", grounding)
+
+    # Continue with normal processing for columns 2 and 3
     watsonx_llm = ModelInference(
         api_client=client, 
         model_id=get_active_model(),
@@ -264,9 +302,11 @@ def fetch_response(user_input, milvus_client, emb, vector_index_properties, vect
     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)
-        response = st.write_stream(stream)
-       # response = st.write_stream(stream, f"<span style='color: {color};'>", unsafe_allow_html=True)
+        if chat_history != st.session_state.chat_history_1:  # Only generate responses for columns 2 and 3
+            stream = generate_response(watsonx_llm, prompt_data, params)
+            response = st.write_stream(stream)
+        else:
+            response = grounding  # For column 1, we already displayed the content
 
     if genparam.TOKEN_CAPTURE_ENABLED:
         chat_number = len(chat_history) // 2
@@ -324,23 +364,35 @@ def main():
         # Create three columns
         col1, col2, col3 = st.columns(3)
         
+###-------------START
+        # First column - PATH-er B. (Document Display)
         with col1:
             st.markdown("<div class='chat-container'>", unsafe_allow_html=True)
             st.subheader(f"{genparam.BOT_1_AVATAR} {genparam.BOT_1_NAME}")
             st.markdown("<div class='chat-messages'>", unsafe_allow_html=True)
             
-            # Display only bot responses
+            # Display previous messages
             for message in st.session_state.chat_history_1:
-                if message["role"] != "user":  # Only show bot messages
-                    with st.chat_message(message["role"], avatar=genparam.BOT_1_AVATAR):
+                if message["role"] == "user":
+                    with st.chat_message(message["role"], avatar=genparam.USER_AVATAR):
                         st.markdown(message['content'])
+                else:
+                    # Parse and display stored documents
+                    documents = message['content'].split("\n\n")[2:]  # Skip count line
+                    for doc in documents:
+                        if doc.strip():
+                            parts = doc.split("\n")
+                            doc_name = parts[0].replace("Document: ", "")
+                            content = parts[1].replace("Content: ", "")
+                            st.markdown(f"**{doc_name}**")
+                            st.code(content)
             
-            # Add user message to history but don't display
+            # Add user message and get new response
             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,
-                get_active_vector_index()
+                st.secrets["vector_index_id_1"]  # Use first vector index
             )
             system_prompt = genparam.BOT_1_PROMPT
             
@@ -355,24 +407,23 @@ def main():
             )
             st.session_state.chat_history_1.append({"role": genparam.BOT_1_NAME, "content": response, "avatar": genparam.BOT_1_AVATAR})
             st.markdown("</div></div>", unsafe_allow_html=True)
-            
+
+        # Second column - MOD-ther S. (Uses documents from first vector index)
         with col2:
             st.markdown("<div class='chat-container'>", unsafe_allow_html=True)
             st.subheader(f"{genparam.BOT_2_AVATAR} {genparam.BOT_2_NAME}")
             st.markdown("<div class='chat-messages'>", unsafe_allow_html=True)
             
-            # Display only bot responses
             for message in st.session_state.chat_history_2:
-                if message["role"] != "user":  # Only show bot messages
+                if message["role"] != "user":
                     with st.chat_message(message["role"], avatar=genparam.BOT_2_AVATAR):
                         st.markdown(message['content'])
             
-            # Add user message to history but don't display
             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,
-                get_active_vector_index()
+                st.secrets["vector_index_id_1"]  # Use first vector index
             )
             system_prompt = genparam.BOT_2_PROMPT
             
@@ -388,23 +439,22 @@ def main():
             st.session_state.chat_history_2.append({"role": genparam.BOT_2_NAME, "content": response, "avatar": genparam.BOT_2_AVATAR})
             st.markdown("</div></div>", unsafe_allow_html=True)
             
+        # Third column - SYS-ter V. (Uses second vector index and chat history from second column)
         with col3:
             st.markdown("<div class='chat-container'>", unsafe_allow_html=True)
             st.subheader(f"{genparam.BOT_3_AVATAR} {genparam.BOT_3_NAME}")
             st.markdown("<div class='chat-messages'>", unsafe_allow_html=True)
             
-            # Display only bot responses
             for message in st.session_state.chat_history_3:
-                if message["role"] != "user":  # Only show bot messages
+                if message["role"] != "user":
                     with st.chat_message(message["role"], avatar=genparam.BOT_3_AVATAR):
                         st.markdown(message['content'])
             
-            # Add user message to history but don't display
             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_2"]
+                st.secrets["vector_index_id_2"]  # Use second vector index
             )
             system_prompt = genparam.BOT_3_PROMPT
             
@@ -419,11 +469,9 @@ def main():
             )
             st.session_state.chat_history_3.append({"role": genparam.BOT_3_NAME, "content": response, "avatar": genparam.BOT_3_AVATAR})
             st.markdown("</div></div>", unsafe_allow_html=True)
+###-------------END
 
         # Update sidebar with new question
-        st.sidebar.markdown("---")  # Add divider
+        st.sidebar.markdown("---")
         st.sidebar.markdown("**Latest Question:**")
-        st.sidebar.markdown(f"_{user_input}_")
-
-if __name__ == "__main__":
-    main()
+        st.sidebar.markdown(f"_{user_input}_")