Add reasoning-model support, multi-file output, pipeline profiles, and persistent chats

- Streaming now hides model reasoning: inline <think> blocks (tag-split
  tolerant ThinkTagFilter) and reasoning_content/reasoning deltas stream
  into a collapsed "Model reasoning" expander instead of the output
- Live tok/s caption updates every ~24 chunks during generation; final
  stats formatting shared via core/api.py::format_stream_stats
- Output page parses every fenced block: multi-file results render as
  per-file tabs with a zip download and save-all into a timestamped
  folder (filenames sniffed from the line preceding each fence,
  path-sanitized)
- Refine/regenerate keeps a version stack: unified diff against the
  previous code with one-click revert; a rough token estimate warns when
  the refine payload may exceed small local context windows
- New swap_policy setting (auto/never) gates every cross-role VRAM
  unload for GPUs that fit both models at once
- Chats persist to chats.json (new core/chats.py); Prompter-chat replies
  gain "Use as pipeline prompt", jumping to the review step with the
  fenced prompt extracted and the preceding message as the task
- Pipeline profiles: named snapshots of backends/models/params/system
  prompts saved in config, applied in one click from the task page
- 39 new tests (think filter + SSE reasoning paths, extract_files,
  profiles, swap policy, chat store, coder-role and category-switch
  preset loading); suite now at 142

Author: atharva557

.gitignore

@@ -1,5 +1,6 @@
 config.json
 history.json
+chats.json
 output/
 __pycache__/
 *.pyc

app.py

@@ -10,12 +10,12 @@
 
 import streamlit as st
 
-from core.config import config_exists, load_config, get_role_endpoint
+from core.config import config_exists, load_config, get_role_endpoint, swap_enabled
 from core.api import (
     test_connection,
     unload_model,
     is_cloud,
-    estimate_cost,
+    format_stream_stats,
     BACKEND_LABELS,
     schedule_unload,
     cancel_unload,
@@ -103,6 +103,9 @@ def init_session_state():
         # Which role's model most recently handled a request (pipeline or
         # chat), so all generation paths can swap VRAM cooperatively.
         "last_model_role": None,
+        # Previous code outputs of this run (newest last), so a refine or
+        # regenerate that makes things worse can be compared and reverted.
+        "code_versions": [],
     }
     for key, value in defaults.items():
         if key not in st.session_state:
@@ -126,6 +129,7 @@ def reset_pipeline():
     st.session_state["prompt_refine_instruction"] = ""
     st.session_state["needs_swap"] = False
     st.session_state["show_settings"] = False
