From aa8b8e76953f82857c0bda2bb6e4e698c97b7db9 Mon Sep 17 00:00:00 2001 From: Prince Patel Date: Fri, 28 Aug 2026 18:38:16 +0530 Subject: [PATCH] feat(ui): add real-time token streaming with st.write_stream (closes #1) --- app.py | 32 ++++++++++++++++-------------- src/core/llm_service.py | 28 ++++++++++++++++++++++---- src/stages/base_stage.py | 14 ++++++++++--- src/stages/stage10_test_cases.py | 8 +++++++- src/stages/stage9_prototype.py | 27 +++++++++++++++++++++---- tests/test_streaming.py | 34 ++++++++++++++++++++++++++++++++ 6 files changed, 116 insertions(+), 27 deletions(-) create mode 100644 tests/test_streaming.py diff --git a/app.py b/app.py index a63fe0f..9c1b938 100644 --- a/app.py +++ b/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,8 +339,9 @@ 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: @@ -348,7 +349,8 @@ def render_hitl_stage(title: str, stage_obj, state_key: str, input_data: str, pl 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() @@ -356,10 +358,10 @@ def render_hitl_stage(title: str, stage_obj, state_key: str, input_data: str, pl 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: diff --git a/src/core/llm_service.py b/src/core/llm_service.py index f462f27..44e3808 100644 --- a/src/core/llm_service.py +++ b/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 { diff --git a/src/stages/base_stage.py b/src/stages/base_stage.py index 238681c..8d7bafc 100644 --- a/src/stages/base_stage.py +++ b/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})") diff --git a/src/stages/stage10_test_cases.py b/src/stages/stage10_test_cases.py index f32e1bd..9e2ac0a 100644 --- a/src/stages/stage10_test_cases.py +++ b/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) diff --git a/src/stages/stage9_prototype.py b/src/stages/stage9_prototype.py index b7489a7..b2d37f1 100644 --- a/src/stages/stage9_prototype.py +++ b/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) diff --git a/tests/test_streaming.py b/tests/test_streaming.py new file mode 100644 index 0000000..5c39c51 --- /dev/null +++ b/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!"