Update neo_sages.py

neo_sages.py

@@ -64,8 +64,8 @@ def initialize_session_state():
         st.session_state.first_question = False 
     if "counter" not in st.session_state:
         st.session_state["counter"] = 0
-    if 'token_capture' not in st.session_state:
-        st.session_state.token_capture = []
+    if 'token_statistics' not in st.session_state:
+        st.session_state.token_statistics = []
 
 # three_column_style = """
 #     <style>
@@ -300,20 +300,18 @@ def fetch_response(user_input, milvus_client, emb, vector_index_properties, vect
         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)
         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)
+            
+            # Only capture tokens for MOD-ther S. and SYS-ter V.
+            if genparam.TOKEN_CAPTURE_ENABLED and chat_history != st.session_state.chat_history_1:
+                token_stats = capture_tokens(prompt_data, response, bot_name)
+                if token_stats:
+                    st.session_state.token_statistics.append(token_stats)
         else:
             response = grounding  # For column 1, we already displayed the content
 
-    if genparam.TOKEN_CAPTURE_ENABLED:
-        chat_number = len(chat_history) // 2
-        token_calculations = capture_tokens(prompt_data, response, chat_number)
-        if token_calculations:
-            st.sidebar.code(token_calculations)
-
     return response
 
 def capture_tokens(prompt_data, response, chat_number):
@@ -330,10 +328,13 @@ def capture_tokens(prompt_data, response, chat_number):
     output_tokens = watsonx_llm.tokenize(prompt=response)["result"]["token_count"]
     total_tokens = input_tokens + output_tokens
 
-    st.session_state.token_capture.append(f"chat {chat_number}: {input_tokens} + {output_tokens} = {total_tokens}")
-
-    token_calculations = "\n".join(st.session_state.token_capture)
-    return token_calculations
+    return {
+        "bot_name": bot_name,
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "total_tokens": total_tokens,
+        "timestamp": time.strftime("%H:%M:%S")
+    }
 
 def main():
     initialize_session_state()
@@ -343,16 +344,53 @@ def main():
     
     # Sidebar configuration
     st.sidebar.header('The Solutioning Sages')
-    st.sidebar.write('')
+    st.sidebar.divider()
     
-    # Display user chat history in sidebar
-    st.sidebar.subheader("Your Questions")
-    for i, message in enumerate(st.session_state.chat_history_1):
-        if message["role"] == "user":
-            st.sidebar.markdown(f"**Question {i//2 + 1}:** {message['content']}")
+    # Display token statistics in sidebar
+    st.sidebar.subheader("Token Usage Statistics")
     
-    st.sidebar.write('')
-    st.sidebar.write('')
+    # Group token statistics by interaction (for MOD-ther S. and SYS-ter V. only)
+    if st.session_state.token_statistics:
+        current_timestamp = None
+        interaction_count = 0
+        stats_by_time = {}
+        
+        # Group stats by timestamp
+        for stat in st.session_state.token_statistics:
+            if stat["timestamp"] not in stats_by_time:
+                stats_by_time[stat["timestamp"]] = []
+            stats_by_time[stat["timestamp"]].append(stat)
+        
+        # Display grouped stats
+        for timestamp, stats in stats_by_time.items():
+            interaction_count += 1
+            st.sidebar.markdown(f"**Interaction {interaction_count}** ({timestamp})")
+            
+            # Calculate total tokens for this interaction
+            total_input = sum(stat['input_tokens'] for stat in stats)
+            total_output = sum(stat['output_tokens'] for stat in stats)
+            total = total_input + total_output
+            
+            # Display individual bot statistics
+            for stat in stats:
+                st.sidebar.markdown(
+                    f"_{stat['bot_name']}_  \n"
+                    f"Input: {stat['input_tokens']} tokens  \n"
+                    f"Output: {stat['output_tokens']} tokens  \n"
+                    f"Total: {stat['total_tokens']} tokens"
+                )
+            
+            # Display interaction totals
+            st.sidebar.markdown("**Interaction Totals:**")
+            st.sidebar.markdown(
+                f"Total Input: {total_input} tokens  \n"
+                f"Total Output: {total_output} tokens  \n"
+                f"Total Usage: {total} tokens"
+            )
+            st.sidebar.markdown("---")
+    
+    st.sidebar.markdown("")
+
 
     if not check_password():
         st.stop()
@@ -364,7 +402,6 @@ 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)
@@ -469,7 +506,7 @@ 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("---")