Update app.py

app.py

@@ -1,164 +1,1023 @@
-# Provibe App (v1.1) - Streamlit Implementation
-
+# -*- coding: utf-8 -*-
 import streamlit as st
+import requests
 from openai import OpenAI
+import google.generativeai as genai
+import anthropic
+import streamlit.components.v1 as components
+import re
+# from supabase import create_client # Keep commented if not using history/auth
+import base64
+import time
+import json # For potentially structured Q&A
+
+# ---------- Helper: Safe Rerun ----------
+def safe_rerun():
+    """Safely trigger a Streamlit rerun if the function exists."""
+    if hasattr(st, "experimental_rerun"):
+        st.experimental_rerun()
+    elif hasattr(st, "rerun"):
+        st.rerun()
+    else:
+        st.warning("Rerun function not available. Please update Streamlit.")
+
+# ---------- Setup & API Client Initialization ----------
+# (Keep your existing API key loading logic here - unchanged)
+openai_client = None
+genai_client = None
+deepseek_api_key = None
+claude_client = None
+secrets_available = {"openai": False, "gemini": False, "deepseek": False, "claude": False}
+secret_errors = []
+
+# (Keep your existing API key loading and client initialization logic here)
+# ... (Omitted for brevity, assume it's the same as before) ...
+# OpenAI API Key
+try:
+    openai_api_key = st.secrets.get("OPENAI_API_KEY")
+    if openai_api_key:
+        openai_client = OpenAI(api_key=openai_api_key)
+        secrets_available["openai"] = True
+    else:
+        secret_errors.append("Streamlit Secret `OPENAI_API_KEY` not found.")
+except KeyError:
+    secret_errors.append("Streamlit Secret `OPENAI_API_KEY` not found.")
+except Exception as e:
+    secret_errors.append(f"Error initializing OpenAI client: {e}")
+
+# Gemini API Key (Google GenAI)
+try:
+    gemini_api_key = st.secrets.get("GEMINI_API_KEY")
+    if gemini_api_key:
+        genai.configure(api_key=gemini_api_key)
+        genai_client = genai
+        secrets_available["gemini"] = True
+    else:
+        secret_errors.append("Streamlit Secret `GEMINI_API_KEY` not found.")
+except KeyError:
+    secret_errors.append("Streamlit Secret `GEMINI_API_KEY` not found.")
+except Exception as e:
+    secret_errors.append(f"Error initializing Google GenAI client: {e}")
+
+# DeepSeek API Key
+try:
+    deepseek_api_key = st.secrets.get("DEEPSEEK_API_KEY")
+    if deepseek_api_key:
+        secrets_available["deepseek"] = True
+    else:
+        secret_errors.append("Streamlit Secret `DEEPSEEK_API_KEY` not found.")
+except KeyError:
+    secret_errors.append("Streamlit Secret `DEEPSEEK_API_KEY` not found.")
+except Exception as e:
+    secret_errors.append(f"Error reading DeepSeek API key: {e}")
+
+# CLAUDE API Key and Client Initialization
+try:
+    claude_api_key = st.secrets.get("CLAUDE_API_KEY")
+    if claude_api_key:
+        claude_client = anthropic.Anthropic(api_key=claude_api_key)
+        secrets_available["claude"] = True
+    else:
+        secret_errors.append("Streamlit Secret `CLAUDE_API_KEY` not found.")
+except KeyError:
+    secret_errors.append("Streamlit Secret `CLAUDE_API_KEY` not found.")
+except Exception as e:
+    secret_errors.append(f"Error initializing Claude client: {e}")
+
+any_secret_loaded = any(secrets_available.values())
+
+
+# ---------- Model Configuration ----------
+# (Keep your existing SUPPORTED_MODELS dictionary population logic here - unchanged)
+# ... (Omitted for brevity, assume it's the same as before) ...
+SUPPORTED_MODELS = {}
+
+# OpenAI Models
+if secrets_available["openai"] and openai_client:
+    SUPPORTED_MODELS.update({
+        "GPT-4o (OpenAI)": {"id": "gpt-4o", "provider": "openai", "client": openai_client},
+        "GPT-4o Mini (OpenAI)": {"id": "gpt-4o-mini", "provider": "openai", "client": openai_client},
+        "GPT-4 Turbo (OpenAI)": {"id": "gpt-4-turbo", "provider": "openai", "client": openai_client},
+        "GPT-4 (OpenAI)": {"id": "gpt-4", "provider": "openai", "client": openai_client},
+        "GPT-3.5 Turbo (OpenAI)": {"id": "gpt-3.5-turbo", "provider": "openai", "client": openai_client},
+    })
+
+# Gemini Models
+if secrets_available["gemini"] and genai_client:
+    SUPPORTED_MODELS.update({
+        "Gemini 1.5 Pro (Google)": {"id": "gemini-1.5-pro-latest", "provider": "gemini", "client": genai_client},
+        "Gemini 1.5 Flash (Google)": {"id": "gemini-1.5-flash-latest", "provider": "gemini", "client": genai_client},
+        "Gemini 1.0 Pro (Google)": {"id": "gemini-1.0-pro", "provider": "gemini", "client": genai_client},
+    })
+
+# DeepSeek Models
+if secrets_available["deepseek"] and deepseek_api_key:
+    SUPPORTED_MODELS.update({
+        "DeepSeek Chat": {"id": "deepseek-chat", "provider": "deepseek", "client": None},
+        "DeepSeek Coder": {"id": "deepseek-coder", "provider": "deepseek", "client": None},
+    })
+
+# Claude Models
+if secrets_available["claude"] and claude_client:
+    SUPPORTED_MODELS.update({
+        "Claude 3.5 Sonnet (Anthropic)": {"id": "claude-3-5-sonnet-20240620", "provider": "claude", "client": claude_client},
+        "Claude 3 Haiku (Anthropic)": {"id": "claude-3-haiku-20240307", "provider": "claude", "client": claude_client},
+        "Claude 3 Opus (Anthropic)": {"id": "claude-3-opus-20240229", "provider": "claude", "client": claude_client},
+    })
+
+# Determine default model based on preference and availability
+DEFAULT_MODEL_PREFERENCE = [
+    "GPT-4o Mini (OpenAI)",
+    "Gemini 1.5 Flash (Google)",
+    "Claude 3 Haiku (Anthropic)",
+    "DeepSeek Chat",
+    "GPT-3.5 Turbo (OpenAI)",
+]
+DEFAULT_MODEL = next((m for m in DEFAULT_MODEL_PREFERENCE if m in SUPPORTED_MODELS), None)
+if not DEFAULT_MODEL and SUPPORTED_MODELS:
+    DEFAULT_MODEL = next(iter(SUPPORTED_MODELS)) # Fallback to the first available
+
+
+# ---------- Helper Functions for Generation ----------
+# (Keep your existing _generate_with_... provider functions here - unchanged)
+# ... (Omitted for brevity, assume they are the same as before) ...
+def _generate_with_openai_provider(client, model_id, prompt, max_tokens, system_message=None):
+    messages = []
+    if system_message:
+        messages.append({"role": "system", "content": system_message})
+    messages.append({"role": "user", "content": prompt})
+    try:
+        response = client.chat.completions.create(
+            model=model_id,
+            messages=messages,
+            temperature=0.6,
+            max_tokens=max_tokens
+        )
+        return response.choices[0].message.content
+    except Exception as e:
+        st.error(f"❌ OpenAI API Error ({model_id}): {e}")
+        return f"Error: OpenAI API call failed for {model_id}. Details: {e}"
+
+def _generate_with_gemini_provider(client, model_id, prompt, max_tokens, system_message=None):
+    full_prompt = f"{system_message}\n\n{prompt}" if system_message else prompt
+    try:
+        model = client.GenerativeModel(
+            model_id,
+            safety_settings={
+                'HARM_CATEGORY_HARASSMENT': 'block_none',
+                'HARM_CATEGORY_HATE_SPEECH': 'block_none',
+                'HARM_CATEGORY_SEXUALLY_EXPLICIT': 'block_none',
+                'HARM_CATEGORY_DANGEROUS_CONTENT': 'block_none',
+            },
+            generation_config=client.types.GenerationConfig(temperature=0.7)
+        )
+        response = model.generate_content(full_prompt)
+
+        if response.parts:
+             return "".join(part.text for part in response.parts if hasattr(part, 'text'))
+        elif hasattr(response, 'text') and response.text:
+             return response.text
+        elif response.prompt_feedback.block_reason:
+            reason = response.prompt_feedback.block_reason
+            st.warning(f"Gemini response blocked ({model_id}). Reason: {reason}")
+            return f"Error: Response blocked by API safety filters ({model_id}): {reason}"
+        else:
+            if response.candidates and response.candidates[0].finish_reason != "STOP":
+                 st.warning(f"Gemini generation stopped unexpectedly ({model_id}). Reason: {response.candidates[0].finish_reason}")
+                 return f"Error: Generation stopped unexpectedly ({model_id}). Reason: {response.candidates[0].finish_reason}"
+            else:
+                 st.warning(f"Gemini returned an empty or unexpected response ({model_id}).")
+                 return f"Error: Gemini returned an empty response for {model_id}."
+
+    except Exception as e:
+        st.error(f"❌ Gemini SDK error ({model_id}): {e}")
+        error_detail = getattr(e, 'message', str(e))
+        if "API key not valid" in error_detail:
+            return f"Error: Invalid Gemini API Key ({model_id}). Please check your Streamlit secrets."
+        return f"Error: Gemini SDK call failed for {model_id}. Details: {error_detail}"
+
+
+def _generate_with_deepseek_provider(api_key, model_id, prompt, max_tokens, system_message=None):
+    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+    messages = []
+    if system_message:
+        messages.append({"role": "system", "content": system_message})
+    messages.append({"role": "user", "content": prompt})
+    payload = {
+        "model": model_id,
+        "messages": messages,
+        "temperature": 0.6,
+        "max_tokens": max_tokens
+    }
+    try:
+        response = requests.post("https://api.deepseek.com/chat/completions", headers=headers, json=payload, timeout=90)
+        response.raise_for_status()
+        response_data = response.json()
+        if ("choices" in response_data and response_data["choices"] and
+            "message" in response_data["choices"][0] and
+            "content" in response_data["choices"][0]["message"]):
+            return response_data["choices"][0]["message"]["content"]
+        else:
+            st.warning(f"DeepSeek returned an unexpected response structure ({model_id}): {response_data}")
+            return f"Error: DeepSeek returned an unexpected structure for {model_id}."
+    except requests.exceptions.RequestException as e:
+        st.error(f"❌ DeepSeek API Request Error ({model_id}): {e}")
+        return f"Error: DeepSeek API request failed for {model_id}. Details: {e}"
+    except Exception as e:
+        st.error(f"❌ DeepSeek Error processing response ({model_id}): {e}")
+        return f"Error: DeepSeek processing failed for {model_id}. Details: {e}"
+
+def _generate_with_claude_provider(client, model_id, prompt, max_tokens, system_message=None):
+    try:
+        message = client.messages.create(
+            model=model_id,
+            max_tokens=max_tokens,
+            system=system_message if system_message else None,
+            messages=[
+                {"role": "user", "content": prompt}
+            ]
+        )
+        content = ""
+        if message.content:
+           content = "\n".join([block.text for block in message.content if hasattr(block, "text")])
+        return content
+    except Exception as e:
+        st.error(f"❌ Claude API Error ({model_id}): {e}")
+        if isinstance(e, anthropic.AuthenticationError):
+             return f"Error: Claude authentication failed ({model_id}). Check your API key."
+        return f"Error: Claude API call failed for {model_id}. Details: {e}"
 
