Update app.py

app.py

@@ -1,695 +1,111 @@
+# Provibe App (v1.0) - Streamlit Implementation
+
 import streamlit as st
-import requests
 from openai import OpenAI
-import google.generativeai as genai
-import anthropic  # Import the Anthropic client library for Claude
-import streamlit.components.v1 as components
-import re
-# from supabase import create_client  # Keep commented if not using history/auth
-import base64  # Needed for download buttons
-import time   # For potential delays/retries if needed
-
-# ---------- Helper: Safe Rerun ----------
-def safe_rerun():
-    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 ----------
-openai_client = None
-genai_client = None  # For Google GenAI
-deepseek_api_key = None
-claude_client = None  # For Anthropic Claude
-secrets_available = {"openai": False, "gemini": False, "deepseek": False, "claude": False}
-secret_errors = []  # Store errors for display
-
-# 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  # Store the configured module
-        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 using Anthropic library
-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}")
-
-# Check if at least one API key is loaded
-any_secret_loaded = any(secrets_available.values())
-
-# ---------- Model Configuration ----------
-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},
-        "ChatGPT o1 (OpenAI)": {"id": "chatgpt-o1", "provider": "openai", "client": openai_client},
-        "ChatGPT o3 (OpenAI)": {"id": "chatgpt-o3", "provider": "openai", "client": openai_client},
-    })
-
-# Gemini Models (Google)
-if secrets_available["gemini"] and genai_client:
-    SUPPORTED_MODELS.update({
-        "Gemini 2.0 (Google)": {"id": "gemini-2.0-latest", "provider": "gemini", "client": genai_client},
-        "Gemini 2.5 (Google)": {"id": "gemini-2.5-latest", "provider": "gemini", "client": genai_client},
-        "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},
-    })
+import time
 
-# Claude Models (Anthropic)
-if secrets_available["claude"] and claude_client:
-    SUPPORTED_MODELS.update({
-        "Claude 3.7 Sonnet (Anthropic)": {"id": "claude-3-7-sonnet-20250219", "provider": "claude", "client": claude_client},
-        "Claude 3.5 Haiku (Anthropic)": {"id": "claude-3-5-haiku-20241022", "provider": "claude", "client": claude_client},
-    })
+# 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
 
-DEFAULT_MODEL_PREFERENCE = [
-    "GPT-4o Mini (OpenAI)",
-    "Gemini 1.5 Flash (Google)",
-    "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))
-
-# ---------- Helper Functions for Generation ----------
-
-def _generate_with_openai_provider(client, model_id, prompt, max_tokens):
-    try:
-        response = client.chat.completions.create(
-            model=model_id,
-            messages=[
-                {"role": "system", "content": "You are an expert product manager and software architect. Be detailed and thorough."},
-                {"role": "user", "content": prompt}
-            ],
-            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):
-    try:
-        model = client.GenerativeModel(
-            model_id,
-            safety_settings={
-                'HARM_CATEGORY_HARASSMENT': 'block_medium_and_above',
-                'HARM_CATEGORY_HATE_SPEECH': 'block_medium_and_above',
-                'HARM_CATEGORY_SEXUALLY_EXPLICIT': 'block_medium_and_above',
-                'HARM_CATEGORY_DANGEROUS_CONTENT': 'block_medium_and_above',
-            },
-            generation_config=client.types.GenerationConfig(temperature=0.7)
-        )
-        response = model.generate_content(prompt)
-        if hasattr(response, 'text') and response.text:
-            return response.text
-        elif response.parts:
-            return "".join(part.text for part in response.parts if hasattr(part, 'text'))
-        elif response.prompt_feedback.block_reason:
-            st.warning(f"Gemini response blocked ({model_id}). Reason: {response.prompt_feedback.block_reason}")
-            return f"Response blocked by API safety filters ({model_id}): {response.prompt_feedback.block_reason}"
-        else:
-            st.warning(f"Gemini returned an empty or unexpected response ({model_id}). Response: {response}")
-            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))
-        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):
-    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
-    payload = {
-        "model": model_id,
-        "messages": [
-            {"role": "system", "content": "You are an expert product manager and software architect. Be detailed and thorough."},
-            {"role": "user", "content": prompt}
-        ],
-        "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):
-    try:
-        message = client.messages.create(
-            model=model_id,
-            max_tokens=max_tokens,
-            messages=[
-                {"role": "user", "content": prompt}
-            ]
-        )
-        content = message.content
-        # If content is a list of TextBlock objects, extract and join their text
-        if isinstance(content, list):
-            content = "\n".join([block.text for block in content if hasattr(block, "text")])
-        elif not isinstance(content, str):
-            content = str(content)
-        return content
-    except Exception as e:
-        st.error(f"❌ Claude API Error ({model_id}): {e}")
-        return f"Error: Claude API call failed for {model_id}. Details: {e}"
-
-def generate_with_selected_model(selected_model_name, prompt, max_tokens=2000):
-    if not any_secret_loaded or not SUPPORTED_MODELS:
-        return "Error: No API keys loaded or models available."
-    if selected_model_name not in SUPPORTED_MODELS:
-        st.error(f"Selected model '{selected_model_name}' is not configured or unavailable.")
-        selected_model_name = DEFAULT_MODEL
-        if not selected_model_name:
-            return "Error: Default model also unavailable."
-        st.warning(f"Falling back to default model: {DEFAULT_MODEL}")
-
-    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 = None
-    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)
-    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)
-    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)
-    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)
-    else:
-        st.error(f"Unknown provider '{provider}' configured for model '{selected_model_name}'.")
-        result = f"Error: Unknown provider {provider}."
-    end_time = time.time()
-    duration = end_time - start_time
-    # st.caption(f"Generation took {duration:.2f} seconds.")
-
-    if isinstance(result, str) and result.startswith("Error:"):
+# ---------- Helper Functions ----------
+def generate_ai_response(prompt, max_tokens=2000):
+    if not openai_client:
+        st.error("OpenAI API key not configured.")
         return None
-    return result
-
-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).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
-
-    components.html(
-        f"""
-        <div id="mermaid-container-{key}" style="background-color: white; padding: 10px; border-radius: 5px; overflow: auto;">
-          <pre class="mermaid" id="mermaid-{key}">
-          {cleaned_code}
-          </pre>
-        </div>
-        <script type="module">
-          import mermaid from 'https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.esm.min.mjs';
-          try {{
-            mermaid.initialize({{ startOnLoad: false, theme: 'default' }});
-            const el = document.getElementById('mermaid-{key}');
-            if (el) {{
-               await mermaid.run({{ nodes: [el] }});
-            }} else {{
-               console.error("Mermaid target element not found: mermaid-{key}");
-            }}
-          }} catch (e) {{
-            console.error("Mermaid rendering error (Key: {key}):", e);
-          }}
-        </script>
-        """,
-        height=500, scrolling=True,
-    )
-
-# ---------- Initialize Session State ----------
-if 'elaboration_pending' not in st.session_state:
-    st.session_state.elaboration_pending = False
-if 'elaboration_complete' not in st.session_state:
-    st.session_state.elaboration_complete = False
-if 'generating_prd' not in st.session_state:
-    st.session_state.generating_prd = False
-if 'generating_optional' not in st.session_state:
-    st.session_state.generating_optional = 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:
-    st.session_state.tech_stack_hint = ""
-if 'model_choice' not in st.session_state:
-    st.session_state.model_choice = DEFAULT_MODEL
-if 'elaborated_idea_raw' not in st.session_state:
-    st.session_state.elaborated_idea_raw = None
-if 'confirmed_product_idea' not in st.session_state:
-    st.session_state.confirmed_product_idea = None
-if 'prd_content' not in st.session_state:
-    st.session_state.prd_content = None
-if 'selected_docs_to_generate' not in st.session_state:
-    st.session_state.selected_docs_to_generate = {}
-if 'generated_docs' not in st.session_state:
-    st.session_state.generated_docs = {}
-
-# ---------- Define Document Options ----------
-doc_options = {
-    "frontend": {
-        "label": "Frontend Architecture",
-        "prompt": lambda idea, hint: f"Based on the following refined product description, describe the frontend tech stack and architecture.\n\nProduct Description:\n{idea}\n\nConsider these preferences if relevant: {hint}",
-        "display_func": lambda content, key: st.markdown(content),
-        "download_filename": "frontend_architecture.md",
-        "mime": "text/markdown",
-        "render_func": None,
-        "code_language": None,
-    },
-    "backend": {
-        "label": "Backend Architecture",
-        "prompt": lambda idea, hint: f"Based on the following refined product description, describe the backend system design.\n\nProduct Description:\n{idea}\n\nConsider these preferences if relevant: {hint}",
-        "display_func": lambda content, key: st.markdown(content),
-        "download_filename": "backend_architecture.md",
-        "mime": "text/markdown",
-        "render_func": None,
-        "code_language": None,
-    },
-    "project_structure": {
-        "label": "Project Folder Structure",
-        "prompt": lambda idea, hint: f"""
-Based on the following refined product description and any technical hints provided, generate a suggested file and folder structure for the project... [Same prompt as before] ...
-
-**Confirmed Product Description:**
----
-{idea}
----
-**Optional Preferences/Hints Provided:**
-{hint if hint else "None provided"}
-""",
-        "display_func": None,
-        "download_filename": "project_structure.txt",
-        "mime": "text/plain",
-        "render_func": None,
-        "code_language": "bash",
-    },
-    "user_flow": {
-        "label": "User Flow (Text)",
-        "prompt": lambda idea, hint: f"Based on the following refined product description, describe the primary user flow.\n\nProduct Description:\n{idea}\n\nConsider these preferences if relevant: {hint}",
-        "display_func": lambda content, key: st.markdown(content),
-        "download_filename": "user_flow.md",
-        "mime": "text/markdown",
-        "render_func": None,
-        "code_language": None,
-    },
-    "user_flow_mermaid": {
-        "label": "User Flow Diagram (Mermaid)",
-        "prompt": lambda idea, hint: f"Based on the following refined product description, create a user flow diagram in Mermaid.js flowchart format... [Same prompt as before] ...\n\nProduct Description:\n{idea}\n\nConsider these preferences if relevant: {hint}",
-        "display_func": None,
-        "download_filename": "user_flow_diagram.mmd",
-        "mime": "text/plain",
-        "render_func": render_mermaid_diagram,
-        "code_language": "mermaid",
-    },
-    "system_arch_mermaid": {
-        "label": "System Architecture Diagram (Mermaid)",
-        "prompt": lambda idea, hint: f"Based on the following refined product description, create a system architecture diagram in Mermaid.js format... [Same prompt as before] ...\n\nProduct Description:\n{idea}\n\nConsider these preferences if relevant: {hint}",
-        "display_func": None,
-        "download_filename": "system_architecture.mmd",
-        "mime": "text/plain",
-        "render_func": render_mermaid_diagram,
-        "code_language": "mermaid",
-    },
-    "db_schema": {
-        "label": "Database Schema",
-        "prompt": lambda idea, hint: f"Based on the following refined product description, design a relational database schema... [Same prompt as before] ...\n\nProduct Description:\n{idea}\n\nConsider these preferences if relevant: {hint}",
-        "display_func": lambda content, key: st.code(content, language='sql'),
-        "download_filename": "database_schema.sql",
-        "mime": "text/x-sql",
-        "render_func": None,
-        "code_language": "sql",
-    },
-    "sql_queries": {
-        "label": "SQL Queries",
-        "prompt": lambda idea, hint: f"Based on the following refined product description and likely schema, write 5-10 sample SQL CRUD queries.\n\nProduct Description:\n{idea}\n\nConsider these preferences if relevant: {hint}",
-        "display_func": None,
-        "download_filename": "sample_queries.sql",
-        "mime": "text/x-sql",
-        "render_func": None,
-        "code_language": "sql",
-    },
-}
-
-# ---------- UI Layout ----------
-st.set_page_config(layout="wide", page_title="Idea to PRD Generator")
-st.title("🚀 Product Idea → Refinement → PRD & Docs Generator")
-st.caption("Enter your initial idea, let the AI refine it, edit if needed, then generate detailed documents.")
-
-# Display any errors loading secrets prominently at the top
-if secret_errors:
-    st.error("API Key Configuration Issues Found:")
-    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. Cannot proceed.")
-    st.stop()
-
-# --- Columns for Side-by-Side Layout ---
-col1, col2 = st.columns(2)
-
-# ---------- Step 1: Input Initial Idea (in Column 1) ----------
-with col1:
-    st.header("1. Define Your Initial Idea")
-    with st.form(key="idea_form"):
-        initial_idea_input = st.text_area(
-            "💡 Enter your raw product idea here:", height=180,
-            value=st.session_state.initial_product_idea,
-            help="Describe your app concept, target users, and core goals."
-        )
-        tech_hint_input = st.text_input(
-            "⚙️ Optional: Tech Stack or Specific Preferences",
-            placeholder="e.g., Use React Native, focus on scalability",
-            value=st.session_state.tech_stack_hint,
-            help="Mention any preferred technologies or constraints."
-        )
-        available_model_names = list(SUPPORTED_MODELS.keys())
-        default_index = 0
-        current_model_choice = st.session_state.get('model_choice', DEFAULT_MODEL)
-        if current_model_choice in available_model_names:
-            default_index = available_model_names.index(current_model_choice)
-        elif DEFAULT_MODEL in available_model_names:
-            default_index = available_model_names.index(DEFAULT_MODEL)
-            st.session_state.model_choice = DEFAULT_MODEL
-        model_choice_input = st.selectbox(
-            "🧠 Choose AI model:",
-            options=available_model_names,
-            index=default_index,
-            key="model_choice_select",
-            help="Select the AI model to use for all generation steps."
-        )
-        analyze_button = st.form_submit_button(
-            label="1️⃣ Analyze & Refine Idea",
-            disabled=st.session_state.elaboration_pending or st.session_state.generating_prd,
-            use_container_width=True
-        )
-    # --- Update session state when the analyze button is pressed ---
-    if analyze_button:
-        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.elaboration_pending = True
-        safe_rerun()
-
-# ---------- Step 2: Review and Confirm Elaborated Idea (in Column 2) ----------
-with col2:
-    st.header("2. Review & Confirm Refined Idea")
-    if not st.session_state.elaboration_complete and not st.session_state.elaboration_pending:
-        st.info("The refined idea and analysis will appear here after you click 'Analyze & Refine Idea'.")
-
-    if st.session_state.elaboration_pending:
-        with st.spinner(f"🧠 AI ({st.session_state.model_choice}) is analyzing..."):
-            elaboration_prompt = f"""
-Analyze the following product idea meticulously... [USE THE SAME ELABORATION PROMPT AS BEFORE - includes Optimized Description, Suggested Platform(s), Key Features, Target Audience, Initial Technical Considerations] ...
----
-**Original 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"}
----
-Important: Base your analysis and suggestions *only* on the 'Original Idea' and 'Optional Preferences/Hints'.
-"""
-            max_tokens_elaboration = 1200
-            st.session_state.elaborated_idea_raw = generate_with_selected_model(
-                st.session_state.model_choice, elaboration_prompt, max_tokens=max_tokens_elaboration
-            )
-            st.session_state.elaboration_pending = False
-            if st.session_state.elaborated_idea_raw:
-                st.session_state.elaboration_complete = True
-            else:
-                st.error("Failed to elaborate on the idea. Check API errors above or model availability.")
-                st.session_state.elaboration_complete = False
-            safe_rerun()
-
-    if st.session_state.elaboration_complete:
-        st.info("Review and edit the AI's analysis below. This text will be used to generate all documents.")
-        edited_idea = st.text_area(
-            "✏️ **Edit Refined Description & Analysis:**",
-            value=st.session_state.elaborated_idea_raw,
-            height=450,
-            key="elaborated_idea_edit_area",
-            help="Modify the description, features, platform, etc. This text becomes the basis for the PRD."
-        )
-        if st.button("2️⃣ Confirm & Generate PRD", key="confirm_prd_button", disabled=st.session_state.generating_prd, use_container_width=True):
-            if not edited_idea.strip():
-                st.warning("The refined description cannot be empty.")
-            else:
-                st.session_state.confirmed_product_idea = edited_idea
-                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}
-                st.session_state.generating_prd = True
-                st.session_state.elaboration_complete = False
-                st.session_state.generating_optional = False
-                safe_rerun()
-
-st.markdown("---")
-
-# ---------- PRD Generation Logic (Below columns) ----------
-if st.session_state.generating_prd and st.session_state.confirmed_product_idea:
-    st.header("3. Generating Product Requirements Document (PRD)")
-    st.info(f"Generating PRD using {st.session_state.model_choice}...")
-    with st.spinner("Please wait... This might take a moment."):
-        prd_prompt = f"""
-Write a comprehensive and detailed Product Requirements Document (PRD) based *strictly* on the provided 'Confirmed Product Description'... [USE THE SAME OPTIMIZED PRD PROMPT AS BEFORE]...
-
-**Confirmed Product Description:**
----
-{st.session_state.confirmed_product_idea}
----
-**Optional Preferences/Hints (Consider if relevant):**
-{st.session_state.tech_stack_hint if st.session_state.tech_stack_hint else "None provided"}
-"""
-        prd_content_result = generate_with_selected_model(
-            st.session_state.model_choice, prd_prompt, max_tokens=3500
-        )
-        if prd_content_result:
-            st.session_state.prd_content = prd_content_result
-        else:
-            st.session_state.prd_content = None
-            st.error("Failed to generate PRD. Please check API errors or model availability.")
-        st.session_state.generating_prd = False
-        safe_rerun()
-
-# ---------- Display PRD and Offer Optional Docs ----------
-if st.session_state.prd_content:
-    st.header("4. Product Requirements Document (PRD)")
-    st.markdown(st.session_state.prd_content)
-    # Convert prd_content to a string if it's a list
-    prd_text = st.session_state.prd_content
-    if isinstance(prd_text, list):
-        prd_text = "\n".join(prd_text)
-    st.download_button(
-        label="📥 Download PRD",
-        data=prd_text.encode('utf-8'),
-        file_name="prd.md",
-        mime="text/markdown",
-        key="download_prd"
+    response = openai_client.chat.completions.create(
+        model="gpt-4-turbo",
+        messages=[{"role": "user", "content": prompt}],
+        max_tokens=max_tokens
     )
+    return response.choices[0].message.content
+
+# ---------- Session State Initialization ----------
+for key in ['step', 'idea', 'refined_idea', 'answers', 'documents', 'status']:
+    if key not in st.session_state:
+        st.session_state[key] = None
+if 'docs_ready' not in st.session_state:
+    st.session_state['docs_ready'] = {}
+
+# ---------- UI ----------
+st.title("🚀 Provibe - AI-Powered Product Documentation Generator")
+
+# Step-by-step Wizard
+
+# Step 1: Idea Input
+if st.session_state['step'] is None or st.session_state['step'] == 1:
+    st.header("Step 1: Enter Your Product Idea")
+    idea_input = st.text_area("Describe your initial product idea:", height=150)
+    if st.button("Refine Idea") and idea_input.strip():
+        st.session_state['idea'] = idea_input
+        st.session_state['step'] = 2
+        st.rerun()
 
-    st.markdown("---")
-    st.header("5. Generate Additional Documents (Optional)")
-    st.write("Select documents to generate based on the **confirmed product description**.")
-    num_options = len(doc_options)
-    num_cols = 3
-    cols = st.columns(num_cols)
-    doc_keys = list(doc_options.keys())
-    for i, key in enumerate(doc_keys):
-        config = doc_options[key]
-        with cols[i % num_cols]:
-            is_checked = st.session_state.selected_docs_to_generate.get(key, False)
-            st.session_state.selected_docs_to_generate[key] = st.checkbox(
-                config["label"],
-                value=is_checked,
-                key=f"checkbox_{key}",
-                disabled=st.session_state.generating_optional
-            )
+# Step 2: Idea Refinement
+if st.session_state['step'] == 2:
+    st.header("Step 2: AI-Refined Idea")
+    with st.spinner("Refining your idea..."):
+        refined_prompt = f"Refine and enhance the clarity of this product idea:\n\n{st.session_state['idea']}"
+        refined = generate_ai_response(refined_prompt, 1000)
+        st.session_state['refined_idea'] = st.text_area("Refined Idea:", refined, height=200)
+    if st.button("Proceed to Detailed Questions"):
+        st.session_state['step'] = 3
+        st.rerun()
 
-    if st.button("3️⃣ Generate Selected Documents", key="generate_optional_button", disabled=st.session_state.generating_optional, use_container_width=True):
-        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:
-            for key in selected_keys:
-                if key in st.session_state.generated_docs:
-                    st.session_state.generated_docs[key] = None
-            st.session_state.generating_optional = True
-            safe_rerun()
+# Step 3: AI-Generated Questions
+if st.session_state['step'] == 3:
+    st.header("Step 3: Product Details")
+    questions_prompt = f"Generate 5 critical questions to detail the product idea:\n\n{st.session_state['refined_idea']}"
+    questions = generate_ai_response(questions_prompt, 500).split('\n')
+    answers = []
+    with st.form("detail_form"):
+        for i, q in enumerate(questions):
+            answer = st.text_input(f"{q.strip()}")
+            answers.append(answer)
+        submitted = st.form_submit_button("Generate Development Plan")
+        if submitted and all(answers):
+            st.session_state['answers'] = answers
+            st.session_state['step'] = 4
+            st.rerun()
+
+# Step 4: Development Plan Generation
+if st.session_state['step'] == 4:
+    st.header("Step 4: AI-Generated Development Plan")
+    with st.spinner("Creating your development plan..."):
+        detail_prompt = f"Create a concise development plan based on:\nIdea: {st.session_state['refined_idea']}\nDetails: {st.session_state['answers']}"
+        dev_plan = generate_ai_response(detail_prompt, 1500)
+        st.markdown(dev_plan)
+    if st.button("Generate Documentation"):
+        st.session_state['step'] = 5
+        st.rerun()
 
-# ---------- Optional Docs Generation Logic ----------
-if st.session_state.generating_optional and st.session_state.confirmed_product_idea:
-    selected_keys_to_generate = [k for k, v in st.session_state.selected_docs_to_generate.items() if v and st.session_state.generated_docs.get(k) is None]
-    if selected_keys_to_generate:
-        st.info(f"⏳ Generating selected documents using {st.session_state.model_choice}...")
-        progress_bar = st.progress(0)
-        total_selected = len(selected_keys_to_generate)
-        generation_successful = True
-        for i, key in enumerate(selected_keys_to_generate):
-            config = doc_options[key]
-            st.write(f"Generating {config['label']}... ({i+1}/{total_selected})")
-            with st.spinner(f"Thinking about {config['label']}..."):
-                prompt = config["prompt"](st.session_state.confirmed_product_idea, st.session_state.tech_stack_hint)
-                content = generate_with_selected_model(st.session_state.model_choice, prompt, max_tokens=2000)
-                if content:
-                    st.session_state.generated_docs[key] = content
-                else:
-                    st.session_state.generated_docs[key] = f"Error generating {config['label']}."
-                    generation_successful = False
-                    st.error(f"Failed to generate {config['label']}.")
-            progress_bar.progress((i + 1) / total_selected)
-            time.sleep(0.1)
-        st.session_state.generating_optional = False
-        progress_bar.empty()
-        if generation_successful:
-            st.success("✅ Selected documents generated!")
-        else:
-            st.warning("Some documents could not be generated.")
-        safe_rerun()
-    else:
-        st.session_state.generating_optional = False
-        safe_rerun()
+# Step 5: Document Generation
+if st.session_state['step'] == 5:
+    st.header("Step 5: Document Generation")
+    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 Selected Documents"):
+        st.session_state['documents'] = selected_docs
+        st.session_state['status'] = {doc: False for doc in selected_docs}
+        for doc in selected_docs:
+            with st.spinner(f"Generating {doc}..."):
+                doc_prompt = f"Generate a detailed {doc.lower()} for the following product:\n\n{st.session_state['refined_idea']}\nDetails: {st.session_state['answers']}"
+                st.session_state['docs_ready'][doc] = generate_ai_response(doc_prompt, 2000)
+                st.session_state['status'][doc] = True
+        st.rerun()
 
-# ---------- Display Generated Optional Documents ----------
-if any(st.session_state.generated_docs.values()):
-    st.markdown("---")
-    st.header("6. Generated Documents")
-    for key in doc_options.keys():
-        content = st.session_state.generated_docs.get(key)
+# Display generated documents
+if st.session_state['docs_ready']:
+    st.header("Generated Documents")
+    for doc, content in st.session_state['docs_ready'].items():
         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)
