Update main.py

main.py

@@ -2,7 +2,7 @@ import gradio as gr
 from sqlalchemy.exc import SQLAlchemyError
 from utils import InnovativeIdea, init_db, get_db, SessionLocal, get_llm_response
 from data_models import IdeaForm
-from chatbot import InnovativeIdeaChatbot
+from chatbot import InnovativeIdeaChatbot 
 from config import MODELS, DEFAULT_SYSTEM_PROMPT, STAGES
 import logging
 import os
@@ -90,12 +90,12 @@ def create_gradio_interface():
         return chat_history, *[form_data.get(stage["field"], "") for stage in STAGES]
 
     def start_over():
-        chat_history, form_data = innovative_chatbot.start_over()
+        chat_history, form_data, initial_stage = innovative_chatbot.start_over()
         return (
             chat_history,  # Update the chatbot with the new chat history
             "",  # Clear the message input
             *[form_data.get(stage["field"], "") for stage in STAGES],  # Reset all form fields
-            gr.update(value=default_stage)  # Reset the stage selection
+            gr.update(value=initial_stage)  # Reset the stage selection
         )
 
     with gr.Blocks(theme=gr.themes.Soft()) as demo: