Update chatbot.py

chatbot.py

@@ -1,230 +1,382 @@
-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 config import MODELS, DEFAULT_SYSTEM_PROMPT, STAGES
+from typing import List, Tuple, Dict, Any, Optional
 import logging
-import os
-from typing import Dict, Any 
 import re
-import time
+from data_models import IdeaForm, IDEA_STAGES
+from config import DEFAULT_SYSTEM_PROMPT, STAGES
+from utils import (
+    get_llm_response, extract_form_data, save_idea_to_database,
+    load_idea_from_database, update_idea_in_database,
+    get_db, clear_database, init_db, create_tables,
+    perform_web_search, optimize_search_query,
+    SessionLocal, InnovativeIdea
+)
 
-# Set up logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+class InnovativeIdeaChatbot:
+    def __init__(self):
+        create_tables()
+        init_db()
+        self.idea_form = IdeaForm()
+        self.chat_history = []
+        self.idea_id = None
+        self.current_stage = None
+        self.api_key = None
+        self.add_system_message(self.get_initial_greeting())
 
-def create_gradio_interface():
-    innovative_chatbot = InnovativeIdeaChatbot()
-    default_stage = STAGES[0]["name"]
+    def get_initial_greeting(self) -> str:
+        greeting = """
+Welcome to the Innovative Idea Generator! I'm Myamoto, your AI assistant designed to help you refine and develop your innovative ideas.
+Here's how we'll work together:
+1. We'll go through 10 stages to explore different aspects of your idea.
+2. At each stage, I'll ask you questions and provide feedback to help you think deeper about your concept.
+3. You can ask me questions at any time or request more information on a topic.
+4. If you want to perform a web search for additional information, just start your message with '@' followed by your search query.
+5. When you're ready to move to the next stage, simply type 'next'.
+Let's start by exploring your innovative idea! What's the name of your idea, or would you like help coming up with one?
+"""
+        self.greeted = True
+        return greeting
 
-    # Initialize the database
-    try:
-        init_db()
-        db = next(get_db())
-        initial_idea = db.query(InnovativeIdea).first()
-        if initial_idea is None:
-            logging.info("No initial idea found in the database. Creating a new one.")
-            initial_idea = InnovativeIdea()
-            db.add(initial_idea)
-            db.commit()
-            db.refresh(initial_idea)
-        
-        # Create form_fields while the session is still open
-        form_fields = {
-            stage["name"]: gr.Textbox(
-                label=stage["question"],
-                placeholder=stage["example"],
-                value=getattr(initial_idea, stage["field"], ""),
-                visible=(stage["name"] == default_stage),
-                interactive=False
-            ) for stage in STAGES
-        }
-        
-        # Now we can safely close the session
-        db.close()
-    except SQLAlchemyError as e:
-        logging.error(f"Database initialization failed: {str(e)}")
-        raise RuntimeError(f"Failed to initialize database: {str(e)}")
-
-    def chatbot_function(message, history, model, system_prompt, thinking_budget, current_stage):
+    def add_system_message(self, message: str):
+        self.chat_history.append(("System", message))
+
+    def set_api_key(self, api_key: str):
+        self.api_key = api_key
+
+    def activate_stage(self, stage_name: str) -> Optional[str]:
+        self.current_stage = stage_name
+        for stage in STAGES:
+            if stage["name"] == stage_name:
+                return f"Let's work on the '{stage_name}' stage. {stage['question']}"
+        return None
+
+    def process_stage_input(self, stage_name: str, message: str, model: str, system_prompt: str, thinking_budget: int) -> Tuple[List[Tuple[str, str]], Dict[str, Any]]:
+        if self.current_stage != stage_name:
+            activation_message = self.activate_stage(stage_name)
+            if activation_message is None:
+                error_message = f"Error: Unable to activate stage '{stage_name}'. Please check if the stage name is correct."
+                self.chat_history.append(("System", error_message))
+                return self.chat_history, self.idea_form.dict()
+            self.chat_history.append(("System", activation_message))
+
+        # Check for web search request
+        if message.startswith('@'):
+            search_query = message[1:].strip()
+            optimized_query = optimize_search_query(search_query, model)
+            search_results = perform_web_search(optimized_query)
+            self.chat_history.append(("Human", message))
+            self.chat_history.append(("AI", f"Here are the search results for '{optimized_query}':\n\n{search_results}"))
+            return self.chat_history, self.idea_form.dict()
+
+        # Generate the prompt for the current stage
+        stage_prompt = self.generate_prompt_for_stage(stage_name)
+
+        # Use the DEFAULT_SYSTEM_PROMPT from config.py
+        formatted_system_prompt = DEFAULT_SYSTEM_PROMPT.format(
+            current_stage=stage_name,
+            stage_prompt=stage_prompt
+        )
+
+        # Combine the formatted system prompt and user's input
+        combined_prompt = f"{formatted_system_prompt}\n\nUser input: {message}"
+
+        # Get LLM response
+        llm_response = get_llm_response(combined_prompt, model, thinking_budget, self.api_key)
+
+        # Parse the LLM response to extract only the user-facing content
+        parsed_response = self.parse_llm_response(llm_response)
+
+        # Add the interaction to chat history
+        self.chat_history.append(("Human", message))
+        self.chat_history.append(("AI", parsed_response))
+
+        # Extract form data from the LLM response
+        form_data = extract_form_data(llm_response)
+        
+        # Update the idea form
+        if stage_name in form_data:
+            setattr(self.idea_form, stage_name.lower().replace(" ", "_"), form_data[stage_name])
+
+        return self.chat_history, self.idea_form.dict()
+
+    def parse_llm_response(self, response: str) -> str:
+        # Remove content within <form_data> tags
+        response = re.sub(r'<form_data>.*?</form_data>', '', response, flags=re.DOTALL)
+        
+        # Remove content within <reflection> tags
+        response = re.sub(r'<reflection>.*?</reflection>', '', response, flags=re.DOTALL)
+        
+        # Remove content within <analysis> tags
+        response = re.sub(r'<analysis>.*?</analysis>', '', response, flags=re.DOTALL)
+        
+        # Remove content within <summary> tags
+        response = re.sub(r'<summary>.*?</summary>', '', response, flags=re.DOTALL)
+        
+        # Remove content within <step> tags
+        response = re.sub(r'<step>.*?</step>', '', response, flags=re.DOTALL)
+        
+        # Remove any remaining HTML-like tags
+        response = re.sub(r'<[^>]+>', '', response)
+        
+        # Remove extra whitespace and newlines
+        response = re.sub(r'\s+', ' ', response).strip()
+        
+        return response
+
+    def fill_out_form(self, current_stage: str, model: str, thinking_budget: int) -> Dict[str, str]:
+        form_data = {}
+        for stage in STAGES:
+            stage_name = stage["name"]
+            if stage_name == current_stage:
+                # Generate new data for the current stage
+                form_data[stage["field"]] = self.generate_form_data(stage_name, model, thinking_budget)
+            else:
+                # Use existing data for other stages
+                form_data[stage["field"]] = getattr(self.idea_form, stage["field"], "")
+        
+        # Update the idea form
+        for stage in STAGES:
+            setattr(self.idea_form, stage["field"], form_data[stage["field"]])
+        
+        # Save to database
         try:
-            # If this is the first message, get the initial greeting
-            if not history:
-                initial_greeting = innovative_chatbot.get_initial_greeting()
-                history.append((None, initial_greeting))
-                yield history, "", ""
-                return
+            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
 
-            for partial_response in innovative_chatbot.process_stage_input_stream(current_stage, message, model, system_prompt, thinking_budget):
-                chat_history, form_data = partial_response
-                history.append((message, chat_history[-1][1]))
-                yield history, form_data.get(current_stage, ""), ""
+    def generate_prompt_for_stage(self, stage: str) -> str:
+        for s in IDEA_STAGES:
+            if s.name == stage:
+                return f"We are currently working on the '{stage}' stage. {s.question}"
+        return f"We are currently working on the '{stage}' stage. Please provide relevant information."
 
