Update chatbot.py

chatbot.py

@@ -138,10 +138,18 @@ Let's start by exploring your innovative idea! What's the name of your idea, or
             setattr(self.idea_form, stage["field"], form_data[stage["field"]])
         
         # Save to database
-        if self.idea_id:
-            update_idea_in_database(self.idea_id, self.idea_form, self.db)
-        else:
-            self.idea_id = save_idea_to_database(self.idea_form, self.db)
+        try:
+            new_session = SessionLocal()
+            if self.idea_id:
+                update_idea_in_database(self.idea_id, self.idea_form, new_session)
+            else:
+                self.idea_id = save_idea_to_database(self.idea_form, new_session)
+            new_session.commit()
+        except Exception as e:
+            logging.error(f"Error saving idea to database: {str(e)}")
+            new_session.rollback()
+        finally:
+            new_session.close()
 
         return form_data
 
@@ -157,15 +165,60 @@ Let's start by exploring your innovative idea! What's the name of your idea, or
         self.idea_id = None
         self.current_stage = None
         self.add_system_message(self.get_initial_greeting())
-        clear_database(self.db)
+        try:
+            new_session = SessionLocal()
+            clear_database(new_session)
+            new_session.commit()
+        except Exception as e:
+            logging.error(f"Error clearing database: {str(e)}")
+            new_session.rollback()
+        finally:
+            new_session.close()
         return self.chat_history, self.idea_form.dict()
 
+    def start_over(self):
+        self.chat_history = []
+        self.idea_form = IdeaForm()
+        self.current_stage = None
+        self.add_system_message(self.get_initial_greeting())
+        
+        try:
+            new_session = SessionLocal()
+            # Clear the existing database
+            clear_database(new_session)
+            
+            # Create a new empty idea
+            new_idea = InnovativeIdea()
+            new_session.add(new_idea)
+            new_session.commit()
+            new_session.refresh(new_idea)
+            
+            # Update the idea_id
+            self.idea_id = new_idea.id
+            
+            new_session.close()
+        except Exception as e:
+            logging.error(f"Error in start_over: {str(e)}")
+            if 'new_session' in locals():
+                new_session.rollback()
+                new_session.close()
+        
+        return self.chat_history, self.idea_form.dict(), STAGES[0]["name"]
+
     def update_idea_form(self, stage_name: str, form_data: str):
         setattr(self.idea_form, stage_name.lower().replace(" ", "_"), form_data)
-        if self.idea_id:
-            update_idea_in_database(self.idea_id, self.idea_form, self.db)
-        else:
-            self.idea_id = save_idea_to_database(self.idea_form, self.db)
+        try:
+            new_session = SessionLocal()
+            if self.idea_id:
+                update_idea_in_database(self.idea_id, self.idea_form, new_session)
+            else:
+                self.idea_id = save_idea_to_database(self.idea_form, new_session)
+            new_session.commit()
+        except Exception as e:
+            logging.error(f"Error updating idea form: {str(e)}")
+            new_session.rollback()
+        finally:
+            new_session.close()
 
     def generate_form_data(self, stage: str, model: str, thinking_budget: int) -> str:
         # Prepare the conversation history for the LLM
@@ -204,23 +257,6 @@ Let's start by exploring your innovative idea! What's the name of your idea, or
         
         return form_data.get(stage, "")
 
-    def start_over(self):
-        self.reset()
-        clear_database(self.db)
-        self.idea_id = None
-        
-        # Reinitialize with a new empty idea
-        new_idea = InnovativeIdea()
-        self.db.add(new_idea)
-        self.db.commit()
-        self.idea_id = new_idea.id
-        
-        self.idea_form = IdeaForm()
-        self.chat_history = []
-        self.add_system_message(self.get_initial_greeting())
-        
-        return self.chat_history, self.idea_form.dict()
-
     def process_stage_input_stream(self, stage_name: str, message: str, model: str, system_prompt: str, thinking_budget: int):
         if self.current_stage != stage_name:
             activation_message = self.activate_stage(stage_name)
@@ -281,10 +317,18 @@ Let's start by exploring your innovative idea! What's the name of your idea, or
             setattr(self.idea_form, stage.field, form_data[stage.name])
         
         # Save to database
-        if self.idea_id:
-            update_idea_in_database(self.idea_id, self.idea_form, self.db)
-        else:
-            self.idea_id = save_idea_to_database(self.idea_form, self.db)
+        try:
+            new_session = SessionLocal()
+            if self.idea_id:
+                update_idea_in_database(self.idea_id, self.idea_form, new_session)
+            else:
+                self.idea_id = save_idea_to_database(self.idea_form, new_session)
+            new_session.commit()
+        except Exception as e:
+            logging.error(f"Error saving idea to database: {str(e)}")
+            new_session.rollback()
+        finally:
+            new_session.close()
 
     def generate_form_data_stream(self, stage: str, model: str, thinking_budget: int):
         conversation = "\n".join([f"{role}: {message}" for role, message in self.chat_history])