Update app.py

app.py

@@ -1,925 +1,164 @@
-# -*- coding: utf-8 -*-
+# Provibe App (v1.1) - Streamlit Implementation
+
 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
-
-# ---------- 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 = []
-
-# 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)
-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)
-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):
-    # Note: Gemini API v1 doesn't explicitly support system_message like OpenAI.
-    # It's often prepended to the user prompt or handled via specific model tuning.
-    # We will prepend it for basic compatibility.
-    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', # More permissive for testing
-                '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) # max_output_tokens isn't directly controllable here for safety reasons
-        )
-        # Gemini requires specific handling for potentially blocked content.
-        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: # Handle simpler text responses if applicable
-             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:
-             # Handle cases where the response might be empty but not explicitly blocked
-             # Check candidate information if available
-            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))
-        # Attempt to parse specific Gemini API errors if possible
-        if "API key not valid" in error_detail:
-            return f"Error: Invalid Gemini API Key ({model_id}). Please check your Streamlit secrets."
-        # Add more specific error handling if needed
-        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:
-        # Anthropic API v3 uses system parameter directly
-        message = client.messages.create(
-            model=model_id,
-            max_tokens=max_tokens,
-            system=system_message if system_message else None, # Add system message here
-            messages=[
-                {"role": "user", "content": prompt}
-            ]
-        )
-        # Handle response content which is a list of blocks
-        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}")
-        # Provide more specific feedback if it's an authentication error
-        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}"
-
-
-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 # Return None on critical failure
-
-    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 # Return None on critical failure
-        st.warning(f"Falling back from '{original_choice}' to default model: {DEFAULT_MODEL}")
-        # Update session state if fallback occurs
-        st.session_state.model_choice = DEFAULT_MODEL
-
-    model_config = SUPPORTED_MODELS[selected_model_name]
-    provider = model_config["provider"]
-    model_id = model_config["id"]
-    client = model_config.get("client") # Might be None for DeepSeek
-
-    st.info(f"Generating with: **{selected_model_name}**")
-    start_time = time.time()
-    result = f"Error: Provider '{provider}' not implemented." # Default error
-
-    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.")
-
-    # Check if the result indicates an error before returning
-    if isinstance(result, str) and result.startswith("Error:"):
-        st.error(f"Generation failed for {selected_model_name}. Check logs or API status.")
-        return None # Explicitly return None for errors
-    return result
-
-# --- Mermaid Diagram Helper ---
-def is_valid_mermaid(code):
-    if not isinstance(code, str): return False
-    code_lower = code.strip().lower()
-    # Simplified check for common diagram types
-    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
-
-    # Clean potentially markdown-formatted code
-    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
-
-    # Ensure unique IDs for multiple diagrams
-    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 {{
-                // Ensure mermaid is loaded (using import for modern browsers)
-                const mermaid = (await import('https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.esm.min.mjs')).default;
-                mermaid.initialize({{ startOnLoad: false, theme: 'default' }});
-                // Wait for the element to be ready in the DOM
-                const checkElement = setInterval(() => {{
-                    const el = document.getElementById('{mermaid_id}');
-                    if (el) {{
-                        clearInterval(checkElement);
-                        mermaid.run({{ nodes: [el] }});
-                    }}
-                }}, 100);
-                // Timeout to prevent infinite loop
-                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,
+# 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
+
+# ---------- Helper Functions ----------
+def generate_ai_response(prompt, max_tokens=2000):
+    if not openai_client:
+        st.error("OpenAI API key not configured.")
+        return None
+    response = openai_client.chat.completions.create(
+        model="gpt-4-turbo",
+        messages=[{"role": "user", "content": prompt}],
+        max_tokens=max_tokens
     )
-
-# ---------- Initialize Session State (Provibe Workflow) ----------
-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 to disable buttons during AI calls
-    st.session_state.processing = False
-if 'initial_product_idea' not in st.session_state:
-    st.session_state.initial_product_idea = ""
-if 'tech_stack_hint' not in st.session_state: # Keep tech hint if useful for prompts
-    st.session_state.tech_stack_hint = ""
-if 'model_choice' not in st.session_state:
-    st.session_state.model_choice = DEFAULT_MODEL
-if 'refined_idea_content' not in st.session_state: # Stores the AI-refined idea
-    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 'prd_content' not in st.session_state: # Stores generated PRD
-    st.session_state.prd_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) ----------
-# Key: unique identifier, used in session state
-# label: Display name for checkbox
-# prompt_func: Lambda function to generate the specific prompt for this doc type.
-#              Takes confirmed_idea and tech_hint as input.
-# display_func: Optional function to display content (e.g., st.markdown). Defaults to st.markdown.
-# download_filename: Filename for download button.
-# mime: MIME type for download.
-# render_func: Optional function for special rendering (e.g., Mermaid).
-# code_language: Optional language hint for st.code display.
-# system_message: Optional system message specific to this document type.
-
-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",
-    },
-    # Add other document types as needed, following the same structure
+    return response.choices[0].message.content
+
+# ---------- Session State Initialization ----------
+initial_state = {
+    'step': 1,
+    'idea': '',
+    'refined_idea': '',
+    'questions': [],
+    'answers': {},
+    'documents': [],
+    'docs_ready': {}
 }
 