-# Initialize OpenAI client
-openai_api_key = st.secrets.get("OPENAI_API_KEY")
-openai_client = OpenAI(api_key=openai_api_key) if openai_api_key else None
+def generate_with_selected_model(selected_model_name, prompt, max_tokens=2000, system_message=None):
+    """Generates text using the chosen model, handling provider specifics."""
+    if not any_secret_loaded or not SUPPORTED_MODELS:
+        st.error("Error: No API keys loaded or models available. Configure secrets.")
+        return None
+
+    if selected_model_name not in SUPPORTED_MODELS:
+        st.error(f"Selected model '{selected_model_name}' is not configured or unavailable.")
+        original_choice = selected_model_name
+        selected_model_name = DEFAULT_MODEL
+        if not selected_model_name:
+            st.error("Fatal: Default model is also unavailable. Cannot proceed.")
+            return None
+        st.warning(f"Falling back from '{original_choice}' to default model: {DEFAULT_MODEL}")
+        st.session_state.model_choice = DEFAULT_MODEL # Update state on fallback
+
+    model_config = SUPPORTED_MODELS[selected_model_name]
+    provider = model_config["provider"]
+    model_id = model_config["id"]
+    client = model_config.get("client")
+
+    st.info(f"Generating with: **{selected_model_name}**")
+    start_time = time.time()
+    result = f"Error: Provider '{provider}' not implemented."
 