-            # Update the database with the new form data
-            db = SessionLocal()
-            idea = db.query(InnovativeIdea).first()
-            for key, value in form_data.items():
-                setattr(idea, key, value)
-            db.commit()
-            db.close()
+    def reset(self):
+        self.chat_history = []
+        self.idea_form = IdeaForm()
+        self.idea_id = None
+        self.current_stage = None
+        self.add_system_message(self.get_initial_greeting())
+        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"An error occurred in chatbot_function: {str(e)}", exc_info=True)
-            yield history + [(None, f"An error occurred: {str(e)}")], "", ""
-
-    def fill_form(stage, model, thinking_budget):
-        form_data = innovative_chatbot.fill_out_form(stage, model, thinking_budget)
-        return [form_data.get(stage["field"], "") for stage in STAGES]
-
-    def clear_chat():
-        # Reset the database to an empty form
-        db = SessionLocal()
-        idea = db.query(InnovativeIdea).first()
-        empty_form = IdeaForm()
-        for key, value in empty_form.dict().items():
-            setattr(idea, key, value)
-        db.commit()
-        db.close()
-        
-        chat_history, form_data = innovative_chatbot.reset()
-        return chat_history, *[form_data.get(stage["field"], "") for stage in STAGES]
-
-    def 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=initial_stage)  # Reset the stage selection
+            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)
+        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
+        conversation = "\n".join([f"{role}: {message}" for role, message in self.chat_history])
+        
+        stage_prompt = self.generate_prompt_for_stage(stage)
+        
+        formatted_system_prompt = DEFAULT_SYSTEM_PROMPT.format(
+            current_stage=stage,
+            stage_prompt=stage_prompt
         )
 