-            elif config["render_func"]:
-                config["render_func"](content, key=f"render_{key}")
-            elif config["code_language"]:
-                st.code(content, language=config["code_language"])
-            elif config["display_func"]:
-                config["display_func"](content, key=f"display_{key}")
-            else:
-                st.markdown(content)
-            if not is_error:
-                try:
-                    st.download_button(
-                        label=f"📥 Download {config['label']}",
-                        data=content.encode('utf-8') if isinstance(content, str) else "\n".join(content).encode('utf-8'),
-                        file_name=config["download_filename"],
-                        mime=config["mime"],
-                        key=f"download_{key}"
-                    )
-                except Exception as e:
-                    st.warning(f"Could not prepare download: {e}")
-                if config["render_func"] and config["code_language"] == "mermaid":
-                    if st.checkbox(f"🔍 Show raw Mermaid code", key=f"show_raw_mermaid_{key}"):
-                        cleaned_code_raw = re.sub(r"^```mermaid\s*\n?", "", content, flags=re.IGNORECASE).strip()
-                        cleaned_code_raw = re.sub(r"\n?```\s*$", "", cleaned_code_raw).strip()
-                        st.code(cleaned_code_raw, language="mermaid")
-                elif config["code_language"] and not config["render_func"]:
-                    if st.checkbox(f"🔍 Show raw {config['code_language']} code", key=f"show_raw_code_{key}"):
-                        st.code(content, language=config["code_language"])
-            st.markdown("---")
+            st.subheader(doc)
+            st.markdown(content)
+            st.download_button(
+                f"Download {doc}",
+                data=content,
+                file_name=f"{doc.replace(' ', '_').lower()}.md",
+                mime="text/markdown"
+            )
 
-# ---------- Footer/Info ----------
-st.markdown("---")
-footer_model_choice = st.session_state.get('model_choice', 'N/A')
-st.caption(f"Using model: **{footer_model_choice}**. Remember to review all generated content carefully.")
+st.caption("Powered by Provibe AI")