-# ---------- Helper Functions ----------
-def generate_ai_response(prompt, max_tokens=2000):
-    if not openai_client:
-        st.error("OpenAI API key not configured.")
+    try:
+        if provider == "openai":
+            if not client: result = f"Error: OpenAI client not initialized for {selected_model_name}."
+            else: result = _generate_with_openai_provider(client, model_id, prompt, max_tokens, system_message)
+        elif provider == "gemini":
+            if not client: result = f"Error: Gemini client not initialized for {selected_model_name}."
+            else: result = _generate_with_gemini_provider(client, model_id, prompt, max_tokens, system_message)
+        elif provider == "deepseek":
+            if not deepseek_api_key: result = f"Error: DeepSeek API key not available for {selected_model_name}."
+            else: result = _generate_with_deepseek_provider(deepseek_api_key, model_id, prompt, max_tokens, system_message)
+        elif provider == "claude":
+            if not client: result = f"Error: Claude client not initialized for {selected_model_name}."
+            else: result = _generate_with_claude_provider(client, model_id, prompt, max_tokens, system_message)
+    except Exception as e:
+        st.error(f"❌ Unexpected error during generation with {selected_model_name}: {e}")
+        result = f"Error: Unexpected failure during {provider} generation. Details: {e}"
+
+    end_time = time.time()
+    duration = end_time - start_time
+    # st.caption(f"Generation took {duration:.2f} seconds.") # Less verbose
+
+    if isinstance(result, str) and result.startswith("Error:"):
+        # Error already logged by provider function
         return None
