Update chatbot.py

chatbot.py

@@ -1,249 +1,202 @@
-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
-import logging
-import os
-from typing import Dict, Any 
-
-# Set up logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-
-def create_gradio_interface():
-    innovative_chatbot = InnovativeIdeaChatbot()
-    default_stage = STAGES[0]["name"]
-
-    # Initialize the database
-    try:
+from typing import List, Tuple, Dict, Any, Optional
+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
+)
+
+class InnovativeIdeaChatbot:
+    def __init__(self):
+        create_tables()
         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):
-        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))
-                return history, "", ""
-
-            chat_history, form_data = innovative_chatbot.process_stage_input(current_stage, message, model, system_prompt, thinking_budget)
-            
-            history.append((message, chat_history[-1][1]))
-            
-            # 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()
-            
-            return history, form_data.get(current_stage, ""), ""  # Return an empty string to clear the input
-        except Exception as e:
-            logging.error(f"An error occurred in chatbot_function: {str(e)}", exc_info=True)
-            return history + [(None, f"An error occurred: {str(e)}")], "", ""  # Return an empty string to clear the input
-
-    def fill_form(stage, model, thinking_budget):
-        form_data = innovative_chatbot.fill_out_form(stage, model, thinking_budget)
-        return form_data
-
-    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()
-        
-        return innovative_chatbot.reset()
-
-    def update_form_field(stage_name: str, value: str) -> str:
-        db = SessionLocal()
-        idea = db.query(InnovativeIdea).first()
-        field_name = next(stage["field"] for stage in STAGES if stage["name"] == stage_name)
-        setattr(idea, field_name, value)
-        db.commit()
-        db.close()
-        return value
-
-    def update_placeholders(form_data: Dict[str, str]) -> Dict[str, Any]:
-        return {
-            stage["name"]: gr.update(value=form_data.get(stage["name"], ""), placeholder=form_data.get(stage["name"], stage["example"]))
-            for stage in STAGES
-        }
-
-    def start_over():
-        chat_history, form_data = innovative_chatbot.start_over()
-        return (
-            [],  # Clear the chatbot
-            "",  # 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
+        self.idea_form = IdeaForm()
+        self.chat_history = []
+        self.idea_id = None
+        self.db = SessionLocal()
+        self.current_stage = None
+        self.api_key = None
+        self.add_system_message(self.get_initial_greeting())
+
+    def __del__(self):
+        if hasattr(self, 'db'):
+            self.db.close()
+
+    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
+
+    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
         )
 
-    def fill_and_update_form(stage, model, thinking_budget):
-        # Get the conversation history
-        chat_history = innovative_chatbot.chat_history
-        
-        # Get the specific question for the current stage
-        stage_info = next((s for s in STAGES if s["name"] == stage), None)
-        if not stage_info:
-            return [gr.update() for _ in STAGES]  # Return no updates if stage not found
-        
-        stage_question = stage_info["question"]
-        
-        # Use the generate_form_data method from InnovativeIdeaChatbot
-        form_data = innovative_chatbot.generate_form_data(stage, model, thinking_budget)
-        
-        # Update the database and form fields
-        updated_fields = {}
-        updated_value = update_form_field(stage, form_data)
-        updated_fields[stage] = updated_value
-        
-        # Convert the dictionary to a list of individual updates
-        return [gr.update(value=updated_fields.get(s["name"], ""), 
-                          placeholder=updated_fields.get(s["name"], s["example"])) 
-                for s in STAGES]
+        # Combine the formatted system prompt and user's input
+        combined_prompt = f"{formatted_system_prompt}\n\nUser input: {message}"
 
-    with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# Innovative Idea Generator")
+        # Get LLM response
+        llm_response = get_llm_response(combined_prompt, model, thinking_budget, self.api_key)
+
+        # Add the interaction to chat history
+        self.chat_history.append(("Human", message))
+        self.chat_history.append(("AI", llm_response))
+
+        # Extract form data from the LLM response
+        form_data = extract_form_data(llm_response)
         
-        mode = gr.Radio(["Chatbot", "Direct Input"], label="Mode", value="Chatbot")
+        # 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 fill_out_form(self, current_stage: str, model: str, thinking_budget: int) -> Dict[str, str]:
+        form_data = {}
+        for stage in IDEA_STAGES:
+            stage_name = stage.name
+            if stage_name == current_stage:
+                # Generate new data for the current stage
+                form_data[stage_name] = self.generate_form_data(stage_name, model, thinking_budget)
+            else:
+                # Use existing data for other stages
+                form_data[stage_name] = getattr(self.idea_form, stage.field, "")
         
-        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])
+        # Update the idea form
+        for stage in IDEA_STAGES:
+            setattr(self.idea_form, stage.field, form_data[stage.name])
         
-        submit.click(chatbot_function, 
-                     inputs=[msg, chatbot, model, system_prompt, thinking_budget, stages],
-                     outputs=[chatbot, form_fields[default_stage], msg])
+        # 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)
+
+        return form_data
+
+    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."
+
+    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())
+        clear_database(self.db)
+        return self.chat_history, self.idea_form.dict()
+
+    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)
+
+    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])
         
-        fill_form_btn.click(fill_and_update_form,
-                            inputs=[stages, model, thinking_budget],
-                            outputs=list(form_fields.values()))
+        stage_prompt = self.generate_prompt_for_stage(stage)
         
-        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())
+        formatted_system_prompt = DEFAULT_SYSTEM_PROMPT.format(
+            current_stage=stage,
+            stage_prompt=stage_prompt
         )
 
-        # Update API key when changed
-        api_key.change(innovative_chatbot.set_api_key, inputs=[api_key])
+        prompt = f"""
+        {formatted_system_prompt}
 
-        # 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)]
-            else:
-                return [gr.update(visible=True)] * 3 + [gr.update(interactive=False)] * len(STAGES) + [gr.update(visible=False)]
+        Based on the following conversation, extract the relevant information for the '{stage}' stage of the innovative idea:
 
-        mode.change(
-            toggle_mode,
-            inputs=[mode],
-            outputs=[chatbot, msg, submit] + list(form_fields.values()) + [submit_form_btn]
-        )
+        {conversation}
 
-        # 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]
-        )
+        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:
 
-        # 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]
-        )
+        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.
+
+        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.
+        """
 
-        # 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])
+        # Get LLM response
+        llm_response = get_llm_response(prompt, model, thinking_budget, self.api_key)
 
-    return demo
+        # Extract form data from the LLM response
+        form_data = extract_form_data(llm_response)
+        
+        return form_data.get(stage, "")
 
-def main():
-    try:
-        demo = create_gradio_interface()
-        return demo
-    except Exception as e:
-        logging.error(f"Failed to initialize application: {str(e)}", exc_info=True)
-        raise
+    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()
 
+# Ensure this is at the end of the file
 if __name__ == "__main__":
-    try:
-        demo = main()
-        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
+    # Any test code or standalone functionality for chatbot.py
+    pass