Update src/streamlit_app.py

src/streamlit_app.py

@@ -102,8 +102,6 @@ SAMPLE_ASSISTANT_MSG = {
 # =========================
 if "messages" not in st.session_state:
     st.session_state.messages = []
-if "needs_answer" not in st.session_state:
-    st.session_state.needs_answer = False
 if "seeded" not in st.session_state:
     st.session_state.seeded = False
 
@@ -114,7 +112,7 @@ if (not st.session_state.seeded) and (len(st.session_state.messages) == 0):
 
 
 # =========================
-# HEADER (prevent logo flash)
+# HEADER
 # =========================
 col1, col2 = st.columns([1, 7], vertical_alignment="center")
 with col1:
@@ -146,10 +144,9 @@ with st.sidebar:
     )
 
     if st.button("Clear chat"):
-            st.session_state.messages = []
-            st.session_state.needs_answer = False
-            st.session_state.seeded = True
-            st.rerun()    
+        st.session_state.messages = []
+        st.session_state.seeded = True
+        st.rerun()
 
 
 # =========================
@@ -312,7 +309,6 @@ for i, msg in enumerate(st.session_state.messages):
             )
 
         with metrics_col:
-            
             with st.popover("ℹ️", help="Inspect"):
                 st.markdown("**Context**")
                 st.code(msg.get("context", ""), language="text")
@@ -323,25 +319,16 @@ for i, msg in enumerate(st.session_state.messages):
 
 
 # =========================
-# INPUT
+# INPUT + GENERATE (NO RERUN / NO FLASH)
 # =========================
 query = st.chat_input("Ask a question...")
 
 if query:
+    # Persist user message
     st.session_state.messages.append({"role": "user", "content": query})
-    st.session_state.needs_answer = True
-    st.rerun()
 
-
-# =========================
-# GENERATE (once per user message)
-# =========================
-if (
-    st.session_state.needs_answer
-    and st.session_state.messages
-    and st.session_state.messages[-1]["role"] == "user"
-):
-    question = st.session_state.messages[-1]["content"]
+    # Generate immediately in the same run (no st.rerun)
+    question = query
 
     # Web search
     web_ctx, search_time = web_search_snippets(question)
@@ -351,7 +338,7 @@ if (
     prompt = f"{context}\n{system_prompt}\n{question}\n"
     input_tokens = count_tokens(prompt)
 
-    # Run traced answer
+    # Assistant response
     with st.chat_message("assistant"):
         # Row 1: message + feedback (feedback disabled until persisted)
         msg_col, fb_col = st.columns([14, 1], vertical_alignment="center")
@@ -364,12 +351,14 @@ if (
         answer = traced_answer(context, system_prompt, question)
         trace_id = get_trace_id_if_available()
 
-        # Stream into the message column
+        # Typewriter render (reduce updates to avoid jitter)
         buf = ""
-        for ch in answer:
+        for j, ch in enumerate(answer, 1):
             buf += ch
-            placeholder.markdown(buf)
-            time.sleep(0.002)
+            if j % 6 == 0:  # update every 6 chars
+                placeholder.markdown(buf)
+                time.sleep(0.001)
+        placeholder.markdown(buf)
 
         gen_time = time.perf_counter() - start
         output_tokens = count_tokens(answer)
@@ -394,9 +383,8 @@ if (
                 st.code(question, language="text")
                 st.markdown("**Prompt**")
                 st.code(prompt, language="text")
-            
 
-    # Persist assistant message
+    # Persist assistant message (so it shows on subsequent runs)
     st.session_state.messages.append(
         {
             "role": "assistant",
@@ -413,7 +401,4 @@ if (
             "trace_id": trace_id,
             "feedback": None,
         }
-    )
-
-    st.session_state.needs_answer = False
-    st.rerun()
+    )