+for key, default in initial_state.items():
+    if key not in st.session_state:
+        st.session_state[key] = default
 
-# ---------- UI Layout (Provibe Workflow) ----------
-st.set_page_config(layout="wide", page_title="Provibe Prompt Tester")
-st.title("🧪 Provibe Prompt Tester")
-st.caption("Test and refine prompts for the Provibe document generation workflow.")
-
-# 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. Please configure necessary Streamlit secrets (e.g., OPENAI_API_KEY, GEMINI_API_KEY). Cannot proceed.")
-    st.stop() # Stop execution if no models are usable
-
-# --- Workflow Steps ---
-
-# ---------- 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", # Ensure a unique key
-            help="This model will be used for refinement and document generation."
-        )
-        submit_idea_button = st.form_submit_button(
-            label="➡️ Refine Idea with AI",
-            use_container_width=True,
-            disabled=st.session_state.processing
-        )
-
-    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
-        st.session_state.refined_idea_content = None # Clear previous refinement
-        st.session_state.confirmed_idea_content = None # Clear previous confirmation
-        st.session_state.prd_content = None
-        st.session_state.generated_docs = {}
-        st.session_state.current_step = "refine_idea"
-        st.session_state.processing = True # Start processing
-        safe_rerun()
-    elif submit_idea_button:
-        st.warning("Please enter a product idea.")
+# ---------- UI ----------
+st.title("🚀 Provibe - AI-Powered Product Documentation Generator")
 
+# Step-by-step Wizard
 
-# ---------- Step 2: AI Refinement ----------
-if st.session_state.current_step == "refine_idea":
-    st.header("Step 2: Refining Idea...")
-    st.info(f"Using **{st.session_state.model_choice}** to refine the idea. Please wait.")
-    with st.spinner("AI is thinking..."):
+# 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")
 
-        # --- PROMPT: Define the Idea Refinement Prompt ---
-        refinement_prompt = f"""
-# --- PROMPT: Insert your specific Idea Refinement prompt here ---
-# Example: Analyze the following product idea and technical hints. Generate a concise, well-structured 'Refined Product Description' that clarifies the core value proposition, key features, and target audience. This description will be the basis for generating all other documents. Focus on clarity and completeness for a development team.
-# --- End Idea Refinement Prompt ---
+    if submit_idea and idea_input.strip():
+        st.session_state.idea = idea_input
+        st.session_state.step = 2
+        st.rerun()
 
-**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 Refinement Prompt ---
+# 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)
 
-        system_message_refine = "You are an AI assistant helping to refine a product idea into a clear specification."
-        max_tokens_refine = 1200
+    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")
 
-        refined_content = generate_with_selected_model(
-            st.session_state.model_choice,
-            refinement_prompt,
-            max_tokens=max_tokens_refine,
-            system_message=system_message_refine
-        )
+    if submit_refined and st.session_state.refined_idea.strip():
+        st.session_state.step = 3
+        st.rerun()
 
-        if refined_content:
-            st.session_state.refined_idea_content = refined_content
-            st.session_state.current_step = "review_idea"
+# 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.error("Failed to refine the idea. Check API errors or model selection.")
-            # Option to go back or retry 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 Idea ----------
-if st.session_state.current_step == "review_idea":
-    st.header("Step 3: Review and Confirm Refined Idea")
-    if st.session_state.refined_idea_content:
-        st.info("Review the AI-refined description below. Edit it as needed. This **final text** will be used to generate all documents.")
-        edited_idea = st.text_area(
-            "✏️ **Edit Refined Description:**",
-            value=st.session_state.refined_idea_content,
-            height=350,
-            key="refined_idea_edit_area",
-            help="Make any necessary corrections or additions."
-        )
-
-        button_col1, button_col2 = st.columns(2)
-        with button_col1:
-            confirm_button = st.button(
-                "✅ Confirm & Proceed to Generate Docs",
-                key="confirm_idea_button",
-                use_container_width=True,
-                disabled=st.session_state.processing
-            )
-        with button_col2:
-             back_button = st.button(
-                 "⬅️ Back to Idea Input",
-                 key="back_to_input_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.prd_content = None
-                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"
-             st.session_state.refined_idea_content = None # Clear refinement
-             safe_rerun()
+            st.warning("Please answer all questions before proceeding.")
 
-    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: 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:
 
+            Refined Idea:
+            {st.session_state.refined_idea}
 
-# ---------- 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 description:**")
-        with st.expander("View Confirmed Description", expanded=False):
-             st.markdown(f"> {st.session_state.confirmed_idea_content}") # Display quoted
+            Detailed Answers:
+            {answers_formatted}
+            """
+            st.session_state.development_plan = generate_ai_response(detail_prompt, 1500)
 