-    response = openai_client.chat.completions.create(
-        model="gpt-4-turbo",
-        messages=[{"role": "user", "content": prompt}],
-        max_tokens=max_tokens
+    return result
+
+# --- Mermaid Diagram Helper ---
+# (Keep your existing Mermaid helper functions here - unchanged)
+# ... (Omitted for brevity, assume they are the same as before) ...
+def is_valid_mermaid(code):
+    if not isinstance(code, str): return False
+    code_lower = code.strip().lower()
+    return bool(re.search(r"^\s*(%%.*?\n)*\s*(graph|flowchart|sequenceDiagram|classDiagram|stateDiagram|erDiagram|gantt|pie|gitGraph)", code_lower, re.MULTILINE))
+
+def render_mermaid_diagram(mermaid_code, key):
+    if not isinstance(mermaid_code, str) or not mermaid_code.strip():
+        st.warning(f"Mermaid code is empty or invalid (Key: {key}).")
+        return
+
+    cleaned_code = re.sub(r"^```mermaid\s*\n?", "", mermaid_code, flags=re.IGNORECASE | re.MULTILINE).strip()
+    cleaned_code = re.sub(r"\n?```\s*$", "", cleaned_code).strip()
+
+    if not is_valid_mermaid(cleaned_code):
+        st.warning(f"⚠️ Mermaid diagram might not render correctly (Key: {key}). Check syntax. Displaying raw code.")
+        st.code(cleaned_code, language="mermaid")
+        return
+
+    container_id = f"mermaid-container-{key}"
+    mermaid_id = f"mermaid-{key}"
+
+    components.html(
+        f"""
+        <div id="{container_id}" style="background-color: white; padding: 10px; border-radius: 5px; overflow: auto;">
+            <pre class="mermaid" id="{mermaid_id}">
+            {cleaned_code}
+            </pre>
+        </div>
+        <script type="module">
+            try {{
+                const mermaid = (await import('https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.esm.min.mjs')).default;
+                mermaid.initialize({{ startOnLoad: false, theme: 'default' }});
+                const checkElement = setInterval(() => {{
+                    const el = document.getElementById('{mermaid_id}');
+                    if (el) {{
+                        clearInterval(checkElement);
+                        mermaid.run({{ nodes: [el] }});
+                    }}
+                }}, 100);
+                setTimeout(() => clearInterval(checkElement), 5000);
+            }} catch (e) {{
+                console.error("Mermaid rendering error (Key: {key}):", e);
+                const container = document.getElementById('{container_id}');
+                if(container) container.innerHTML = "<p style='color:red;'>Error rendering Mermaid diagram. Check browser console.</p>";
+            }}
+        </script>
+        """,
+        height=500, scrolling=True,
     )
-    return response.choices[0].message.content
-
-# ---------- Session State Initialization ----------
-initial_state = {
-    'step': 1,
-    'idea': '',
-    'refined_idea': '',
-    'questions': [],
-    'answers': {},
-    'documents': [],
-    'docs_ready': {}
+
+# ---------- Initialize Session State (Provibe Workflow) ----------
+# Core workflow steps
+if 'current_step' not in st.session_state:
+    st.session_state.current_step = "input_idea" # input_idea -> refine_idea -> review_idea -> generate_docs -> display_docs
+if 'processing' not in st.session_state: # General flag for disabling buttons
+    st.session_state.processing = False
+
+# Input data
+if 'initial_product_idea' not in st.session_state:
+    st.session_state.initial_product_idea = ""
+if 'tech_stack_hint' not in st.session_state:
+    st.session_state.tech_stack_hint = ""
+if 'model_choice' not in st.session_state:
+    st.session_state.model_choice = DEFAULT_MODEL
+
+# Refinement stage data (NEW states added here)
+if 'refinement_sub_step' not in st.session_state:
+    # Tracks progress within the 'refine_idea' step
+    # Possible values: 'generate_questions', 'await_answers', 'generate_final_refinement'
+    st.session_state.refinement_sub_step = 'generate_questions'
+if 'clarifying_questions' not in st.session_state:
+    # Stores the list of questions generated by the AI
+    st.session_state.clarifying_questions = []
+if 'user_answers' not in st.session_state:
+    # Stores user answers, dictionary mapping question index to answer string
+    st.session_state.user_answers = {}
+
+# Output/Generated data
+if 'refined_idea_content' not in st.session_state: # Stores the final AI-refined idea (after Q&A)
+    st.session_state.refined_idea_content = None
+if 'confirmed_idea_content' not in st.session_state: # Stores the user-confirmed/edited idea
+    st.session_state.confirmed_idea_content = None
+if 'selected_docs_to_generate' not in st.session_state: # Stores user selection for optional docs
+    st.session_state.selected_docs_to_generate = {}
+if 'generated_docs' not in st.session_state: # Stores content of generated optional docs
+    st.session_state.generated_docs = {}
+
+
+# ---------- Define Document Options (Align with Provibe Output) ----------
+# (Keep your existing doc_options dictionary here - unchanged)
+# ... (Omitted for brevity, assume it's the same as before, including PRD option) ...
+doc_options = {
+    "prd": {
+        "label": "Product Requirements Document (PRD)",
+        "prompt_func": lambda idea, hint: f"""
+# --- PROMPT: Insert your specific PRD generation prompt here ---
+# Example: Write a comprehensive Product Requirements Document (PRD) based strictly on the following confirmed product description. Include sections like Introduction, Goals, Target Audience, Features (with details), User Stories, Design Considerations, Non-Functional Requirements, Open Issues, and Future Considerations. Ensure the PRD is detailed, clear, and actionable for a development team.
+# --- End PRD Prompt ---
+
+**Confirmed Product Description:**
+---
+{idea}
+---
+**Optional Preferences/Hints (Consider if relevant):**
+{hint if hint else "None provided"}
+""",
+        "system_message": "You are an expert Product Manager tasked with writing a detailed and professional PRD.",
+        "max_tokens": 3500, # Allow more tokens for PRD
+        "display_func": lambda content, key: st.markdown(content),
+        "download_filename": "prd.md",
+        "mime": "text/markdown",
+    },
+    "user_flow_text": {
+        "label": "User Flow (Text Description)",
+        "prompt_func": lambda idea, hint: f"""
+# --- PROMPT: Insert your specific User Flow (Text) generation prompt here ---
+# Example: Based on the product description below, outline the primary user flow step-by-step, from initial interaction to achieving the core goal. Describe each step clearly.
+# --- End User Flow (Text) Prompt ---
+
+**Product Description:**
+---
+{idea}
+---
+**Preferences/Hints:** {hint if hint else "None provided"}
+""",
+        "system_message": "You are a UX designer describing a key user journey.",
+        "max_tokens": 1000,
+        "display_func": lambda content, key: st.markdown(content),
+        "download_filename": "user_flow.md",
+        "mime": "text/markdown",
+    },
+    "user_flow_mermaid": {
+        "label": "User Flow Diagram (Mermaid)",
+        "prompt_func": lambda idea, hint: f"""
+# --- PROMPT: Insert your specific User Flow (Mermaid) generation prompt here ---
+# Example: Generate a Mermaid flowchart diagram representing the primary user flow for the product described below. Use standard flowchart syntax (graph TD, nodes, arrows). Ensure the diagram is clear and accurately reflects the user journey. Start the code block with ```mermaid and end it with ```. Do not include any other text before or after the code block.
+# --- End User Flow (Mermaid) Prompt ---
+
+**Product Description:**
+---
+{idea}
+---
+**Preferences/Hints:** {hint if hint else "None provided"}
+""",
+        "system_message": "You are an expert in creating Mermaid diagrams, specifically flowcharts for user journeys.",
+        "max_tokens": 1000,
+        "render_func": render_mermaid_diagram, # Special rendering
+        "code_language": "mermaid",
+        "download_filename": "user_flow_diagram.mmd",
+        "mime": "text/plain",
+    },
+    "frontend_arch": {
+        "label": "Frontend Architecture Notes",
+        "prompt_func": lambda idea, hint: f"""
+# --- PROMPT: Insert your Frontend Architecture prompt here ---
+# Example: Based on the product description and hints, suggest a suitable frontend architecture. Describe key components, recommended libraries/frameworks (consider hints like 'React Native'), state management approach, and potential component breakdown.
+# --- End Frontend Architecture Prompt ---
+
+**Product Description:**
+---
+{idea}
+---
+**Preferences/Hints:** {hint if hint else "None provided"}
+""",
+        "system_message": "You are a frontend architect designing a web/mobile application.",
+         "max_tokens": 1500,
+        "display_func": lambda content, key: st.markdown(content),
+        "download_filename": "frontend_architecture.md",
+        "mime": "text/markdown",
+    },
+    "backend_arch": {
+        "label": "Backend Architecture Notes",
+        "prompt_func": lambda idea, hint: f"""
+# --- PROMPT: Insert your Backend Architecture prompt here ---
+# Example: Based on the product description and hints, propose a backend architecture. Discuss potential API design (e.g., RESTful), choice of language/framework, database considerations (type, scaling), authentication/authorization strategy, and key microservices or modules if applicable.
+# --- End Backend Architecture Prompt ---
+
+**Product Description:**
+---
+{idea}
+---
+**Preferences/Hints:** {hint if hint else "None provided"}
+""",
+        "system_message": "You are a backend/systems architect designing the server-side logic and infrastructure.",
+         "max_tokens": 1500,
+        "display_func": lambda content, key: st.markdown(content),
+        "download_filename": "backend_architecture.md",
+        "mime": "text/markdown",
+    },
+     "system_arch_mermaid": {
+        "label": "System Architecture Diagram (Mermaid)",
+        "prompt_func": lambda idea, hint: f"""
+# --- PROMPT: Insert your System Architecture (Mermaid) prompt here ---
+# Example: Generate a Mermaid diagram illustrating the high-level system architecture for the product described below. Include key components like frontend client, backend API, database, authentication service, and any major third-party integrations mentioned or implied. Use appropriate Mermaid diagram syntax (e.g., graph TD or C4 model elements if suitable). Start the code block with ```mermaid and end it with ```. Do not include any other text before or after the code block.
+# --- End System Architecture (Mermaid) Prompt ---
+
+**Product Description:**
+---
+{idea}
+---
+**Preferences/Hints:** {hint if hint else "None provided"}
+""",
+        "system_message": "You create system architecture diagrams using Mermaid syntax.",
+        "max_tokens": 1000,
+        "render_func": render_mermaid_diagram,
+        "code_language": "mermaid",
+        "download_filename": "system_architecture.mmd",
+        "mime": "text/plain",
+    },
+    "db_schema": {
+        "label": "Database Schema (SQL)",
+        "prompt_func": lambda idea, hint: f"""
+# --- PROMPT: Insert your Database Schema (SQL) prompt here ---
+# Example: Based on the product description, design a preliminary relational database schema. Provide SQL `CREATE TABLE` statements for the primary entities, including relevant columns, data types, primary keys, and foreign key relationships. Assume a PostgreSQL syntax unless hints suggest otherwise.
+# --- End Database Schema (SQL) Prompt ---
+
+**Product Description:**
+---
+{idea}
+---
+**Preferences/Hints:** {hint if hint else "None provided"}
+""",
+        "system_message": "You are a database administrator designing a schema.",
+        "max_tokens": 1500,
+        "display_func": lambda content, key: st.code(content, language='sql'), # Use code block for SQL
+        "code_language": "sql",
+        "download_filename": "database_schema.sql",
+        "mime": "text/x-sql",
+    },
+    "project_structure": {
+        "label": "Project Folder Structure",
+        "prompt_func": lambda idea, hint: f"""
+# --- PROMPT: Insert your Project Structure prompt here ---
+# Example: Suggest a logical file and folder structure for a project implementing the described product. Consider frontend, backend, shared components, tests, etc., based on the description and any tech stack hints. Present it as a simple tree structure.
+# --- End Project Structure Prompt ---
+
+**Product Description:**
+---
+{idea}
+---
+**Preferences/Hints:** {hint if hint else "None provided"}
+""",
+        "system_message": "You are suggesting a clean project layout for a software development team.",
+        "max_tokens": 800,
+        "display_func": lambda content, key: st.code(content, language='bash'), # Use code block for structure
+        "code_language": "bash",
+        "download_filename": "project_structure.txt",
+        "mime": "text/plain",
+    },
 }
 
-for key, default in initial_state.items():
-    if key not in st.session_state:
-        st.session_state[key] = default
 
-# ---------- UI ----------
-st.title("🚀 Provibe - AI-Powered Product Documentation Generator")
+# ---------- UI Layout (Provibe Workflow) ----------
+st.set_page_config(layout="wide", page_title="Provibe Prompt Tester")
+st.title("🧪 Provibe Prompt Tester (with Q&A Refinement)")
+st.caption("Test and refine prompts for the Provibe document generation workflow, including interactive Q&A.")
 
-# Step-by-step Wizard
+# Display API Key Errors
+if secret_errors:
+    st.error("API Key Configuration Issues:")
+    for error in secret_errors:
+        st.error(f"- {error}")
+if not any_secret_loaded or not SUPPORTED_MODELS:
+    st.error("No API keys loaded or LLM models available. Configure secrets.")
+    st.stop()
 
-# Step 1: Idea Input
-if st.session_state.step == 1:
-    st.header("Step 1: Enter Your Product Idea")
-    with st.form("idea_form"):
-        idea_input = st.text_area("Describe your initial product idea:", height=150)
-        submit_idea = st.form_submit_button("Refine Idea")
+# --- Workflow Steps ---
 
-    if submit_idea and idea_input.strip():
-        st.session_state.idea = idea_input
-        st.session_state.step = 2
-        st.rerun()
+# ---------- Step 1: Input Initial Idea ----------
+if st.session_state.current_step == "input_idea":
+    st.header("Step 1: Input Product Idea")
+    with st.form(key="idea_form"):
+        initial_idea_input = st.text_area(
+            "💡 Enter the initial product idea:", height=150,
+            value=st.session_state.initial_product_idea,
+            help="The raw concept or description."
+        )
+        tech_hint_input = st.text_input(
+            "⚙️ Optional: Tech Stack Hints or Constraints",
+            placeholder="e.g., Use React, target mobile, needs offline support",
+            value=st.session_state.tech_stack_hint,
+            help="Any preferences to guide AI generation."
+        )
+        available_model_names = list(SUPPORTED_MODELS.keys())
+        default_model_key = st.session_state.get('model_choice', DEFAULT_MODEL)
+        default_index = available_model_names.index(default_model_key) if default_model_key in available_model_names else 0
+        model_choice_input = st.selectbox(
+            "🧠 Choose AI model for all steps:",
+            options=available_model_names,
+            index=default_index,
+            key="model_select",
+            help="This model will be used for refinement and document generation."
+        )
+        submit_idea_button = st.form_submit_button(
+            label="➡️ Start Interactive Refinement",
+            use_container_width=True,
+            disabled=st.session_state.processing
+        )
 
-# Step 2: Idea Refinement
-elif st.session_state.step == 2:
-    st.header("Step 2: AI-Refined Idea")
-    if not st.session_state.refined_idea:
-        with st.spinner("Refining your idea..."):
-            refined_prompt = f"Refine and clarify this product idea:\n\n{st.session_state.idea}"
-            st.session_state.refined_idea = generate_ai_response(refined_prompt, 1000)
+    if submit_idea_button and initial_idea_input:
+        st.session_state.initial_product_idea = initial_idea_input
+        st.session_state.tech_stack_hint = tech_hint_input
+        st.session_state.model_choice = model_choice_input
+        # Reset states for the refinement process
+        st.session_state.clarifying_questions = []
+        st.session_state.user_answers = {}
+        st.session_state.refined_idea_content = None
+        st.session_state.confirmed_idea_content = None
+        st.session_state.generated_docs = {}
+        st.session_state.selected_docs_to_generate = {}
+        # Set the next step and the initial sub-step for refinement
+        st.session_state.current_step = "refine_idea"
+        st.session_state.refinement_sub_step = "generate_questions"
+        st.session_state.processing = True # Start processing
+        safe_rerun()
+    elif submit_idea_button:
+        st.warning("Please enter a product idea.")
 
-    with st.form("refine_form"):
-        st.session_state.refined_idea = st.text_area("Refined Idea:", st.session_state.refined_idea, height=200)
-        submit_refined = st.form_submit_button("Proceed to Detailed Questions")
 
-    if submit_refined and st.session_state.refined_idea.strip():
-        st.session_state.step = 3
-        st.rerun()
+# ---------- Step 2: Interactive Refinement (Q&A) ----------
+if st.session_state.current_step == "refine_idea":
+    st.header("Step 2: Interactive Idea Refinement")
 
-# Step 3: AI-Generated Questions
-elif st.session_state.step == 3:
-    st.header("Step 3: Product Details")
-
-    if not st.session_state.questions:
-        with st.spinner("Generating critical questions..."):
-            questions_prompt = f"Generate 5 critical questions to clarify details for this product:\n\n{st.session_state.refined_idea}"
-            raw_questions = generate_ai_response(questions_prompt, 500)
-            st.session_state.questions = [q.strip("- ").strip() for q in raw_questions.strip().split("\n") if q.strip()]
-
-    with st.form("questions_form"):
-        all_answered = True
-        for idx, question in enumerate(st.session_state.questions):
-            answer = st.text_input(f"Q{idx+1}: {question}", key=f"question_{idx}")
-            st.session_state.answers[idx] = answer
-            if not answer.strip():
-                all_answered = False
-        submit_answers = st.form_submit_button("Generate Development Plan")
-
-    if submit_answers:
-        if all_answered:
-            st.session_state.step = 4
-            st.rerun()
-        else:
-            st.warning("Please answer all questions before proceeding.")
+    # --- Sub-Step 2a: Generate Clarifying Questions ---
+    if st.session_state.refinement_sub_step == "generate_questions":
+        st.info(f"Using **{st.session_state.model_choice}** to generate clarifying questions. Please wait.")
+        with st.spinner("AI is preparing questions..."):
 
-# Step 4: Development Plan Generation
-elif st.session_state.step == 4:
-    st.header("Step 4: AI-Generated Development Plan")
-    answers_formatted = "\n".join(f"{st.session_state.questions[i]}: {a}" for i, a in st.session_state.answers.items())
-    if 'development_plan' not in st.session_state:
-        with st.spinner("Creating your development plan..."):
-            detail_prompt = f"""
-            Create a concise, clear development plan for this product:
+            # --- PROMPT: Define the Question Generation Prompt ---
+            question_gen_prompt = f"""
+# --- PROMPT: Insert your Question Generation prompt here ---
+# Example: Based on the initial product idea and hints below, generate 3-5 specific clarifying questions for the user. These questions should help elicit more detail about key features, target audience, technical constraints, or core functionality needed to write a better product specification. Output *only* the questions, each on a new line, starting with '- '. Do not include numbering or any other text.
+# --- End Question Generation Prompt ---
 
-            Refined Idea:
-            {st.session_state.refined_idea}
+**Initial Product Idea:**
+---
+{st.session_state.initial_product_idea}
+---
+**Optional Preferences/Hints Provided:**
+{st.session_state.tech_stack_hint if st.session_state.tech_stack_hint else "None provided"}
+"""
+            # --- End Question Generation Prompt ---
 
-            Detailed Answers:
-            {answers_formatted}
-            """
-            st.session_state.development_plan = generate_ai_response(detail_prompt, 1500)
+            system_message_qa = "You are an AI assistant helping to clarify a product idea by asking relevant questions."
+            max_tokens_qa = 300
 
-    st.markdown(st.session_state.development_plan)
+            questions_raw = generate_with_selected_model(
+                st.session_state.model_choice,
+                question_gen_prompt,
+                max_tokens=max_tokens_qa,
+                system_message=system_message_qa
+            )
 
-    if st.button("Proceed to Document Generation"):
-        st.session_state.step = 5
-        st.rerun()
+            if questions_raw and not questions_raw.startswith("Error:"):
+                # Parse the questions (assuming one question per line, maybe starting with '- ')
+                st.session_state.clarifying_questions = [
+                    q.strip('- ') for q in questions_raw.strip().split('\n') if q.strip() and q.strip() != '-'
+                ]
+                if st.session_state.clarifying_questions:
+                    st.session_state.user_answers = {i: "" for i in range(len(st.session_state.clarifying_questions))} # Initialize empty answers
+                    st.session_state.refinement_sub_step = "await_answers"
+                else:
+                     st.warning("AI generated questions but they seem empty or incorrectly formatted. Proceeding without Q&A.")
+                     # Fallback: Skip Q&A and go directly to final refinement based only on initial idea
+                     st.session_state.refinement_sub_step = "generate_final_refinement"
+                     st.session_state.clarifying_questions = [] # Ensure it's empty
+                     st.session_state.user_answers = {}
+
+            else:
+                st.error("Failed to generate clarifying questions. Check API errors or model selection.")
+                st.warning("Proceeding to refine based only on the initial idea (skipping Q&A).")
+                st.session_state.refinement_sub_step = "generate_final_refinement" # Skip Q&A on failure
+                st.session_state.clarifying_questions = [] # Ensure it's empty
+                st.session_state.user_answers = {}
+
+            st.session_state.processing = False # Done generating questions (or failed)
+            safe_rerun()
+
+    # --- Sub-Step 2b: Display Questions and Collect Answers ---
+    elif st.session_state.refinement_sub_step == "await_answers":
+        st.info("Please answer the following questions to help refine the product idea:")
+        with st.form("answers_form"):
+            # Display generated questions and input fields for answers
+            for i, question in enumerate(st.session_state.clarifying_questions):
+                st.session_state.user_answers[i] = st.text_area(
+                    f"❓ {question}",
+                    key=f"answer_{i}",
+                    value=st.session_state.user_answers.get(i, ""), # Preserve answers on rerun
+                    height=100
+                )
+
+            submit_answers_button = st.form_submit_button(
+                "➡️ Submit Answers & Generate Refined Description",
+                use_container_width=True,
+                disabled=st.session_state.processing
+            )
+
+        if submit_answers_button:
+            # Basic check if answers are provided (optional)
+            # if not all(st.session_state.user_answers.values()):
+            #     st.warning("Please try to answer all questions for the best result.")
+            # else:
+            st.session_state.refinement_sub_step = "generate_final_refinement"
+            st.session_state.processing = True # Start final refinement generation
+            safe_rerun()
+
+        # Option to go back
+        if st.button("⬅️ Back to Idea Input (Discard Q&A)", disabled=st.session_state.processing):
+            st.session_state.current_step = "input_idea"
+            # Clear Q&A state
+            st.session_state.clarifying_questions = []
+            st.session_state.user_answers = {}
+            safe_rerun()
+
+
+    # --- Sub-Step 2c: Generate Final Refined Description (using Q&A) ---
+    elif st.session_state.refinement_sub_step == "generate_final_refinement":
+        st.info(f"Using **{st.session_state.model_choice}** to generate the final refined description based on the idea and your answers. Please wait.")
+        with st.spinner("AI is synthesizing the refined description..."):
+
+            # Prepare Q&A string for the prompt
+            qa_summary = "\n".join([
+                f"Q: {st.session_state.clarifying_questions[i]}\nA: {answer}"
+                for i, answer in st.session_state.user_answers.items() if answer # Include only answered questions
+            ]) if st.session_state.user_answers else "No questions were answered."
 
-# Step 5: Document Generation
-elif st.session_state.step == 5:
-    st.header("Step 5: Generate Detailed Documentation")
-    doc_types = ["Product Requirements Document", "User Flow", "System Architecture", "Database Schema"]
-    selected_docs = st.multiselect("Select documents to generate:", doc_types, default=doc_types)
-
-    if st.button("Generate Documents"):
-        st.session_state.documents = selected_docs
-        for doc in selected_docs:
-            with st.spinner(f"Generating {doc}..."):
-                doc_prompt = f"""
-                Generate a detailed {doc.lower()} for this product:
-
-                Idea:
-                {st.session_state.refined_idea}
-
-                Details:
-                {"; ".join(st.session_state.answers.values())}
-                """
-                content = generate_ai_response(doc_prompt, 2000)
-                st.session_state.docs_ready[doc] = content
-        st.success("All selected documents generated!")
-        st.session_state.step = 6
-        st.rerun()
 
-# Display generated documents
-if st.session_state.step == 6 and st.session_state.docs_ready:
-    st.header("📄 Generated Documents")
-    for doc, content in st.session_state.docs_ready.items():
-        st.subheader(doc)
-        st.markdown(content)
-        st.download_button(
-            f"Download {doc}",
-            data=content,
-            file_name=f"{doc.replace(' ', '_').lower()}.md",
-            mime="text/markdown",
-            key=f"download_{doc}"
+            # --- PROMPT: Define the Final Refinement Prompt (using Q&A) ---
+            final_refinement_prompt = f"""
+# --- PROMPT: Insert your Final Refinement prompt (using Q&A) here ---
+# Example: Based on the initial product idea, user preferences, and the following question-answer pairs, generate a concise yet comprehensive 'Refined Product Description'. Synthesize all the information into a well-structured description covering the core value proposition, key features, target audience, and any clarified technical aspects. This description will be the basis for generating all subsequent documents.
+# --- End Final Refinement Prompt ---
+
+**Initial Product Idea:**
+---
+{st.session_state.initial_product_idea}
+---
+**Optional Preferences/Hints Provided:**
+{st.session_state.tech_stack_hint if st.session_state.tech_stack_hint else "None provided"}
+---
+**Clarifying Questions & User Answers:**
+---
+{qa_summary}
+---
+"""
+            # --- End Final Refinement Prompt ---
+
+            system_message_final_refine = "You are an AI assistant synthesizing information into a final product specification."
+            max_tokens_final_refine = 1500 # Allow slightly more tokens for synthesis
+
+            final_refined_content = generate_with_selected_model(
+                st.session_state.model_choice,
+                final_refinement_prompt,
+                max_tokens=max_tokens_final_refine,
+                system_message=system_message_final_refine
+            )
+
+            if final_refined_content and not final_refined_content.startswith("Error:"):
+                st.session_state.refined_idea_content = final_refined_content
+                st.session_state.current_step = "review_idea" # Move to the next main step
+            else:
+                st.error("Failed to generate the final refined description.")
+                # Option to retry or go back might be added here
+                st.session_state.current_step = "input_idea" # Go back if failed
+
+            st.session_state.processing = False # End processing
+            safe_rerun()
+
+
+# ---------- Step 3: Review and Confirm Final Idea ----------
+if st.session_state.current_step == "review_idea":
+    st.header("Step 3: Review and Confirm Final Refined Idea")
+    if st.session_state.refined_idea_content:
+        st.info("Review the AI's final refined description below (generated using your answers). Edit it as needed. This **final text** will be used to generate all documents.")
+
+        # Display Q&A for context if available
+        if st.session_state.clarifying_questions and st.session_state.user_answers:
+             with st.expander("View Q&A used for this refinement"):
+                 for i, q in enumerate(st.session_state.clarifying_questions):
+                     st.markdown(f"**Q:** {q}")
+                     st.markdown(f"**A:** {st.session_state.user_answers.get(i, '_No answer_')}")
+                 st.markdown("---")
+
+
+        edited_idea = st.text_area(
+            "✏️ **Edit Final Refined Description:**",
+            value=st.session_state.refined_idea_content,
+            height=350,
+            key="final_refined_idea_edit_area",
+            help="Make any necessary corrections or additions."
         )
 
-    if st.button("🔄 Start New Idea"):
-        for key in initial_state:
-            st.session_state[key] = initial_state[key]
-        st.session_state.step = 1
-        st.rerun()
+        button_col1, button_col2 = st.columns(2)
+        with button_col1:
+            confirm_button = st.button(
+                "✅ Confirm & Proceed to Generate Docs",
+                key="confirm_final_idea_button",
+                use_container_width=True,
+                disabled=st.session_state.processing
+            )
+        with button_col2:
+             back_button = st.button(
+                 "⬅️ Back to Idea Input (Start Over)",
+                 key="back_to_input_final_button",
+                 use_container_width=True,
+                 disabled=st.session_state.processing
+             )
+
+        if confirm_button:
+            if not edited_idea.strip():
+                st.warning("The refined description cannot be empty.")
+            else:
+                st.session_state.confirmed_idea_content = edited_idea
+                # Reset generation states
+                st.session_state.generated_docs = {}
+                st.session_state.selected_docs_to_generate = {k: False for k in doc_options} # Reset selections
+                st.session_state.current_step = "generate_docs"
+                safe_rerun()
+        if back_button:
+             st.session_state.current_step = "input_idea"
+             # Clear refinement & Q&A state
+             st.session_state.refined_idea_content = None
+             st.session_state.clarifying_questions = []
+             st.session_state.user_answers = {}
+             safe_rerun()
+
+    else:
+        st.error("No refined idea content found. Please go back to Step 1.")
+        if st.button("⬅️ Back to Idea Input"):
+            st.session_state.current_step = "input_idea"
+            safe_rerun()
+
+
+# ---------- Step 4: Select and Generate Documents ----------
+if st.session_state.current_step == "generate_docs":
+    st.header("Step 4: Generate Product Documents")
+    if st.session_state.confirmed_idea_content:
+        st.markdown("**Based on this confirmed final description:**")
+        with st.expander("View Confirmed Description", expanded=False):
+             st.markdown(f"> {st.session_state.confirmed_idea_content}")
+
+        st.subheader("Select Documents to Generate:")
+        num_doc_options = len(doc_options)
+        cols = st.columns(min(num_doc_options, 3))
+        doc_keys = list(doc_options.keys())
+
+        for i, key in enumerate(doc_keys):
+            config = doc_options[key]
+            with cols[i % 3]:
+                if key not in st.session_state.selected_docs_to_generate:
+                     st.session_state.selected_docs_to_generate[key] = False
+                st.session_state.selected_docs_to_generate[key] = st.checkbox(
+                    config["label"],
+                    value=st.session_state.selected_docs_to_generate.get(key, False),
+                    key=f"checkbox_{key}",
+                    disabled=st.session_state.processing
+                )
+
+        generate_button = st.button(
+            "🚀 Generate Selected Documents",
+            key="generate_docs_button",
+            use_container_width=True,
+            disabled=st.session_state.processing
+        )
+        back_to_review_button = st.button(
+             "⬅️ Back to Review Final Idea",
+              key="back_to_review_final_button",
+              use_container_width=True,
+              disabled=st.session_state.processing
+        )
+
+
+        if generate_button:
+            selected_keys = [k for k, v in st.session_state.selected_docs_to_generate.items() if v]
+            if not selected_keys:
+                st.warning("Please select at least one document type to generate.")
+            else:
+                st.session_state.processing = True
+                st.session_state.generated_docs = {}
+                st.info(f"⏳ Generating {len(selected_keys)} selected document(s) using {st.session_state.model_choice}...")
+                progress_bar = st.progress(0)
+                generation_successful = True
+
+                for i, key in enumerate(selected_keys):
+                    config = doc_options[key]
+                    st.write(f"   - Generating {config['label']}...")
+                    with st.spinner(f"AI processing {config['label']}..."):
+                        prompt = config["prompt_func"](
+                            st.session_state.confirmed_idea_content,
+                            st.session_state.tech_stack_hint
+                        )
+                        system_msg = config.get("system_message")
+                        max_tok = config.get("max_tokens", 2000)
+
+                        content = generate_with_selected_model(
+                            st.session_state.model_choice,
+                            prompt,
+                            max_tokens=max_tok,
+                            system_message=system_msg
+                        )
+
+                        if content and not content.startswith("Error:"):
+                            st.session_state.generated_docs[key] = content
+                        else:
+                            st.session_state.generated_docs[key] = f"Error: Failed to generate {config['label']}."
+                            generation_successful = False
+                            st.error(f"   - Failed to generate {config['label']}. See logs above.")
+                    progress_bar.progress((i + 1) / len(selected_keys))
+                    time.sleep(0.1)
+
+                progress_bar.empty()
+                st.session_state.processing = False
+
+                if generation_successful:
+                    st.success("✅ Document generation complete!")
+                else:
+                    st.warning("⚠️ Some documents could not be generated.")
+
+                st.session_state.current_step = "display_docs"
+                safe_rerun()
+
+        if back_to_review_button:
+            st.session_state.current_step = "review_idea"
+            # Keep confirmed idea, but allow editing again
+            # The refined_idea_content should still hold the content before editing
+            safe_rerun()
+
+    else:
+        st.error("Confirmed idea content is missing. Please restart the process from Step 1.")
+        if st.button("⬅️ Restart Process"):
+            # Reset key states
+            st.session_state.current_step = "input_idea"
+            st.session_state.initial_product_idea = ""
+            st.session_state.tech_stack_hint = ""
+            st.session_state.refined_idea_content = None
+            st.session_state.confirmed_idea_content = None
+            st.session_state.clarifying_questions = []
+            st.session_state.user_answers = {}
+            st.session_state.generated_docs = {}
+            st.session_state.selected_docs_to_generate = {}
+            safe_rerun()
+
+
+# ---------- Step 5: Display Generated Documents ----------
+if st.session_state.current_step == "display_docs":
+    st.header("Step 5: Generated Documents")
+
+    if not st.session_state.generated_docs:
+         st.info("No documents were generated in the previous step.")
+    else:
+        st.markdown("**Review the generated documents below:**")
+        display_order = [key for key in doc_options if key in st.session_state.generated_docs]
+
+        for key in display_order:
+            content = st.session_state.generated_docs.get(key)
+            if content:
+                config = doc_options[key]
+                st.subheader(f"📄 {config['label']}")
+                is_error = isinstance(content, str) and content.startswith("Error:")
+
+                if is_error:
+                    st.error(content)
+                else:
+                    # Display/Render content
+                    if config.get("render_func"):
+                        try: config["render_func"](content, key=f"render_{key}")
+                        except Exception as e: st.error(f"Render Error: {e}"); st.code(content)
+                    elif config.get("display_func"):
+                        try: config["display_func"](content, key=f"display_{key}")
+                        except Exception as e: st.error(f"Display Error: {e}"); st.text(content)
+                    else: st.markdown(content)
+
+                    # Download button
+                    try:
+                        download_data = content.encode('utf-8') if isinstance(content, str) else str(content).encode('utf-8')
+                        st.download_button(
+                            label=f"📥 Download {config['label']}", data=download_data,
+                            file_name=config["download_filename"], mime=config.get("mime", "text/plain"),
+                            key=f"download_{key}"
+                        )
+                    except Exception as e: st.warning(f"Download Error: {e}")
+
+                    # Show raw content option
+                    if config.get("render_func") or config.get("code_language"):
+                         if st.checkbox(f"🔍 Show raw content for {config['label']}", key=f"show_raw_{key}", value=False):
+                              st.code(content, language=config.get("code_language", None))
+
+                st.markdown("---")
+
+    # Navigation buttons
+    button_col1, button_col2 = st.columns(2)
+    with button_col1:
+        generate_more_button = st.button("🔄 Generate Different Documents", key="generate_more_button", use_container_width=True)
+    with button_col2:
+        restart_all_button = st.button("⏪ Start New Idea", key="restart_all_button", use_container_width=True)
+
+    if generate_more_button:
+         st.session_state.current_step = "generate_docs"
+         st.session_state.generated_docs = {}
+         safe_rerun()
+
+    if restart_all_button:
+        st.session_state.current_step = "input_idea"
+        st.session_state.initial_product_idea = ""
+        st.session_state.tech_stack_hint = ""
+        st.session_state.refined_idea_content = None
+        st.session_state.confirmed_idea_content = None
+        st.session_state.clarifying_questions = []
+        st.session_state.user_answers = {}
+        st.session_state.generated_docs = {}
+        st.session_state.selected_docs_to_generate = {}
+        safe_rerun()
+
+
+# ---------- Footer ----------
+st.markdown("---")
+footer_model_choice = st.session_state.get('model_choice', 'N/A')
+st.caption(f"Using model: **{footer_model_choice}** | Workflow Step: **{st.session_state.get('current_step', 'N/A')}**"
+           f"{' (Sub-step: ' + st.session_state.get('refinement_sub_step', 'N/A') + ')' if st.session_state.get('current_step') == 'refine_idea' else ''}")
 
-st.caption("Powered by Provibe AI 🚀")