+    st.session_state["code_versions"] = []
     # Drop widget state so inputs re-initialize cleanly
     for key in (
         "task_input_area",
@@ -399,6 +403,7 @@ def run_streaming_generation(
     first_token = True
     chunk_count = 0
     usage: dict = {}
+    reasoning: dict = {}
     start_time = time.perf_counter()
 
     try:
@@ -414,16 +419,31 @@ def run_streaming_generation(
             messages=messages,
             api_key=endpoint["api_key"],
             usage_out=usage,
+            reasoning_out=reasoning,
         )
 
-        # Create a placeholder for streaming output
+        # Placeholders in display order: hidden reasoning (filled lazily,
+        # only for thinking models), then the output, then live stats.
+        reasoning_slot = st.empty()
+        reasoning_box = None
         output_placeholder = st.empty()
+        stats_placeholder = st.empty()
 
         for token in stream:
             if first_token:
                 first_token = False
                 if on_first_token:
                     on_first_token()
+            # Thinking models: stream the hidden reasoning into a collapsed
+            # expander instead of polluting the prompt/code output
+            if reasoning.get("text"):
+                if reasoning_box is None:
+                    with reasoning_slot.container():
+                        with st.expander("Model reasoning", expanded=False):
+                            reasoning_box = st.empty()
+                reasoning_box.markdown(reasoning["text"])
+            if not token:
+                continue  # reasoning-only chunk
             chunk_count += 1
             full_output += token
             if partial_key:
@@ -432,30 +452,18 @@ def run_streaming_generation(
                 output_placeholder.code(full_output, line_numbers=True)
             else:
                 output_placeholder.markdown(full_output)
-
-        elapsed = time.perf_counter() - start_time
-        stats = []
-        if chunk_count and elapsed > 0:
-            # One SSE chunk is roughly one token
-            stats.append(
-                f"~{chunk_count} tokens in {elapsed:.1f}s "
-                f"({chunk_count / elapsed:.1f} tok/s)"
-            )
-        if usage.get("input_tokens") is not None or usage.get("output_tokens") is not None:
-            # Exact billed counts (cloud backends report these)
-            stats.append(
-                f"{usage.get('input_tokens', '?')} in / "
-                f"{usage.get('output_tokens', '?')} out tokens"
-            )
-            cost = estimate_cost(
-                endpoint["model"],
-                usage.get("input_tokens"),
-                usage.get("output_tokens"),
-            )
-            if cost is not None:
-                stats.append(f"~${cost:.4f} est.")
-        if stats:
-            st.caption(" · ".join(stats))
+            if chunk_count % 24 == 0:
+                elapsed = time.perf_counter() - start_time
+                if elapsed > 0:
+                    stats_placeholder.caption(
+                        f"~{chunk_count} tokens · {chunk_count / elapsed:.1f} tok/s"
+                    )
+
+        stats_line = format_stream_stats(
+            endpoint["model"], chunk_count, time.perf_counter() - start_time, usage
+        )
+        if stats_line:
+            stats_placeholder.caption(stats_line)
 
         return full_output, ""
 
@@ -534,8 +542,9 @@ def run_prompter(status_placeholder, messages: list[dict] | None = None) -> tupl
     step) instead of the default single-turn task rewrite."""
 
     # A local coder model left loaded by a chat or earlier run would fight the
-    # prompter for VRAM — evict it first. (No-op for a cloud coder.)
-    if st.session_state.get("last_model_role") == "coder":
+    # prompter for VRAM — evict it first. (No-op for a cloud coder, skipped
+    # entirely when the swap policy says both models fit.)
+    if st.session_state.get("last_model_role") == "coder" and swap_enabled(config):
         coder = get_role_endpoint(config, "coder")
         unload_model(coder["base_url"], coder["model"], coder["backend"])
     st.session_state["last_model_role"] = "prompter"
@@ -721,8 +730,9 @@ def _stop_coder():
 
         # Only swap if a local prompter actually ran since the last unload —
         # regenerating from the output page (or a cloud prompter) would
-        # otherwise trigger a pointless unload.
-        if st.session_state.pop("needs_swap", False):
+        # otherwise trigger a pointless unload. The 'never' swap policy
+        # (both models fit in VRAM) skips it too; the flag is still popped.
+        if st.session_state.pop("needs_swap", False) and swap_enabled(config):
             with status_placeholder.container():
                 render_pipeline_checklist([
                     {"label": "Prompt confirmed", "state": "done"},
@@ -803,6 +813,11 @@ def _on_first_token():
             st.session_state["coder_partial"] = result
             st.rerun()  # re-render into the stable error view below
         else:
+            # Keep the code this run replaced, so the output page can show a
+            # diff of what changed and offer a one-click revert.
+            prev_code = st.session_state.get("generated_code", "")
+            if prev_code and prev_code != result:
+                st.session_state.setdefault("code_versions", []).append(prev_code)
             st.session_state["generated_code"] = result
             st.session_state["coder_partial"] = ""
             st.session_state["coder_mode"] = "full"

core/api.py

@@ -76,6 +76,38 @@ def estimate_cost(model: str, input_tokens, output_tokens) -> float | None:
     return None
 
 
+def rough_token_count(*texts: str) -> int:
+    """Very rough token estimate (~4 chars/token) across the given strings.
+    Good enough for a context-size warning; never used for billing."""
+    return sum(len(t) for t in texts if t) // 4
+
+
+def format_stream_stats(
+    model: str, chunk_count: int, elapsed: float, usage: dict
+) -> str:
+    """One-line stats caption for a generation: chunk-based tok/s estimate,
+    plus exact billed tokens and estimated cost when the backend reported
+    usage. Returns '' when there is nothing to show."""
+    stats = []
+    if chunk_count and elapsed > 0:
+        # One SSE chunk is roughly one token
+        stats.append(
+            f"~{chunk_count} tokens in {elapsed:.1f}s "
+            f"({chunk_count / elapsed:.1f} tok/s)"
+        )
+    if usage.get("input_tokens") is not None or usage.get("output_tokens") is not None:
+        stats.append(
+            f"{usage.get('input_tokens', '?')} in / "
+            f"{usage.get('output_tokens', '?')} out tokens"
+        )
+        cost = estimate_cost(
+            model, usage.get("input_tokens"), usage.get("output_tokens")
+        )
+        if cost is not None:
+            stats.append(f"~${cost:.4f} est.")
+    return " · ".join(stats)
+
+
 def chat_completions_url(base_url: str, backend: str) -> str:
     """Chat-completions endpoint for a backend (Anthropic excluded — it
     goes through the official SDK, not an OpenAI-compatible path)."""

core/chats.py

@@ -0,0 +1,41 @@
+"""
+PromptChain — Chat Persistence
+
+Persists the two single-model chats (Prompter / Coder) to chats.json so a
+conversation survives app restarts, mirroring how history.json persists
+pipeline runs. Shared across tabs/sessions; concurrent saves are
+last-writer-wins, which is acceptable for a single-user local app.
+"""
+
+import json
+from pathlib import Path
+
+CHATS_PATH = Path(__file__).parent.parent / "chats.json"
+# Keep persisted conversations from growing without bound
+MAX_MESSAGES = 200
+
+
+def _load_all() -> dict:
+    if not CHATS_PATH.exists():
+        return {}
+    try:
+        data = json.loads(CHATS_PATH.read_text(encoding="utf-8"))
+        return data if isinstance(data, dict) else {}
+    except (json.JSONDecodeError, OSError):
+        return {}
+
+
+def load_chat_messages(role: str) -> list[dict]:
+    """Persisted messages for 'prompter' or 'coder'. [] on any problem."""
+    msgs = _load_all().get(role, [])
+    return msgs if isinstance(msgs, list) else []
+
+
+def save_chat_messages(role: str, messages: list[dict]) -> None:
+    """Persist one role's conversation (trimmed to the newest MAX_MESSAGES)."""
+    data = _load_all()
+    data[role] = messages[-MAX_MESSAGES:]
+    CHATS_PATH.write_text(
+        json.dumps(data, indent=2, ensure_ascii=False),
+        encoding="utf-8",
+    )

core/config.py

@@ -32,6 +32,13 @@ def get_default_config() -> dict:
         # Free a resident local model's VRAM after this many idle minutes;
         # 0 = never. Applies to whichever role (prompter or coder) served last.
         "idle_unload_minutes": 5,
+        # 'auto' unloads one local model before running the other (single-GPU
+        # default); 'never' skips all cross-role unloads for machines whose
+        # VRAM fits both models at once.
+        "swap_policy": "auto",
+        # Named snapshots of a full pipeline setup (endpoints, sampling params,
+        # system prompts) that can be applied in one click from the task page.
+        "pipeline_profiles": {},
         "custom_presets": {},
         # Edits to built-in presets, keyed role -> category -> name -> text.
         # Kept separate so the shipped presets.json is never mutated.
@@ -114,6 +121,12 @@ def get_role_endpoint(config: dict, role: str) -> dict:
         "api_key": _resolve_api_key(config, backend),
     }
 
+def swap_enabled(config: dict) -> bool:
+    """False when the user opted out of VRAM swapping ('never' policy —
+    both models fit in VRAM, so cross-role unloads are pure overhead)."""
+    return config.get("swap_policy", "auto") != "never"
+
+
 def save_config(config: dict) -> None:
     """
     Save config dict to config.json.
@@ -127,6 +140,56 @@ def save_config(config: dict) -> None:
     )
 
 
+# ═══════════════════════════════════════════════════════
+#  Pipeline Profiles
+# ═══════════════════════════════════════════════════════
+
+# Config keys snapshotted into (and restored from) a pipeline profile —
+# everything that defines "which two models run and how", so switching e.g.
+# "local drafting" ↔ "cloud final pass" is one click instead of a Settings trip.
+PROFILE_CONFIG_KEYS = (
+    "prompter_backend", "prompter_base_url", "prompter_model",
+    "coder_backend", "coder_base_url", "coder_model",
+    "prompter_temperature", "coder_temperature",
+    "prompter_max_tokens", "coder_max_tokens",
+)
+
+
+def capture_pipeline_profile(config: dict, prompter_system: str, coder_system: str) -> dict:
+    """Snapshot the current endpoints/params plus the active system prompts."""
+    profile = {key: config.get(key) for key in PROFILE_CONFIG_KEYS}
+    profile["prompter_system"] = prompter_system
+    profile["coder_system"] = coder_system
+    return profile
+
+
+def apply_pipeline_profile(config: dict, profile: dict) -> dict:
+    """Write a profile's endpoint/param snapshot into the config and persist
+    it. System prompts are returned to the caller via the profile itself —
+    they live in session state, not config."""
+    for key in PROFILE_CONFIG_KEYS:
+        if profile.get(key) is not None:
+            config[key] = profile[key]
+    save_config(config)
+    return config
+
+
+def save_pipeline_profile(config: dict, name: str, profile: dict) -> dict:
+    """Create or overwrite a named pipeline profile."""
+    config.setdefault("pipeline_profiles", {})[name] = profile
+    save_config(config)
+    return config
+
+
+def delete_pipeline_profile(config: dict, name: str) -> dict:
+    """Remove a named pipeline profile (no-op if absent)."""
+    profiles = config.get("pipeline_profiles", {})
+    if name in profiles:
+        del profiles[name]
+        save_config(config)
+    return config
+
+
 # ═══════════════════════════════════════════════════════
 #  Preset Management
 # ═══════════════════════════════════════════════════════