-        st.subheader("Select Documents to Generate:")
-        # Use columns for better layout of checkboxes
-        num_doc_options = len(doc_options)
-        cols = st.columns(min(num_doc_options, 3)) # Max 3 columns
-        doc_keys = list(doc_options.keys())
+    st.markdown(st.session_state.development_plan)
 
-        for i, key in enumerate(doc_keys):
-            config = doc_options[key]
-            with cols[i % 3]:
-                # Ensure state exists for each checkbox
-                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 Idea",
-              key="back_to_review_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 = {} # Clear previous results before generating new ones
-                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']}..."):
-                        # Get prompt using the lambda function
-                        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) # Use specific or default max_tokens
-
-                        content = generate_with_selected_model(
-                            st.session_state.model_choice,
-                            prompt,
-                            max_tokens=max_tok,
-                            system_message=system_msg
-                        )
-
-                        if content:
-                            st.session_state.generated_docs[key] = content
-                        else:
-                            # Store error message if generation failed
-                            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) # Small delay for UI update
-
-                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() # Rerun to display results
-
-        if back_to_review_button:
-            st.session_state.current_step = "review_idea"
-            # Keep confirmed idea, but allow editing again
-            st.session_state.refined_idea_content = st.session_state.confirmed_idea_content
-            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.prd_content = None
-            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:**")
-        # Sort keys based on the order in doc_options for consistent display
-        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: # Should always be true if in display_order, but check anyway
-                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:
-                    # Use specific render/display functions if available
-                    if config.get("render_func"):
-                        try:
-                            config["render_func"](content, key=f"render_{key}")
-                        except Exception as render_e:
-                             st.error(f"Error rendering {config['label']}: {render_e}")
-                             st.code(content, language=config.get("code_language", None)) # Fallback to code view
-                    elif config.get("display_func"):
-                        try:
-                            config["display_func"](content, key=f"display_{key}")
-                        except Exception as display_e:
-                             st.error(f"Error displaying {config['label']}: {display_e}")
-                             st.text(content) # Fallback to text view
-                    else: # Default to markdown
-                        st.markdown(content)
-
-                    # Add download button for non-error content
-                    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 download_e:
-                        st.warning(f"Could not prepare download for {config['label']}: {download_e}")
-
-                    # Optional: Show raw code for rendered/code types
-                    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))
+    if st.button("Proceed to Document Generation"):
+        st.session_state.step = 5
+        st.rerun()
 
-                st.markdown("---") # Separator between documents
+# 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()
 
-    # Navigation buttons at the end
-    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
+# 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}"
         )
 
-    if generate_more_button:
-         st.session_state.current_step = "generate_docs" # Go back to selection
-         st.session_state.generated_docs = {} # Clear displayed docs
-         safe_rerun()
-
-    if restart_all_button:
-        # Full reset
-        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.prd_content = None
-        st.session_state.generated_docs = {}
-        st.session_state.selected_docs_to_generate = {}
-        safe_rerun()
-
+    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()
 
-# ---------- 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')}**")
+st.caption("Powered by Provibe AI 🚀")