Update src/app.py

src/app.py

@@ -209,30 +209,15 @@ tab1, tab2, tab3 = st.tabs(["💬 Chat Playground", "📂 Knowledge & Tools", "
 
 # === TAB 1: CHAT ===
 with tab1:
-    st.header("Discussion & Analysis")
-
-    # Chat History Downloader
-    # Only show the button if there is actually a conversation to save
-    if st.session_state.get("messages"):
-        # 1. Format the history log
-        chat_log = f"# ⚓ Navy AI Toolkit - Chat Log\n"
-        chat_log += f"Date: {datetime.now().strftime('%Y-%m-%d %H:%M')}\n"
-        chat_log += f"Model: {st.session_state.get('model_selector_radio', 'Unknown')}\n\n"
-        chat_log += "---\n\n"
-        
-        for msg in st.session_state.messages:
-            role = msg["role"].upper()
-            content = msg["content"]
-            chat_log += f"**{role}**: {content}\n\n"
-            
-        # 2. Render the Download Button (Right below header)
-        st.download_button(
-            label="💾 Download Conversation",
-            data=chat_log,
-            file_name=f"chat_session_{datetime.now().strftime('%Y%m%d_%H%M')}.md",
-            mime="text/markdown"
-        )
+    # 1. LAYOUT: Header + Placeholder for Download Button
+    col_header, col_btn = st.columns([6, 1])
+    with col_header:
+        st.header("Discussion & Analysis")
+    
+    # Reserve a spot for the button so we can render it LATER (after the chat updates)
+    download_placeholder = col_btn.empty()
     
+    # 2. CHAT LOGIC
     if "messages" not in st.session_state: st.session_state.messages = []
     
     # RENDER DEBUG OVERLAY (If enabled in Admin)
@@ -241,11 +226,16 @@ with tab1:
     c1, c2 = st.columns([3, 1])
     with c1: st.caption(f"Active Model: **{st.session_state.get('model_selector_radio', 'Granite')}**")
     with c2: use_rag = st.toggle("Enable Knowledge Base", value=False)
+    
+    # Display existing history
     for msg in st.session_state.messages:
         with st.chat_message(msg["role"]): st.markdown(msg["content"])
+        
+    # Handle New Input
     if prompt := st.chat_input("Input command..."):
         st.session_state.messages.append({"role": "user", "content": prompt})
         with st.chat_message("user"): st.markdown(prompt)
+        
         context_txt = ""
         sys_p = "You are a helpful AI assistant."
         st.session_state.last_context_used = "" # Reset context debug
@@ -267,6 +257,7 @@ with tab1:
                             context_txt += f"<document index='{i+1}' source='{src}'>\n{d.page_content}\n</document>\n"
                         # Debug Capture
                         st.session_state.last_context_used = context_txt
+        
         if context_txt:
             final_prompt = f"User Question: {prompt}\n\n<context>\n{context_txt}\n</context>\n\nInstruction: Answer using the context above."
         else: final_prompt = prompt
@@ -280,10 +271,33 @@ with tab1:
                     m_name = "GPT-4o" if "GPT-4o" in model_choice else model_choice.split()[0]
                     tracker.log_usage(m_name, usage["input"], usage["output"])
                     update_sidebar_metrics()
+        
         st.session_state.messages.append({"role": "assistant", "content": resp})
+        
         if use_rag and context_txt:
             with st.expander("📚 View Context Used"): st.text(context_txt)
 
+    # 3. LATE RENDER: Fill the Download Button Placeholder
+    # This runs AFTER the new message is appended, so the log is complete.
+    if st.session_state.messages:
+        chat_log = f"# ⚓ Navy AI Toolkit - Chat Log\n"
+        chat_log += f"Date: {datetime.now().strftime('%Y-%m-%d %H:%M')}\n"
+        chat_log += f"Model: {st.session_state.get('model_selector_radio', 'Unknown')}\n\n"
+        chat_log += "---\n\n"
+        
+        for msg in st.session_state.messages:
+            role = msg["role"].upper()
+            content = msg["content"]
+            chat_log += f"**{role}**: {content}\n\n"
+            
+        with download_placeholder:
+            st.download_button(
+                label="💾 Save",
+                data=chat_log,
+                file_name=f"chat_{datetime.now().strftime('%Y%m%d_%H%M')}.md",
+                mime="text/markdown"
+            )
+
 # === TAB 2: KNOWLEDGE & TOOLS ===
 with tab2:
     st.header("Document Processor")