-    with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# Innovative Idea Generator")
-        
-        mode = gr.Radio(["Chatbot", "Direct Input"], label="Mode", value="Chatbot")
-        
-        with gr.Row():
-            with gr.Column(scale=2):
-                chatbot = gr.Chatbot(label="Conversation", height=500)
-                msg = gr.Textbox(label="Your input", placeholder="Type your brilliant idea here...")
-                
-                with gr.Row():
-                    submit = gr.Button("Submit")
-                    clear = gr.Button("Clear Chat")
-                    start_over_btn = gr.Button("Start Over")
-
-            with gr.Column(scale=1):
-                stages = gr.Radio(
-                    choices=[stage["name"] for stage in STAGES],
-                    label="Ideation Stages",
-                    value=default_stage
-                )
-                form_fields = {
-                    stage["name"]: gr.Textbox(
-                        label=stage["question"],
-                        placeholder=stage["example"],
-                        value=getattr(initial_idea, stage["field"], ""),
-                        visible=(stage["name"] == default_stage),
-                        interactive=False
-                    ) for stage in STAGES
-                }
-                fill_form_btn = gr.Button("Fill out Form")
-                submit_form_btn = gr.Button("Submit Form", visible=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            model = gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0])
-            system_prompt = gr.Textbox(label="System Prompt", value=DEFAULT_SYSTEM_PROMPT, lines=5)
-            thinking_budget = gr.Slider(minimum=1, maximum=4098, value=2048, step=1, label="Max New Tokens")
-            api_key = gr.Textbox(label="Hugging Face API Key", type="password")
-
-        # Event handlers
-        msg.submit(chatbot_function, 
-                   inputs=[msg, chatbot, model, system_prompt, thinking_budget, stages],
-                   outputs=[chatbot, form_fields[default_stage], msg])
-        
-        submit.click(chatbot_function, 
-                     inputs=[msg, chatbot, model, system_prompt, thinking_budget, stages],
-                     outputs=[chatbot, form_fields[default_stage], msg])
-        
-        fill_form_btn.click(fill_form,
-                            inputs=[stages, model, thinking_budget],
-                            outputs=list(form_fields.values()))
-        
-        clear.click(clear_chat,
-                    outputs=[chatbot] + list(form_fields.values()))
-
-        # Update form field visibility based on selected stage
-        stages.change(
-            lambda s: [gr.update(visible=(stage["name"] == s)) for stage in STAGES],
-            inputs=[stages],
-            outputs=list(form_fields.values())
+        prompt = f"""
+        {formatted_system_prompt}
+
+        Based on the following conversation, extract the relevant information for the '{stage}' stage of the innovative idea:
+
+        {conversation}
+
+        Please provide a concise summary for the '{stage}' stage, focusing only on the information relevant to this stage. 
+        Your response should be structured as follows:
+
+        1. A brief analysis of the conversation related to this stage.
+        2. A concise summary of the key points relevant to this stage.
+        3. A suggested form entry for this stage, enclosed in <form_data></form_data> tags.
+           The form entry should be in the format: "{stage}: Content"
+
+        Remember to keep the form entry concise and directly related to the '{stage}' stage. Do not include information from other stages in the form entry.
+        """
+
+        # Get LLM response
+        llm_response = get_llm_response(prompt, model, thinking_budget, self.api_key)
+
+        # Extract form data from the LLM response
+        form_data = extract_form_data(llm_response)
+        
+        return form_data.get(stage, "")
+
+    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)
+            if activation_message is None:
+                error_message = f"Error: Unable to activate stage '{stage_name}'. Please check if the stage name is correct."
+                self.chat_history.append(("System", error_message))
+                yield self.chat_history, self.idea_form.dict()
+                return
+
+            self.chat_history.append(("System", activation_message))
+
+        # Check for web search request
+        if message.startswith('@'):
+            search_query = message[1:].strip()
+            optimized_query = optimize_search_query(search_query, model)
+            search_results = perform_web_search(optimized_query)
+            self.chat_history.append(("Human", message))
+            self.chat_history.append(("AI", f"Here are the search results for '{optimized_query}':\n\n{search_results}"))
+            yield self.chat_history, self.idea_form.dict()
+            return
+
+        # Generate the prompt for the current stage
+        stage_prompt = self.generate_prompt_for_stage(stage_name)
+
+        formatted_system_prompt = DEFAULT_SYSTEM_PROMPT.format(
+            current_stage=stage_name,
+            stage_prompt=stage_prompt
         )
 
-        # Update API key when changed
-        api_key.change(innovative_chatbot.set_api_key, inputs=[api_key])
+        combined_prompt = f"{formatted_system_prompt}\n\nUser input: {message}"
+
+        # Get LLM response
+        llm_response = get_llm_response(combined_prompt, model, thinking_budget, self.api_key)
+        parsed_response = self.parse_llm_response(llm_response)
+        
+        self.chat_history.append(("Human", message))
+        self.chat_history.append(("AI", parsed_response))
+
+        form_data = extract_form_data(llm_response)
+        
+        if stage_name in form_data:
+            setattr(self.idea_form, stage_name.lower().replace(" ", "_"), form_data[stage_name])
+
+        yield self.chat_history, self.idea_form.dict()
 
-        # Toggle between chatbot and direct input mode
-        def toggle_mode(new_mode):
-            if new_mode == "Direct Input":
-                return [gr.update(visible=False)] * 3 + [gr.update(interactive=True)] * len(STAGES) + [gr.update(visible=True)]
+    def fill_out_form_stream(self, current_stage: str, model: str, thinking_budget: int):
+        form_data = {}
+        for stage in IDEA_STAGES:
+            stage_name = stage.name
+            if stage_name == current_stage:
+                form_data[stage_name] = self.generate_form_data(stage_name, model, thinking_budget)
             else:
