feat(ui): add real-time token streaming with st.write_stream (closes #1)

app.py

@@ -212,11 +212,11 @@ def render_hitl_stage(title: str, stage_obj, state_key: str, input_data: str, pl
             current_val = st.session_state.get(state_key, "")
 
             if not current_val:
-                st.info(f"Click **Execute {title}** to generate output using local AI.", icon=":material/info:")
+                st.info(f"Click **Execute {title}** to stream output in real-time using local AI.", icon=":material/info:")
                 if st.button(f"Generate {title}", type="primary", icon=":material/play_arrow:"):
-                    with st.spinner("Processing with Hermes 3 on GPU..."):
-                        res = stage_obj.execute(input_data)
-                        st.session_state[state_key] = res
+                    with st.spinner("Streaming response token-by-token from GPU..."):
+                        streamed_text = st.write_stream(stage_obj.execute_stream(input_data))
+                        st.session_state[state_key] = streamed_text
                         save_session_state()
                         st.rerun()
             else:
@@ -245,9 +245,9 @@ def render_hitl_stage(title: str, stage_obj, state_key: str, input_data: str, pl
             with bcol1:
                 if st.button("Refine Output", key=f"refine_{state_key}", icon=":material/auto_fix_high:", use_container_width=True):
                     if feedback.strip() and st.session_state.get(state_key):
-                        with st.spinner("Surgically refining output..."):
-                            refined = stage_obj.refine(st.session_state[state_key], feedback)
-                            st.session_state[state_key] = refined
+                        with st.spinner("Surgically streaming refinements..."):
+                            refined_stream = st.write_stream(stage_obj.refine_stream(st.session_state[state_key], feedback))
+                            st.session_state[state_key] = refined_stream
                             save_session_state()
                             st.rerun()
                     else:
@@ -339,27 +339,29 @@ def render_hitl_stage(title: str, stage_obj, state_key: str, input_data: str, pl
         st.subheader("Architect Recommendation & Negotiation")
         if not st.session_state.stage9a_out:
             if st.button("Generate Tech Stack Proposal", type="primary", icon=":material/architecture:"):
-                with st.spinner("Analyzing sprint stories and evaluating trade-offs..."):
-                    st.session_state.stage9a_out = s9.recommend_tech_stack(sprint_in)
+                with st.spinner("Analyzing sprint stories and streaming recommendation..."):
+                    streamed_stack = st.write_stream(s9.recommend_tech_stack_stream(sprint_in))
+                    st.session_state.stage9a_out = streamed_stack
                     save_session_state()
                     st.rerun()
         else:
             st.markdown(st.session_state.stage9a_out)
             dev_feedback = st.text_input("Developer Suggestion / Alternative Stack Preference", placeholder="e.g., Recommend Streamlit for speed, or negotiate Flask/React...")
             if st.button("Send to Architect for Review", icon=":material/send:"):
                 with st.spinner("Architect analyzing trade-offs..."):
-                    st.session_state.stage9a_out = s9.negotiate_tech_stack(st.session_state.stage9a_out, dev_feedback)
+                    streamed_reply = st.write_stream(s9.negotiate_tech_stack_stream(st.session_state.stage9a_out, dev_feedback))
+                    st.session_state.stage9a_out = streamed_reply
                     save_session_state()
                     st.rerun()
 
     with t2:
         st.subheader("Executable Prototype Code Generation")
         if not st.session_state.stage9b_out:
             if st.button("Generate Executable Prototype Code", type="primary", icon=":material/code:"):
-                with st.spinner("Generating modular Python engine and Streamlit application..."):
-                    code_res = s9.generate_code(sprint_in, st.session_state.stage9a_out)
+                with st.spinner("Streaming modular Python engine and Streamlit application..."):
+                    code_res = st.write_stream(s9.generate_code_stream(sprint_in, st.session_state.stage9a_out))
                     st.session_state.stage9b_out = code_res
-                    saved_files = s9.extract_and_save(code_res, config.prototype_dir)
+                    s9.extract_and_save(code_res, config.prototype_dir)
                     save_session_state()
                     st.rerun()
         else:
@@ -386,8 +388,8 @@ def render_hitl_stage(title: str, stage_obj, state_key: str, input_data: str, pl
 
     if not st.session_state.stage10_out:
         if st.button("Generate Test Cases & Pytest Script", type="primary", icon=":material/bug_report:"):
-            with st.spinner("Generating structured test cases and pytest suite..."):
-                test_res = s10.generate_tests(sprint_in, proto_in)
+            with st.spinner("Streaming structured test cases and pytest suite..."):
+                test_res = st.write_stream(s10.generate_tests_stream(sprint_in, proto_in))
                 st.session_state.stage10_out = test_res
                 s10.extract_and_save(test_res, config.prototype_dir)
                 with open(config.artifacts_dir / "test_specifications.md", "w", encoding="utf-8") as f:

src/core/llm_service.py

@@ -1,9 +1,9 @@
-from typing import Optional
+from typing import Optional, Generator
 from langchain_ollama import ChatOllama
 from .config import AppConfig, get_config
 
 class LLMService:
-    """Manages ChatOllama connections, GPU offloading, fallbacks, and invocation helpers."""
+    """Manages ChatOllama connections, GPU offloading, fallbacks, and invocation/streaming helpers."""
 
     def __init__(self, config: Optional[AppConfig] = None):
         self.config = config or get_config()
@@ -37,15 +37,24 @@ def _initialize_client(self) -> ChatOllama:
             return client
 
     def invoke(self, system_prompt: str, user_content: str) -> str:
-        """Executes a standard system + user prompt pair."""
+        """Executes a standard system + user prompt pair synchronously."""
         response = self.client.invoke([
             ("system", system_prompt),
             ("user", user_content)
         ])
         return str(response.content)
 
+    def stream(self, system_prompt: str, user_content: str) -> Generator[str, None, None]:
+        """Streams response tokens in real-time as a Python generator."""
+        chunks = self.client.stream([
+            ("system", system_prompt),
+            ("user", user_content)
+        ])
+        for chunk in chunks:
+            yield str(chunk.content)
+
     def refine(self, original_output: str, user_feedback: str, role_title: str = "Senior Software Engineering Specialist") -> str:
-        """Refines existing output incorporating direct human feedback."""
+        """Refines existing output incorporating direct human feedback synchronously."""
         refinement_system_prompt = (
             f"Role: {role_title}\n"
             "Task: You are surgically refining the previous output based on direct user feedback. "
@@ -55,6 +64,17 @@ def refine(self, original_output: str, user_feedback: str, role_title: str = "Se
         refinement_user_prompt = f"Original Output:\n{original_output}\n\nUser Feedback:\n{user_feedback}"
         return self.invoke(refinement_system_prompt, refinement_user_prompt)
 
+    def stream_refine(self, original_output: str, user_feedback: str, role_title: str = "Senior Software Engineering Specialist") -> Generator[str, None, None]:
+        """Streams refined output tokens in real-time."""
+        refinement_system_prompt = (
+            f"Role: {role_title}\n"
+            "Task: You are surgically refining the previous output based on direct user feedback. "
+            "Strictly incorporate the user's feedback into the output while maintaining clean, structured formatting. "
+            "Preserve all existing valid data while applying the requested revisions."
+        )
+        refinement_user_prompt = f"Original Output:\n{original_output}\n\nUser Feedback:\n{user_feedback}"
+        yield from self.stream(refinement_system_prompt, refinement_user_prompt)
+
     def get_status(self) -> dict:
         """Returns health and runtime status of the LLM connection."""
         return {

src/stages/base_stage.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Optional, Generator
 from ..core.llm_service import LLMService, get_llm_service
 
 class BaseStage(ABC):
@@ -17,9 +17,17 @@ def system_prompt(self) -> str:
         pass
 
     def execute(self, input_data: str) -> str:
-        """Executes the initial prompt for this stage."""
+        """Executes the initial prompt for this stage synchronously."""
         return self.llm.invoke(self.system_prompt, f"Input Data / Context:\n{input_data}")
 
+    def execute_stream(self, input_data: str) -> Generator[str, None, None]:
+        """Streams the initial prompt execution tokens in real-time."""
+        yield from self.llm.stream(self.system_prompt, f"Input Data / Context:\n{input_data}")
+
     def refine(self, current_output: str, feedback: str) -> str:
-        """Refines the current output based on user feedback."""
+        """Refines the current output based on user feedback synchronously."""
         return self.llm.refine(current_output, feedback, role_title=f"Agile Specialist ({self.stage_name})")
+
+    def refine_stream(self, current_output: str, feedback: str) -> Generator[str, None, None]:
+        """Streams refined output tokens in real-time."""
+        yield from self.llm.stream_refine(current_output, feedback, role_title=f"Agile Specialist ({self.stage_name})")

src/stages/stage10_test_cases.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import List
+from typing import List, Generator
 from .base_stage import BaseStage
 from ..prompts.agile_prompts import PROMPT_STAGE_10_TEST_CASES
 from ..utils.file_extractor import extract_and_save_code_files
@@ -13,8 +13,14 @@ def system_prompt(self) -> str:
         return PROMPT_STAGE_10_TEST_CASES
 
     def generate_tests(self, sprint_stories: str, prototype_code: str) -> str:
+        """Generates test cases and pytest script synchronously."""
         prompt_input = f"Selected Sprint Stories:\n{sprint_stories}\n\nGenerated Prototype Implementation:\n{prototype_code}"
         return self.llm.invoke(self.system_prompt, prompt_input)
 
+    def generate_tests_stream(self, sprint_stories: str, prototype_code: str) -> Generator[str, None, None]:
+        """Streams test cases and pytest script generation in real-time."""
+        prompt_input = f"Selected Sprint Stories:\n{sprint_stories}\n\nGenerated Prototype Implementation:\n{prototype_code}"
+        yield from self.llm.stream(self.system_prompt, prompt_input)
+
     def extract_and_save(self, test_output: str, target_dir: Path) -> List[Path]:
         return extract_and_save_code_files(test_output, target_dir)

src/stages/stage9_prototype.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Generator
 from pathlib import Path
 from .base_stage import BaseStage
 from ..prompts.agile_prompts import PROMPT_STAGE_9_TECH_STACK, PROMPT_STAGE_9_PROTOTYPE
@@ -13,11 +13,15 @@ def system_prompt(self) -> str:
         return PROMPT_STAGE_9_PROTOTYPE
 
     def recommend_tech_stack(self, sprint_context: str) -> str:
-        """Executes Sub-stage 9A: Tech stack recommendation."""
+        """Executes Sub-stage 9A: Tech stack recommendation synchronously."""
         return self.llm.invoke(PROMPT_STAGE_9_TECH_STACK, f"Selected Sprint Stories:\n{sprint_context}")
 
+    def recommend_tech_stack_stream(self, sprint_context: str) -> Generator[str, None, None]:
+        """Streams Sub-stage 9A: Tech stack recommendation in real-time."""
+        yield from self.llm.stream(PROMPT_STAGE_9_TECH_STACK, f"Selected Sprint Stories:\n{sprint_context}")
+
     def negotiate_tech_stack(self, previous_discussion: str, developer_input: str) -> str:
-        """Conducts Sub-stage 9A: Architect-Developer negotiation."""
+        """Conducts Sub-stage 9A: Architect-Developer negotiation synchronously."""
         negotiation_system_prompt = (
             "Role: Principal Software Architect\n"
             "Task: Discuss the tech stack choice with the developer. Address their questions, analyze trade-offs of their proposed stack versus the recommended one, "
@@ -26,11 +30,26 @@ def negotiate_tech_stack(self, previous_discussion: str, developer_input: str) -
         negotiation_user_prompt = f"Previous Discussion:\n{previous_discussion}\n\nDeveloper Feedback:\n{developer_input}"
         return self.llm.invoke(negotiation_system_prompt, negotiation_user_prompt)
 
+    def negotiate_tech_stack_stream(self, previous_discussion: str, developer_input: str) -> Generator[str, None, None]:
+        """Streams Sub-stage 9A: Architect-Developer negotiation in real-time."""
+        negotiation_system_prompt = (
+            "Role: Principal Software Architect\n"
+            "Task: Discuss the tech stack choice with the developer. Address their questions, analyze trade-offs of their proposed stack versus the recommended one, "
+            "and finalize the agreed tech stack architecture for the prototype."
+        )
+        negotiation_user_prompt = f"Previous Discussion:\n{previous_discussion}\n\nDeveloper Feedback:\n{developer_input}"
+        yield from self.llm.stream(negotiation_system_prompt, negotiation_user_prompt)
+
     def generate_code(self, sprint_context: str, tech_stack_agreement: str) -> str:
-        """Executes Sub-stage 9B: Generates complete executable prototype code."""
+        """Executes Sub-stage 9B: Generates complete executable prototype code synchronously."""
         prompt_input = f"Approved Sprint Context:\n{sprint_context}\n\nAgreed Tech Stack Architecture:\n{tech_stack_agreement}"
         return self.llm.invoke(PROMPT_STAGE_9_PROTOTYPE, prompt_input)
 
+    def generate_code_stream(self, sprint_context: str, tech_stack_agreement: str) -> Generator[str, None, None]:
+        """Streams Sub-stage 9B: Generates complete executable prototype code in real-time."""
+        prompt_input = f"Approved Sprint Context:\n{sprint_context}\n\nAgreed Tech Stack Architecture:\n{tech_stack_agreement}"
+        yield from self.llm.stream(PROMPT_STAGE_9_PROTOTYPE, prompt_input)
+
     def extract_and_save(self, code_output: str, target_dir: Path) -> List[Path]:
         """Saves generated code files directly to disk."""
         return extract_and_save_code_files(code_output, target_dir)

tests/test_streaming.py

@@ -0,0 +1,34 @@
+import pytest
+import sys
+from pathlib import Path
+
+# Add project root to sys.path
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from src.stages.base_stage import BaseStage
+from src.stages.stage1_stakeholders import Stage1Stakeholders
+
+def test_base_stage_streaming_interface():
+    class MockStreamingLLM:
+        def stream(self, system_prompt, user_content):
+            chunks = ["Hello", " ", "world", "!"]
+            for chunk in chunks:
+                yield chunk
+
+        def stream_refine(self, original_output, user_feedback, role_title=""):
+            chunks = ["Refined", " ", "output", "!"]
+            for chunk in chunks:
+                yield chunk
+
+    mock_llm = MockStreamingLLM()
+    stage = Stage1Stakeholders(mock_llm)
+
+    # Test execute_stream
+    stream_gen = stage.execute_stream("Input text")
+    collected = "".join(list(stream_gen))
+    assert collected == "Hello world!"
+
+    # Test refine_stream
+    refine_gen = stage.refine_stream("Initial", "Feedback")
+    collected_refine = "".join(list(refine_gen))
+    assert collected_refine == "Refined output!"