-                return [gr.update(visible=True)] * 3 + [gr.update(interactive=False)] * len(STAGES) + [gr.update(visible=False)]
+                form_data[stage_name] = getattr(self.idea_form, stage.field, "")
+            yield form_data
 
-        mode.change(
-            toggle_mode,
-            inputs=[mode],
-            outputs=[chatbot, msg, submit] + list(form_fields.values()) + [submit_form_btn]
-        )
+        # Update the idea form
+        for stage in IDEA_STAGES:
+            setattr(self.idea_form, stage.field, form_data[stage.name])
+        
+        # Save to database
+        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()
 
-        # Handle direct form submission
-        submit_form_btn.click(
-            lambda *values: values,
-            inputs=[form_fields[stage["name"]] for stage in STAGES],
-            outputs=[form_fields[stage["name"]] for stage in STAGES]
+    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])
+        
+        stage_prompt = self.generate_prompt_for_stage(stage)
+        
+        formatted_system_prompt = DEFAULT_SYSTEM_PROMPT.format(
+            current_stage=stage,
+            stage_prompt=stage_prompt
         )
 
-        # Add this new event handler for the Start Over button
-        start_over_btn.click(
-            start_over,
-            outputs=[chatbot, msg] + [form_fields[stage["name"]] for stage in STAGES] + [stages]
-        )
+        prompt = f"""
+        {formatted_system_prompt}
 
-        # Add this new event handler to display the initial greeting when the interface loads
-        demo.load(lambda: ([[None, innovative_chatbot.get_initial_greeting()]], ""),
-                  outputs=[chatbot, msg])
+        Based on the following conversation, extract the relevant information for the '{stage}' stage of the innovative idea:
 
-        # Add this new event handler to update form fields when they change
-        for stage in STAGES:
-            form_fields[stage["name"]].change(
-                lambda value, s=stage["name"]: innovative_chatbot.update_form_field(s, value),
-                inputs=[form_fields[stage["name"]]],
-                outputs=[form_fields[stage["name"]]]
-            )
-
-    return demo
-
-def main():
-    try:
-        demo = create_gradio_interface()
-        return demo
-    except ImportError as e:
-        logging.error(f"Import error: {str(e)}", exc_info=True)
-        print(f"An import error occurred: {str(e)}")
-        print("Please check your import statements and ensure there are no circular dependencies.")
-        return None
-    except Exception as e:
-        logging.error(f"Failed to initialize application: {str(e)}", exc_info=True)
-        print(f"An unexpected error occurred: {str(e)}")
-        print("Please check the log file for more details.")
-        return None
+        {conversation}
+
+        Please provide a concise summary for the '{stage}' stage, focusing only on the information relevant to this stage. 
+        Your response should be structured as follows:
+
+        1. A brief analysis of the conversation related to this stage.
+        2. A concise summary of the key points relevant to this stage.
+        3. A suggested form entry for this stage, enclosed in <form_data></form_data> tags.
+           The form entry should be in the format: "{stage}: Content"
 
-if __name__ == "__main__":
-    try:
-        demo = main()
-        if demo:
-            demo.launch()
-    except Exception as e:
-        logging.error(f"Failed to start the application: {str(e)}", exc_info=True)
-        print(f"An error occurred while starting the application: {str(e)}")
-        print("Please check the log file for more details.")
-        # You might want to add a more user-friendly error message or UI here
+        Remember to keep the form entry concise and directly related to the '{stage}' stage. Do not include information from other stages in the form entry.
+        """
+
+        llm_response = get_llm_response(prompt, model, thinking_budget, self.api_key)
+        form_data = extract_form_data(llm_response)
+        return form_data.get(stage, "")
+
+    def update_form_field(self, stage_name: str, value: str):
+        setattr(self.idea_form, stage_name.lower().replace(" ", "_"), value)
+        try:
+            # Create a new session for this operation
+            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 form field: {str(e)}")
+            # If an error occurs, rollback the new session
+            new_session.rollback()
+        finally:
+            # Always close the new session
+            new_session.close()
+        return value