add catch all

ai_tools.json

@@ -65,24 +65,22 @@
   {
     "type": "function",
     "function": {
-      "name": "clarify_arguments",
-      "description": "Ask a question to the user to get more information. This function should be called when the AI needs clarification on arguments to pass to a different function call. Should always be called for snippet_clip_to_continue_from in get_audio_snippet and whole_song in revise_instrumental_tags.",
+      "name": "edit_directly",
+      "description": "A function to be called when the user asks to directly edit either the lyrics or genre via the instrumental tags. Returns the current section or instrumental tags for the user to edit directly.",
       "parameters": {
         "type": "object",
         "properties": {
-          "function_name": {
+          "edit_type": {
             "type": "string",
-            "description": "Function to call that needs clarification on arguments"
+            "enum": ["lyrics", "genre"],
+            "description": "Specifies whether the user wants to edit the lyrics or the genre."
           },
-          "arguments_to_clarify": {
-            "type": "array",
-            "items": {
-              "type": "string"
-            },
-            "description": "List of questions to ask the end user to clarify the arguments we don't know"
+          "current_item_to_edit": {
+            "type": "string",
+            "description": "The current lyrics section or instrumental tags that the user wants to edit."
           }
         },
-        "required": ["function_name", "arguments_to_clarify"]
+        "required": ["edit_type", "current_item_to_edit"]
       }
     }
   },

app.py

@@ -17,7 +17,7 @@ css = """
 }
 """
 
-textbox = gr.Textbox(lines=1, label='Send a message', show_label=False, placeholder='Send a message', scale=4, visible=True)
+textbox = gr.Textbox(lines=2, label='Send a message', show_label=False, placeholder='Send a message', scale=4, visible=True)
 submit = gr.Button("Send", scale=2, visible=True)
 
 
@@ -28,14 +28,14 @@ with gr.Blocks(css=css) as demo:
     with gr.Tabs() as tabs:
         with gr.TabItem("Ideation", id=0): #index is 0
             gr.Markdown("""<center><font size=6>Let's write a song!</font></center>""")
-            gr.Markdown("""<center><font size=4>First, let's try to find an interesting concept. Fill out the fields below and generate a song seed.</font></center>""")
-            gr.Markdown("""<center><font size=3>If you're stuck, check out <a href="https://onestopforwriters.com/emotions" target="_blank">here</a>.</font></center>""")
+            gr.Markdown("""<center><font size=4>But first, let's generate a song seed to provide context to the AI Songwriter.</font></center>""")
+            gr.Markdown("""<center><font size=3>If you're stuck thinking of a song idea, check out <a href="https://onestopforwriters.com/emotions" target="_blank">here</a>.</font></center>""")
             with gr.Row():
-                feeling_input = gr.Textbox(label="What's an emotion(s) that you've been feeling a lot recently? And why?", placeholder='Enter your emotions', scale=2)
+                feeling_input = gr.Textbox(label="What do you want the song to capture? What's an emotion(s) that you've personally felt that should be in the song? More vulnerable, better the song.", placeholder='Enter your story', scale=2)
                 # audio_input = gr.Audio(sources=["upload"], type="numpy", label="Instrumental",
                 #                 interactive=True, elem_id="instrumental-input")
                 
-            generate_seed_button = gr.Button("STEP 1: Generate Song Seed")
+            generate_seed_button = gr.Button("Click to Generate Song Seed")
             concept_desc = gr.Markdown("""<center><font size=4>Here it is! Hit 'Approve' to confirm this concept. Edit the concept directly or hit 'Try Again' to get another suggestion.</font></center>""", visible=False)
             with gr.Row(visible=False) as concept_row:
                 instrumental_output = gr.TextArea(label="Suggested Song Concept", value="", max_lines=3, scale=2)
@@ -58,7 +58,7 @@ with gr.Blocks(css=css) as demo:
                 approve_button.click(open_accordion, inputs=[approve_button], outputs=[accordion])
                   
             with gr.Row():
-                continue_btn = gr.Button("Continue to Next Step", interactive=False)
+                continue_btn = gr.Button("Ready to Create", interactive=False)
 
             
             def clean_song_seed(song_seed):
@@ -71,8 +71,9 @@ with gr.Blocks(css=css) as demo:
             def make_row_visible(x):
                 return gr.Row(visible=True), gr.Markdown("""<center><font size=4>Here it is! Hit 'Approve' to confirm this concept. Edit the concept directly or hit 'Try Again' to get another suggestion.</font></center>""", visible=True)
             def enable_button(x):
-                return gr.Button("Continue to Next Step", interactive=True)
+                return gr.Button("Ready to Create", interactive=True)
             generate_seed_button.click(make_row_visible, inputs=[generate_seed_button], outputs=[concept_row, concept_desc])
+            feeling_input.submit(make_row_visible, inputs=[generate_seed_button], outputs=[concept_row, concept_desc])
             approve_button.click(enable_button, inputs=[approve_button], outputs=[continue_btn])
             
             try_again_button.click(generate_song_seed, inputs=[feeling_input], outputs=[instrumental_output])
@@ -84,7 +85,7 @@ with gr.Blocks(css=css) as demo:
         
         with gr.TabItem("Generation", id=1): #index is 1
             start_song_gen = gr.State(value=False)
-            gr.Markdown("""<center><font size=4>Now, chat with an AI songwriter to make your song! Hit finish when ready to hear full song.</font></center>""")        
+            gr.Markdown("""<center><font size=4>Now, chat with an AI songwriter to make your song!</font></center>""")        
 
             character = gr.State(value="A 18-year old boy who dreams of being a pop star that uplifts people going through the difficulties of life")
 
@@ -100,7 +101,7 @@ with gr.Blocks(css=css) as demo:
 
             with gr.Row():
                 with gr.Column(scale=2):
-                    chatbot_history = gr.Chatbot(type="messages", value=starting_history, label='SongChat', placeholder=None, layout='bubble', bubble_full_width=False, height=500)
+                    chatbot_history = gr.Chatbot(type="messages", value=starting_history, label='SongChat', placeholder=None, layout='bubble', bubble_full_width=False, height=500, show_copy_button=True)
                     with gr.Row():
                         typical_responses = [textbox, submit]
                         
@@ -109,13 +110,16 @@ with gr.Blocks(css=css) as demo:
 
                         button_options = gr.State([])
                         button_dict = gr.State({
-                            "revise lyrics": "Can we revise the lyrics?", 
+                            "revise lyrics": "Can we revise the lyrics?",
+                            "re-revise lyrics": "Can we revise the lyrics?", 
+                            "edit lyrics directly": "Can I edit the lyrics directly for the whole section?",
                             "generate audio snippet": "Can you generate an audio snippet?", 
                             "continue revising" : "Can we continue revising this section?", 
                             "generate audio snippet with new lyrics": "Can you generate an audio snippet with these new lyrics?", 
                             "return to original instrumental": "Can you use the original clip for this section instead?", 
                             "revise genre": "Can we revise the instrumental tags?",
                             "re-revise genre": "Can we revise the instrumental tags?", 
+                            "revise genre directly": "Can I edit the genre directly for the whole song?",
                             "continue to next section": "Looks good! Let's move on to the next section.",
                             "merge snippets": "Can you merge this snippet into its full song?"
                         })
@@ -187,6 +191,7 @@ with gr.Blocks(css=css) as demo:
                 return new_step_number, *modals
             
             def update_textbox(textbox, step_number):
+                print("on step number", step_number)
                 if step_number == 0:
                     return textbox + "\nAsk me another question to inform the verse"
                 elif step_number == 1:
@@ -214,7 +219,7 @@ with gr.Blocks(css=css) as demo:
                 for btn in typical_responses[2:]:
                     btn.click(set_response_buttons, inputs=[button_dict, btn], outputs=[textbox]).then(model_chat, 
                                     inputs=[genre_input, textbox, chatbot_history, messages, generated_audios], 
-                                    outputs=[textbox, chatbot_history, messages, current_section, current_lyrics, curr_tags, clip_to_continue, curr_audio, generated_audios, button_options]).then(reset_textbox, inputs=[textbox], outputs=[textbox]).then(
+                                    outputs=[textbox, chatbot_history, messages, current_section, current_lyrics, curr_tags, clip_to_continue, curr_audio, generated_audios, button_options]).then(
                                     update_response_options, [button_options, button_dict], typical_responses
                             ).then(
                             make_modal_visible, [tutorial_step], [tutorial_step, modal, modal_1, modal_2]
@@ -225,14 +230,14 @@ with gr.Blocks(css=css) as demo:
 
             submit.click(update_textbox, [textbox, tutorial_step], [textbox]).then(model_chat,
                 inputs=[genre_input, textbox, chatbot_history, messages, generated_audios],
-                outputs=[textbox, chatbot_history, messages, current_section, current_lyrics, curr_tags, clip_to_continue, curr_audio, generated_audios, button_options]).then(reset_textbox, inputs=[textbox], outputs=[textbox]).then(
+                outputs=[textbox, chatbot_history, messages, current_section, current_lyrics, curr_tags, clip_to_continue, curr_audio, generated_audios, button_options]).then(
                         update_response_options, [button_options, button_dict], typical_responses
                 ).then(
                             make_modal_visible, [tutorial_step], [tutorial_step, modal, modal_1, modal_2]
                         )
             textbox.submit(update_textbox, [textbox, tutorial_step], [textbox]).then(model_chat, 
                 inputs=[genre_input, textbox, chatbot_history, messages, generated_audios], 
-                outputs=[textbox, chatbot_history, messages, current_section, current_lyrics, curr_tags, clip_to_continue, curr_audio, generated_audios, button_options]).then(reset_textbox, inputs=[textbox], outputs=[textbox]).then(
+                outputs=[textbox, chatbot_history, messages, current_section, current_lyrics, curr_tags, clip_to_continue, curr_audio, generated_audios, button_options]).then(
                         update_response_options, [button_options, button_dict], typical_responses
                 ).then(
                             make_modal_visible, [tutorial_step], [tutorial_step, modal, modal_1, modal_2]
@@ -241,15 +246,15 @@ with gr.Blocks(css=css) as demo:
             
             regen.click(set_regenerate_query, inputs=[textbox, current_section, current_lyrics, curr_tags, clip_to_continue], outputs=[textbox]).then(model_chat,
                 inputs=[genre_input, textbox, chatbot_history, messages, generated_audios],
-                outputs=[textbox, chatbot_history, messages, current_section, current_lyrics, curr_tags, clip_to_continue, curr_audio, generated_audios, button_options]).then(reset_textbox, inputs=[textbox], outputs=[textbox]).then(
+                outputs=[textbox, chatbot_history, messages, current_section, current_lyrics, curr_tags, clip_to_continue, curr_audio, generated_audios, button_options]).then(
                         update_response_options, [button_options, button_dict], typical_responses
                 ).then(
                             make_modal_visible, [tutorial_step], [tutorial_step, modal, modal_1, modal_2]
                         )
 
-            with gr.Row(visible=False):
+            with gr.Row(visible=True):
                 # get_snippet_button = gr.Button("Get Audio Snippet", scale=2)
-                done = gr.Button("Finish Full Song 🎶", scale=4)
+                done = gr.Button("Complete User Study 🎶", scale=4)
                 #autoGPT_checkbox = gr.Checkbox(label="AutoGPT", value=True, info="Auto-generate responses from journal entry", interactive=True, scale=2)
                 #journal_llm_creativity = gr.Slider(label="Journal LLM Temperature", minimum=0, maximum=1, step=0.01, value=1, interactive=True, scale=2)
                 reset_button = gr.Button("Reset", scale=2)

chat.py

@@ -31,6 +31,12 @@ def determine_title(section_name, generated_audios):
         section_name = f"{section_name} {count + 1}"
     return section_name
 
+async def call_with_timeout(coro, timeout):
+    try:
+        return await asyncio.wait_for(coro, timeout)
+    except asyncio.TimeoutError:
+        return "Timeout"
+
 
 def model_chat(genre_input, query: Optional[str], history: Optional[History], messages: Optional[Messages], generated_audios: List[Tuple[str, str, str]], auto=False) -> Tuple[str, History, Messages, str, str, str, str, str, List]:
     if query is None:
@@ -58,306 +64,425 @@ def model_chat(genre_input, query: Optional[str], history: Optional[History], me
     # Step 2: determine if the response from the model includes a tool call.   
     tool_calls = response_message.choices[0].message.tool_calls
     if tool_calls:
-        messages.append({
-            "role": response_message.choices[0].message.role,
-            "content": response_message.choices[0].message.content,
-            "tool_calls": tool_calls,
-            "function_call": response_message.choices[0].message.function_call
-        })
-
-        if len(tool_calls) > 1:
-            for tool_call in tool_calls:
-                tool_message = {
-                    'role': 'tool',
-                    'tool_call_id': tool_call.id,
-                    'name': tool_call.function.name,
-                    'content': "You called two different functions when you can only call one at a time. Did you mean to call revise_section_lyrics_and_instrumental but instead had two different calls for lyrics and instrumental? Communicate this failure to the user and clarify what they are asking for, then only call one tool next time."
-                }
-                messages.append(tool_message)
-
-            # Generate a response using GPT-4o and add it as a message
-            model_response_with_function_call = oai_client.chat.completions.create(
-                model="gpt-4o",
-                messages=messages,
-            )
-            current_response = model_response_with_function_call.choices[0].message.content
-
-            role = "assistant"
-            messages.append({'role': role, 'content': current_response})
-
-            yield '', messages_to_history(messages), messages, '', '', '', '', None, generated_audios, []
-            return
-
-
-        # If true the model will return the name of the tool / function to call and the argument(s)  
-        for tool_call in tool_calls:
-            print(tool_call)
-            tool_call_id = tool_call.id
-            tool_function_name = tool_call.function.name
-            tool_query_args = eval(tool_call.function.arguments)
-
-            print(tool_function_name, tool_query_args)
-
-            with open('ai_tools.json') as f:
-                ai_tools = json.load(f)
-            
-            for tool in ai_tools:
-                if tool['function']['name'] == tool_function_name:
-                    valid_keys = tool['function']['parameters']['properties'].keys()
-                    required_keys = tool['function']['parameters']['required']
-                    break
-            
-            print('query args before', tool_query_args)
-            tool_query_args = {k: v for k, v in tool_query_args.items() if k in valid_keys}
-            print('query args after', tool_query_args)
-            missing_keys = []
-            for key in required_keys:
-                if key not in tool_query_args:
-                    missing_keys.append(key)
-            if len(missing_keys)>0:
-                missing_keys_str = ", ".join(missing_keys)
-                tool_message = {
-                    'role': 'tool',
-                    'tool_call_id': tool_call_id,
-                    'name': tool_function_name,
-                    'content': f"Sorry, the keys {missing_keys_str} from the function you called are missing, communicate this to the user and either get what these args should be or figure out which function to call."
-                }
-
-                new_messages = messages + [tool_message]
+        try:
+
+            messages.append({
+                "role": response_message.choices[0].message.role,
+                "content": response_message.choices[0].message.content,
+                "tool_calls": tool_calls,
+                "function_call": response_message.choices[0].message.function_call
+            })
+
+            if len(tool_calls) > 1:
+                for tool_call in tool_calls:
+                    tool_message = {
+                        'role': 'tool',
+                        'tool_call_id': tool_call.id,
+                        'name': tool_call.function.name,
+                        'content': "You called two different functions when you can only call one at a time. Did you mean to call revise_section_lyrics_and_instrumental but instead had two different calls for lyrics and instrumental? Communicate this failure to the user and clarify what they are asking for, then only call one tool next time."
+                    }
+                    messages.append(tool_message)
 
+                # Generate a response using GPT-4o and add it as a message
                 model_response_with_function_call = oai_client.chat.completions.create(
                     model="gpt-4o",
-                    messages=new_messages,
-                )  # get a new response from the model where it can see the function response
+                    messages=messages,
+                )
                 current_response = model_response_with_function_call.choices[0].message.content
 
                 role = "assistant"
-                new_messages = new_messages + [{'role': role, 'content': current_response}]
-                new_history = messages_to_history(new_messages)
+                messages.append({'role': role, 'content': current_response})
 
-                generated_audios = update_song_links(generated_audios)
-                yield '', new_history, new_messages, '', '', '', '', None, generated_audios, []
-                        
-                        
-            # Step 3: Call the function and retrieve results. Append the results to the messages list. 
-            if tool_function_name == 'ask_question':
-                question = songwriterAssistant.ask_question(messages)
+                yield '', messages_to_history(messages), messages, '', '', '', '', None, generated_audios, []
+                return
 
-                question = question.replace("ask question:", "").replace("ask question ", "").replace("ask question\n", "").replace("ask question", "")
 
-                ## yield question in tool and assistant message
-                tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': question}
+            # If true the model will return the name of the tool / function to call and the argument(s)  
+            for tool_call in tool_calls:
+                print(tool_call)
+                tool_call_id = tool_call.id
+                tool_function_name = tool_call.function.name
+                tool_query_args = eval(tool_call.function.arguments)
 
-                new_messages = messages + [tool_message]
+                print(tool_function_name, tool_query_args)
 
-                question_message = {'role': 'assistant', 'content': question}
-                new_messages = new_messages + [question_message]
-                new_history = messages_to_history(new_messages)
+                with open('ai_tools.json') as f:
+                    ai_tools = json.load(f)
+                
+                for tool in ai_tools:
+                    if tool['function']['name'] == tool_function_name:
+                        valid_keys = tool['function']['parameters']['properties'].keys()
+                        required_keys = tool['function']['parameters']['required']
+                        break
+                
+                print('query args before', tool_query_args)
+                tool_query_args = {k: v for k, v in tool_query_args.items() if k in valid_keys}
+                print('query args after', tool_query_args)
+                missing_keys = []
+                for key in required_keys:
+                    if key not in tool_query_args:
+                        missing_keys.append(key)
+                if len(missing_keys)>0:
+                    missing_keys_str = ", ".join(missing_keys)
+                    tool_message = {
+                        'role': 'tool',
+                        'tool_call_id': tool_call_id,
+                        'name': tool_function_name,
+                        'content': f"Sorry, the keys {missing_keys_str} from the function you called are missing, communicate this to the user and either get what these args should be or figure out which function to call."
+                    }
 
-                generated_audios = update_song_links(generated_audios)
-                yield '', new_history, new_messages, '', '', '', '', None, generated_audios, []
-            
-            elif tool_function_name == 'clarify_arguments':
-                tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': 'arguments to clarify: \n' + '\n'.join(tool_query_args['arguments_to_clarify'])}
+                    new_messages = messages + [tool_message]
 
-                new_messages = messages + [tool_message]
+                    model_response_with_function_call = oai_client.chat.completions.create(
+                        model="gpt-4o",
+                        messages=new_messages,
+                    )  # get a new response from the model where it can see the function response
+                    current_response = model_response_with_function_call.choices[0].message.content
 
-                model_response_with_function_call = oai_client.chat.completions.create(
-                    model="gpt-4o",
-                    messages=new_messages,
-                )  # get a new response from the model where it can see the function response
-                current_response = model_response_with_function_call.choices[0].message.content
+                    role = "assistant"
+                    new_messages = new_messages + [{'role': role, 'content': current_response}]
+                    new_history = messages_to_history(new_messages)
 
-                role = "assistant"
-                new_messages = new_messages + [{'role': role, 'content': current_response}] # + "\n\nWould you like to get an audio snippet? Or continue writing?"}]
-                # new_messages = [msg for msg in new_messages if msg['content'] is not None and msg['role'] in ['user', 'assistant']]
-                new_history = messages_to_history(new_messages)
+                    generated_audios = update_song_links(generated_audios)
+                    yield '', new_history, new_messages, '', '', '', '', None, generated_audios, []
+                            
+                            
+                # Step 3: Call the function and retrieve results. Append the results to the messages list. 
+                if tool_function_name == 'ask_question':
+                    question = songwriterAssistant.ask_question(messages)
 
-                generated_audios = update_song_links(generated_audios)
-                yield '', new_history, new_messages, '', '', '', '', None, generated_audios, []
+                    question = question.replace("ask question:", "").replace("ask question ", "").replace("ask question\n", "").replace("ask question", "")
 
-            
-            elif tool_function_name == 'write_section':
-                snippet_instrumental_tags = tool_query_args.pop('snippet_instrumental_tags', None)
-                snippet_clip_to_continue_from = tool_query_args.pop('snippet_clip_to_continue_from', None)
-                suggested_lyrics = songwriterAssistant.write_section(**tool_query_args)
-                suggested_lyrics = suggested_lyrics.strip('`*-\n')
+                    ## yield question in tool and assistant message
+                    tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': question}
 
-                ## yield suggested lyrics in tool and assistant message
-                tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': suggested_lyrics}
-                # audio_message = {'role': 'assistant', 'content': "Here's what I've come up with:\n" + suggested_lyrics + "\n\nGenerating audio snippet..."}
-                new_messages = messages + [tool_message] #, audio_message
+                    new_messages = messages + [tool_message]
 
-                model_response_with_function_call = oai_client.chat.completions.create(
-                    model="gpt-4o",
-                    messages=new_messages,
-                )  # get a new response from the model where it can see the function response
-                current_response = model_response_with_function_call.choices[0].message.content
+                    question_message = {'role': 'assistant', 'content': question}
+                    new_messages = new_messages + [question_message]
+                    new_history = messages_to_history(new_messages)
 
-                role = "assistant"
-                new_messages = new_messages + [{'role': role, 'content': current_response}] # + "\n\nWould you like to get an audio snippet? Or continue writing?"}]
-                # new_messages = [msg for msg in new_messages if msg['content'] is not None and msg['role'] in ['user', 'assistant']]
-                history = messages_to_history(new_messages)
+                    generated_audios = update_song_links(generated_audios)
+                    yield '', new_history, new_messages, '', '', '', '', None, generated_audios, []
+                
+                elif tool_function_name == 'edit_directly':
+                    edit_type = tool_query_args.get('edit_type', 'item')
+                    current_item_to_edit = tool_query_args.get('current_item_to_edit', '')
+
+                    content_message = f"Sure, you can edit the {edit_type}. Here you go, edit and submit your new version in the chat textbox!"
+                    tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': content_message}
+
+                    new_messages = messages + [tool_message]
 
-                # current_section, current_lyrics, curr_tags, clip_to_continue, curr_audio
-                buttons = ["revise lyrics", "generate audio snippet", "continue to next section"]
+                    assistant_message = {'role': 'assistant', 'content': content_message}
+                    new_messages = new_messages + [assistant_message]
+                    new_history = messages_to_history(new_messages)
 
-                generated_audios = update_song_links(generated_audios)
-                clips_to_continue = gr.Dropdown(label='Clip to continue', value = snippet_clip_to_continue_from, choices=[x[3] for x in generated_audios]+[""], interactive=True)
+                    generated_audios = update_song_links(generated_audios)
+                    yield f"Here are my edits:\n{current_item_to_edit}", new_history, new_messages, '', '', '', '', None, generated_audios, []
+                
+                # elif tool_function_name == 'clarify_arguments':
+                #     tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': 'arguments to clarify: \n' + '\n'.join(tool_query_args['arguments_to_clarify'])}
 
+                #     new_messages = messages + [tool_message]
 
+                #     model_response_with_function_call = oai_client.chat.completions.create(
+                #         model="gpt-4o",
+                #         messages=new_messages,
+                #     )  # get a new response from the model where it can see the function response
+                #     current_response = model_response_with_function_call.choices[0].message.content
 
-                yield '', history, new_messages, tool_query_args['section_name'], suggested_lyrics.split(':')[-1], snippet_instrumental_tags, clips_to_continue, None, generated_audios, buttons
+                #     role = "assistant"
+                #     new_messages = new_messages + [{'role': role, 'content': current_response}] # + "\n\nWould you like to get an audio snippet? Or continue writing?"}]
+                #     # new_messages = [msg for msg in new_messages if msg['content'] is not None and msg['role'] in ['user', 'assistant']]
+                #     new_history = messages_to_history(new_messages)
 
-                ### DO SOMETHING TO UPDATE CURRENT GENERATION for write_section
+                #     generated_audios = update_song_links(generated_audios)
+                #     yield '', new_history, new_messages, '', '', '', '', None, generated_audios, []
 
+                
+                elif tool_function_name == 'write_section':
+                    snippet_instrumental_tags = tool_query_args.pop('snippet_instrumental_tags', None)
+                    snippet_clip_to_continue_from = tool_query_args.pop('snippet_clip_to_continue_from', None)
+                    suggested_lyrics = songwriterAssistant.write_section(**tool_query_args)
+                    suggested_lyrics = suggested_lyrics.strip('`*-\n')
 
-            elif tool_function_name == 'revise_section_lyrics':
-                revised_lyrics = songwriterAssistant.revise_section_lyrics(**tool_query_args)
+                    ## yield suggested lyrics in tool and assistant message
+                    tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': suggested_lyrics}
+                    # audio_message = {'role': 'assistant', 'content': "Here's what I've come up with:\n" + suggested_lyrics + "\n\nGenerating audio snippet..."}
+                    new_messages = messages + [tool_message] #, audio_message
 
-                # if isinstance(revised_lyrics, list):
-                #     revised_lyrics = '\n'.join(revised_lyrics)
-                if isinstance(revised_lyrics, str) and revised_lyrics.startswith("[") and revised_lyrics.endswith("]"):
-                    try:
-                        revised_lyrics = eval(revised_lyrics)
-                        if isinstance(revised_lyrics, list):
-                            revised_lyrics = '\n'.join(revised_lyrics)
-                    except:
-                        pass
+                    model_response_with_function_call = oai_client.chat.completions.create(
+                        model="gpt-4o",
+                        messages=new_messages,
+                    )  # get a new response from the model where it can see the function response
+                    current_response = model_response_with_function_call.choices[0].message.content
 
-                # ## yield revised lyrics in tool and assistant message
-                tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': revised_lyrics}
-                # audio_message = {'role': 'assistant', 'content': "Here's my revised lyrics:\n" + revised_lyrics + "\n\nGenerating audio snippet..."}
-                new_messages = messages + [tool_message] #, audio_message]
+                    role = "assistant"
+                    new_messages = new_messages + [{'role': role, 'content': current_response}] # + "\n\nWould you like to get an audio snippet? Or continue writing?"}]
+                    # new_messages = [msg for msg in new_messages if msg['content'] is not None and msg['role'] in ['user', 'assistant']]
+                    history = messages_to_history(new_messages)
 
-                model_response_with_function_call = oai_client.chat.completions.create(
-                    model="gpt-4o",
-                    messages=new_messages,
-                )  # get a new response from the model where it can see the function response
-                current_response = model_response_with_function_call.choices[0].message.content
+                    # current_section, current_lyrics, curr_tags, clip_to_continue, curr_audio
+                    buttons = ["revise lyrics",  "edit lyrics directly", "generate audio snippet", "continue to next section"]
 
-                buttons = ["revise lyrics again", "generate audio snippet with new lyrics", "continue to next section"]
+                    generated_audios = update_song_links(generated_audios)
+                    clips_to_continue = gr.Dropdown(label='Clip to continue', value = snippet_clip_to_continue_from, choices=[x[3] for x in generated_audios]+[""], interactive=True)
 
-                role = "assistant"
-                new_messages = new_messages + [{'role': role, 'content': current_response + "\n\nWould you like to get an audio snippet? Or continue writing?"}]
-                # new_messages = [msg for msg in new_messages if msg['content'] is not None and msg['role'] in ['user', 'assistant']]
-                history = messages_to_history(new_messages)
-                generated_audios = update_song_links(generated_audios)
-                clips_to_continue = gr.Dropdown(label='Clip to continue', value = "", choices=[x[3] for x in generated_audios]+[""], interactive=True)
 
-                yield '', history, new_messages, tool_query_args['section_name'], revised_lyrics, '', clips_to_continue, None, generated_audios, buttons
 
+                    yield '', history, new_messages, tool_query_args['section_name'], suggested_lyrics.split(':')[-1], snippet_instrumental_tags, clips_to_continue, None, generated_audios, buttons
 
-            elif tool_function_name == 'revise_instrumental_tags':
-                #detangle tool_query_args dict
-                #snippet_lyrics = tool_query_args['snippet_lyrics'] + "\n[End]"
-                snippet_instrumental_tags = tool_query_args.get('current_instrumental_tags', None)
-                user_instrumental_feedback = tool_query_args.get('user_instrumental_feedback', None)
+                    ### DO SOMETHING TO UPDATE CURRENT GENERATION for write_section
 
-                if snippet_instrumental_tags is None or user_instrumental_feedback is None:
-                    tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': 'Arguments are missing. Please clarify your feedback on the instrumental. Note that you cannot revise the genre if you haven\'t generated a snippet.'}
-                    audio_message = {'role': 'assistant', 'content': 'It seems like some information is missing. Could you please provide your feedback on the instrumental? Note that you cannot revise the genre if you haven\'t generated a snippet.'}
-                    new_messages = messages + [tool_message, audio_message]
-                    new_history = messages_to_history(new_messages)
-                    yield '', new_history, new_messages, '', '', '', None, None, generated_audios, []
-                    return
-                # if 'snippet_clip_to_continue_from' not in tool_query_args:
-                #     tool_query_args['snippet_clip_to_continue_from'] = None
-                # snippet_clip_to_continue_from = tool_query_args['snippet_clip_to_continue_from']
-
-                new_instrumental_tags = songwriterAssistant.revise_instrumental_tags(snippet_instrumental_tags, user_instrumental_feedback)
-
-                if isinstance(tool_query_args['sections_written'], str):
-                    current_lyrics = tool_query_args['sections_written']
-                elif isinstance(tool_query_args['sections_written'], list):
-                    current_lyrics = "\n".join(tool_query_args['sections_written'])
-                else:
-                    current_lyrics = ""
-
-                import re
-                sections_list = re.findall(r'\[.*?\]', current_lyrics)
-
-                #current_lyrics = "\n".join(tool_query_args['sections_written'])
-                song_link = make_song(current_lyrics, new_instrumental_tags)
-                ## filter out suno link from tool query arg
-                while "https://audiopipe.suno.ai/?item_id=" not in song_link:
-                    print("BUGGED OUT, trying again...")
-                    time.sleep(5)
-                    song_link = make_song(current_lyrics, new_instrumental_tags)
 
-                clip_id = song_link.split("https://audiopipe.suno.ai/?item_id=")[1]
+                elif tool_function_name == 'revise_section_lyrics':
+                    revised_lyrics = songwriterAssistant.revise_section_lyrics(**tool_query_args)
 
-                tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': f'new instrumental tags: {new_instrumental_tags}, clip id: {clip_id}'}
-                audio_message = {'role': 'assistant', 'content': f'Sure! I\'ve revised the instrumental tags: {new_instrumental_tags}\nCurrent lyrics: {current_lyrics}\n\n <audio controls autoplay><source src="{song_link}" type="audio/mp3"></audio>'}
-                audio_message['content'] += f'\n\nWhat do you think?'
-                new_messages = messages + [tool_message, audio_message]
-                new_history = messages_to_history(new_messages)
+                    # if isinstance(revised_lyrics, list):
+                    #     revised_lyrics = '\n'.join(revised_lyrics)
+                    if isinstance(revised_lyrics, str) and revised_lyrics.startswith("[") and revised_lyrics.endswith("]"):
+                        try:
+                            revised_lyrics = eval(revised_lyrics)
+                            if isinstance(revised_lyrics, list):
+                                revised_lyrics = '\n'.join(revised_lyrics)
+                        except:
+                            pass
 
-                # current_section, current_lyrics, curr_tags, clip_to_continue, curr_audio
-                if len(sections_list) > 0:
-                    section_name = f"Up to {sections_list[-1]}"
-                else:
-                    section_name = "Up to latest section"
-                section_name = determine_title(section_name, generated_audios)
+                    # ## yield revised lyrics in tool and assistant message
+                    tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': revised_lyrics}
+                    # audio_message = {'role': 'assistant', 'content': "Here's my revised lyrics:\n" + revised_lyrics + "\n\nGenerating audio snippet..."}
+                    new_messages = messages + [tool_message] #, audio_message]
 
-                generated_audios.append((song_link, current_lyrics, new_instrumental_tags, section_name, "streaming"))
-                generated_audios = update_song_links(generated_audios)
-                clips_to_continue = gr.Dropdown(label='Clip to continue', value = "", choices=[x[3] for x in generated_audios]+[""], interactive=True)
+                    model_response_with_function_call = oai_client.chat.completions.create(
+                        model="gpt-4o",
+                        messages=new_messages,
+                    )  # get a new response from the model where it can see the function response
+                    current_response = model_response_with_function_call.choices[0].message.content
 
-                buttons = ["return to original instrumental", "re-revise genre", "revise lyrics", "merge snippets", "continue to next section"]
+                    buttons = ["re-revise lyrics", "edit lyrics directly", "generate audio snippet with new lyrics", "continue to next section"]
 
-                yield '', new_history, new_messages, ', '.join(sections_list), current_lyrics, new_instrumental_tags, clips_to_continue, f'<audio controls><source src="{song_link}" type="audio/mp3"></audio>', generated_audios, buttons
-            elif tool_function_name == 'revise_section_lyrics_and_instrumental':
-                snippet_instrumental_tags = tool_query_args.pop('current_instrumental_tags', None)
-                user_instrumental_feedback = tool_query_args.pop('user_instrumental_feedback', None)
-                snippet_clip_to_continue_from = tool_query_args.pop('snippet_clip_to_continue_from', None)
-                
-                # Revise section lyrics
-                revised_lyrics = songwriterAssistant.revise_section_lyrics(**tool_query_args)
-                
-                # Revise instrumental tags
-                
-                new_instrumental_tags = songwriterAssistant.revise_instrumental_tags(snippet_instrumental_tags, user_instrumental_feedback)
+                    role = "assistant"
+                    new_messages = new_messages + [{'role': role, 'content': current_response + "\n\nWould you like to get an audio snippet? Or continue writing?"}]
+                    # new_messages = [msg for msg in new_messages if msg['content'] is not None and msg['role'] in ['user', 'assistant']]
+                    history = messages_to_history(new_messages)
+                    generated_audios = update_song_links(generated_audios)
+                    clips_to_continue = gr.Dropdown(label='Clip to continue', value = "", choices=[x[3] for x in generated_audios]+[""], interactive=True)
 
-                song_link = make_song(revised_lyrics, new_instrumental_tags, snippet_clip_to_continue_from)
-                while "https://audiopipe.suno.ai/?item_id=" not in song_link:
-                    print("BUGGED OUT, trying again...")
-                    time.sleep(5)
-                    song_link = make_song(revised_lyrics, new_instrumental_tags, snippet_clip_to_continue_from)
+                    yield '', history, new_messages, tool_query_args['section_name'], revised_lyrics, '', clips_to_continue, None, generated_audios, buttons
 
-                clip_id = song_link.split("https://audiopipe.suno.ai/?item_id=")[1]
 
-                tool_message_instrumental = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': f'revised lyrics: {revised_lyrics}\nrevised instrumental tags: {new_instrumental_tags}, clip id: {clip_id}'}
-                audio_message = {'role': 'assistant', 'content': f'Sure! I\'ve revised the lyrics and instrumental tags: {revised_lyrics}\nRevised lyrics: {revised_lyrics}\n\n <audio controls autoplay><source src="{song_link}" type="audio/mp3"></audio>'}
-                audio_message['content'] += f'\n\nWhat do you think?'
-                
-                new_messages = messages + [tool_message_instrumental, audio_message]
-                new_history = messages_to_history(new_messages)
+                elif tool_function_name == 'revise_instrumental_tags':
+                    #detangle tool_query_args dict
+                    #snippet_lyrics = tool_query_args['snippet_lyrics'] + "\n[End]"
+                    snippet_instrumental_tags = tool_query_args.get('current_instrumental_tags', None)
+                    user_instrumental_feedback = tool_query_args.get('user_instrumental_feedback', None)
 
-                generated_audios.append((song_link, revised_lyrics, new_instrumental_tags, tool_query_args["section_name"], "streaming"))
-                generated_audios = update_song_links(generated_audios)
-                clips_to_continue = gr.Dropdown(label='Clip to continue', value = snippet_clip_to_continue_from, choices=[x[3] for x in generated_audios]+[""], interactive=True)
+                    if snippet_instrumental_tags is None or user_instrumental_feedback is None:
+                        tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': 'Arguments are missing. Please clarify your feedback on the instrumental. Note that you cannot revise the genre if you haven\'t generated a snippet.'}
+                        audio_message = {'role': 'assistant', 'content': 'It seems like some information is missing. Could you please provide your feedback on the instrumental? Note that you cannot revise the genre if you haven\'t generated a snippet.'}
+                        new_messages = messages + [tool_message, audio_message]
+                        new_history = messages_to_history(new_messages)
+                        yield '', new_history, new_messages, '', '', '', None, None, generated_audios, []
+                        return
+                    # if 'snippet_clip_to_continue_from' not in tool_query_args:
+                    #     tool_query_args['snippet_clip_to_continue_from'] = None
+                    # snippet_clip_to_continue_from = tool_query_args['snippet_clip_to_continue_from']
 
-                buttons = ["return to original instrumental", "re-revise genre", "revise lyrics", "merge snippets", "continue to next section"]
+                    new_instrumental_tags = songwriterAssistant.revise_instrumental_tags(snippet_instrumental_tags, user_instrumental_feedback)
 
-                yield '', new_history, new_messages, tool_query_args["section_name"], revised_lyrics, new_instrumental_tags, clips_to_continue, f'<audio controls><source src="{song_link}" type="audio/mp3"></audio>', generated_audios, buttons
-            
-            elif tool_function_name == 'merge_all_snippets':
-                updated_clip_url, updated_lyrics, updated_tags, clips_list = concat_snippets(tool_query_args['last_snippet_id'])
+                    if isinstance(tool_query_args['sections_written'], str):
+                        current_lyrics = tool_query_args['sections_written']
+                    elif isinstance(tool_query_args['sections_written'], list):
+                        current_lyrics = "\n".join(tool_query_args['sections_written'])
+                    else:
+                        current_lyrics = ""
 
-                if "still streaming" in updated_clip_url:
-                    tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': f'still streaming, try again later'}
-                    audio_message = {'role': 'assistant', 'content': f'Unfortunately the generated clip audio is still being streamed, so you can merge later when it is fully generated.'}
+                    import re
+                    sections_list = re.findall(r'\[.*?\]', current_lyrics)
 
+                    #current_lyrics = "\n".join(tool_query_args['sections_written'])
+                    song_link = make_song(current_lyrics, new_instrumental_tags)
+                    ## filter out suno link from tool query arg
+                    while "https://audiopipe.suno.ai/?item_id=" not in song_link:
+                        print("BUGGED OUT, trying again...")
+                        if song_link == "Timeout":
+                            tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': 'The request to generate the song timed out. Please try again.'}
+                            audio_message = {'role': 'assistant', 'content': 'The request to generate the song timed out. Try again later or rephrase your request.'}
+                            new_messages = messages + [tool_message, audio_message]
+                            new_history = messages_to_history(new_messages)
+                            yield '', new_history, new_messages, '', '', '', None, None, generated_audios, []
+                            return
+                        time.sleep(5)
+                        song_link = make_song(current_lyrics, new_instrumental_tags)
+
+                    clip_id = song_link.split("https://audiopipe.suno.ai/?item_id=")[1]
+
+                    tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': f'new instrumental tags: {new_instrumental_tags}, clip id: {clip_id}'}
+                    audio_message = {'role': 'assistant', 'content': f'Sure! I\'ve revised the instrumental tags: {new_instrumental_tags}\nCurrent lyrics: {current_lyrics}\n\n <audio controls autoplay><source src="{song_link}" type="audio/mp3"></audio>'}
+                    audio_message['content'] += f'\n\nWhat do you think?'
                     new_messages = messages + [tool_message, audio_message]
                     new_history = messages_to_history(new_messages)
 
+                    # current_section, current_lyrics, curr_tags, clip_to_continue, curr_audio
+                    if len(sections_list) > 0:
+                        section_name = f"Up to {sections_list[-1]}"
+                    else:
+                        section_name = "Up to latest section"
+                    section_name = determine_title(section_name, generated_audios)
+
+                    generated_audios.append((song_link, current_lyrics, new_instrumental_tags, section_name, "streaming"))
+                    generated_audios = update_song_links(generated_audios)
                     clips_to_continue = gr.Dropdown(label='Clip to continue', value = "", choices=[x[3] for x in generated_audios]+[""], interactive=True)
 
+                    buttons = ["return to original instrumental", "re-revise genre", "edit genre directly", "revise lyrics", "merge snippets", "continue to next section"]
+
+                    yield '', new_history, new_messages, ', '.join(sections_list), current_lyrics, new_instrumental_tags, clips_to_continue, f'<audio controls><source src="{song_link}" type="audio/mp3"></audio>', generated_audios, buttons
+                elif tool_function_name == 'revise_section_lyrics_and_instrumental':
+                    snippet_instrumental_tags = tool_query_args.pop('current_instrumental_tags', None)
+                    user_instrumental_feedback = tool_query_args.pop('user_instrumental_feedback', None)
+                    snippet_clip_to_continue_from = tool_query_args.pop('snippet_clip_to_continue_from', None)
+                    
+                    # Revise section lyrics
+                    revised_lyrics = songwriterAssistant.revise_section_lyrics(**tool_query_args)
+                    
+                    # Revise instrumental tags
+                    
+                    new_instrumental_tags = songwriterAssistant.revise_instrumental_tags(snippet_instrumental_tags, user_instrumental_feedback)
+
+                    song_link = make_song(current_lyrics, new_instrumental_tags)
+                    ## filter out suno link from tool query arg
+                    while "https://audiopipe.suno.ai/?item_id=" not in song_link:
+                        print("BUGGED OUT, trying again...")
+                        if song_link == "Timeout":
+                            tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': 'The request to generate the song timed out. Please try again.'}
+                            audio_message = {'role': 'assistant', 'content': 'The request to generate the song timed out. Try again later or rephrase your request.'}
+                            new_messages = messages + [tool_message, audio_message]
+                            new_history = messages_to_history(new_messages)
+                            yield '', new_history, new_messages, '', '', '', None, None, generated_audios, []
+                            return
+                        time.sleep(5)
+                        song_link = make_song(current_lyrics, new_instrumental_tags)
+                    clip_id = song_link.split("https://audiopipe.suno.ai/?item_id=")[1]
+
+                    tool_message_instrumental = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': f'revised lyrics: {revised_lyrics}\nrevised instrumental tags: {new_instrumental_tags}, clip id: {clip_id}'}
+                    audio_message = {'role': 'assistant', 'content': f'Sure! I\'ve revised the lyrics and instrumental tags: {revised_lyrics}\nRevised lyrics: {revised_lyrics}\n\n <audio controls autoplay><source src="{song_link}" type="audio/mp3"></audio>'}
+                    audio_message['content'] += f'\n\nWhat do you think?'
+                    
+                    new_messages = messages + [tool_message_instrumental, audio_message]
+                    new_history = messages_to_history(new_messages)
+
+                    generated_audios.append((song_link, revised_lyrics, new_instrumental_tags, tool_query_args["section_name"], "streaming"))
+                    generated_audios = update_song_links(generated_audios)
+                    clips_to_continue = gr.Dropdown(label='Clip to continue', value = snippet_clip_to_continue_from, choices=[x[3] for x in generated_audios]+[""], interactive=True)
+
+                    buttons = ["return to original instrumental", "re-revise genre", "edit genre directly", "revise lyrics", "edit lyrics directly", "merge snippets", "continue to next section"]
+
+                    yield '', new_history, new_messages, tool_query_args["section_name"], revised_lyrics, new_instrumental_tags, clips_to_continue, f'<audio controls><source src="{song_link}" type="audio/mp3"></audio>', generated_audios, buttons
+                
+                elif tool_function_name == 'merge_all_snippets':
+                    updated_clip_url, updated_lyrics, updated_tags, clips_list = concat_snippets(tool_query_args['last_snippet_id'])
+
+                    if updated_clip_url == "Timeout":
+                        # Handle the timeout case
+                        tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': 'The operation timed out. Please try again later.'}
+                        audio_message = {'role': 'assistant', 'content': 'The concatenation operation took too long and was stopped. Please try again later or rephrase your request.'}
+                        new_messages = messages + [tool_message, audio_message]
+                        new_history = messages_to_history(new_messages)
+                        yield '', new_history, new_messages, '', '', '', None, None, generated_audios, []
+                        return
+                    
+                    if "still streaming" in updated_clip_url:
+                        tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': f'still streaming, try again later'}
+                        audio_message = {'role': 'assistant', 'content': f'Unfortunately the generated clip audio is still being streamed, so you can merge later when it is fully generated.'}
+
+                        new_messages = messages + [tool_message, audio_message]
+                        new_history = messages_to_history(new_messages)
+
+                        clips_to_continue = gr.Dropdown(label='Clip to continue', value = "", choices=[x[3] for x in generated_audios]+[""], interactive=True)
+
+
+                        yield '', new_history, new_messages, "", "", "", clips_to_continue, None, generated_audios, []
+
+                    else:
+                        if "https://audiopipe.suno.ai/?item_id=" in updated_clip_url:
+                            updated_clip_id = updated_clip_url.split("https://audiopipe.suno.ai/?item_id=")[1]
+                        elif "https://cdn1.suno.ai/" in updated_clip_url:
+                            updated_clip_id = updated_clip_url.split("https://cdn1.suno.ai/")[1].split(".mp3")[0]
+                        else:
+                            updated_clip_id = "unknown"
+
+                        #pass this info in new tool and assistant message
+                        tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': f'updated clip id: {updated_clip_id}\nupdated lyrics: {updated_lyrics}\nupdated clips path: {clips_list}'}
+                        audio_message = {'role': 'assistant', 'content': f'Sure! All the clips are now merged. <p>updated lyrics: {updated_lyrics}</p><audio controls autoplay><source src="{updated_clip_url}" type="audio/mp3"></audio><p>updated clips path: {clips_list}</p>'}
+
+                        sections_list = [line for line in current_lyrics.split('\n') if line.startswith('[') and line.endswith(']')]
+
+
+                        new_messages = messages + [tool_message, audio_message]
+                        new_history = messages_to_history(new_messages)
+
+                        if len(sections_list) > 0:
+                            section_name = "Merge up to " + sections_list[-1]
+                        else: 
+                            section_name = "Merge up to latest section"
+                        section_name = determine_title(section_name, generated_audios)
+
+                        # current_section, current_lyrics, curr_tags, clip_to_continue, curr_audio
+                        generated_audios.append((updated_clip_url, updated_lyrics, updated_tags, section_name, "streaming"))
 
-                    yield '', new_history, new_messages, "", "", "", clips_to_continue, None, generated_audios, []
+                        generated_audios = update_song_links(generated_audios)
+                        clips_to_continue = gr.Dropdown(label='Clip to continue', value = "", choices=[x[3] for x in generated_audios]+[""], interactive=True)
 
-                else:
+
+                        yield '', new_history, new_messages, section_name, updated_lyrics, updated_tags, clips_to_continue, f'<audio controls><source src="{updated_clip_url}" type="audio/mp3"></audio>', generated_audios, []
+                elif tool_function_name == 'finish_full_song':
+                    ## args are sections_to_be_written, relevant_ideas, last_snippet_id, sni
+                    
+                    ## STEP 0: POP out instrumental args
+                    snippet_instrumental_tags = tool_query_args.pop('snippet_instrumental_tags', None)
+                    snippet_clip_to_continue_from = tool_query_args.pop('snippet_clip_to_continue_from', None)
+
+                    if isinstance(tool_query_args['sections_written'], str):
+                        current_lyrics = tool_query_args['sections_written']
+                    elif isinstance(tool_query_args['sections_written'], list):
+                        current_lyrics = "\n".join(tool_query_args['sections_written'])
+                    else:
+                        current_lyrics = ""
+
+                    ## STEP 1: WRITE ALL LYRICS using songwriterAssistant
+                    remaining_lyrics = songwriterAssistant.write_all_lyrics(**tool_query_args)
+                    full_lyrics = current_lyrics + remaining_lyrics + "\n[End]"
+
+
+                    # current_section, current_lyrics, curr_tags, clip_to_continue, curr_audio
+                    # yield '', history, messages, "Full Song", full_lyrics, snippet_instrumental_tags, snippet_clip_to_continue_from, None, generated_audios, []
+
+                    ## STEP 2: MAKE SONG FOR REMAINING LYRICS
+                    song_link = make_song(remaining_lyrics, snippet_instrumental_tags, snippet_clip_to_continue_from)
+
+                    #tool and assistant message
+                    tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': f'{full_lyrics}'}
+                    audio_message = {'role': 'assistant', 'content': f'New snippet: \n <audio controls autoplay><source src="{song_link}" type="audio/mp3"></audio>'}
+
+                    new_messages = messages + [tool_message, audio_message]
+                    new_history = messages_to_history(new_messages)
+                    generated_audios.append((song_link, remaining_lyrics, snippet_instrumental_tags, "Rest of Song", "streaming"))
+
+                    yield '', new_history, new_messages, "Rest of Song", full_lyrics, snippet_instrumental_tags, snippet_clip_to_continue_from, song_link, generated_audios, []
+
+                    ## STEP 3: MERGE FULL SONG
+                    if snippet_clip_to_continue_from not in [None, ""]:
+                        updated_clip_url = "still streaming"
+                        while "still streaming" in updated_clip_url:
+                            if "https://audiopipe.suno.ai/?item_id=" in song_link:
+                                clip_id = song_link.split("https://audiopipe.suno.ai/?item_id=")[1]
+                            else:
+                                clip_id = updated_clip_url.split("https://cdn1.suno.ai/")[1].split(".mp3")[0]
+                            updated_clip_url, updated_lyrics, updated_tags, clips_list = concat_snippets(clip_id)
+                    else:
+                        updated_clip_url, updated_lyrics, clips_list = song_link, remaining_lyrics, []
+                    ## YIELD UPDATED CLIP URL, LYRICS, AND CLIPS LIST
                     if "https://audiopipe.suno.ai/?item_id=" in updated_clip_url:
                         updated_clip_id = updated_clip_url.split("https://audiopipe.suno.ai/?item_id=")[1]
                     elif "https://cdn1.suno.ai/" in updated_clip_url:
@@ -365,192 +490,149 @@ def model_chat(genre_input, query: Optional[str], history: Optional[History], me
                     else:
                         updated_clip_id = "unknown"
 
-                    #pass this info in new tool and assistant message
+                    #tool and assistant message
                     tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': f'updated clip id: {updated_clip_id}\nupdated lyrics: {updated_lyrics}\nupdated clips path: {clips_list}'}
-                    audio_message = {'role': 'assistant', 'content': f'Sure! All the clips are now merged. <p>updated lyrics: {updated_lyrics}</p><audio controls autoplay><source src="{updated_clip_url}" type="audio/mp3"></audio><p>updated clips path: {clips_list}</p>'}
-
-                    sections_list = [line for line in current_lyrics.split('\n') if line.startswith('[') and line.endswith(']')]
-
+                    audio_message = {'role': 'assistant', 'content': f'All done! Thank you for participating :) \nFinal Lyrics: {full_lyrics} \nFinal song: <audio controls autoplay><source src="{updated_clip_url}" type="audio/mp3"></audio>'}
 
                     new_messages = messages + [tool_message, audio_message]
                     new_history = messages_to_history(new_messages)
-
-                    if len(sections_list) > 0:
-                        section_name = "Merge up to " + sections_list[-1]
-                    else: 
-                        section_name = "Merge up to latest section"
-                    section_name = determine_title(section_name, generated_audios)
-
-                    # current_section, current_lyrics, curr_tags, clip_to_continue, curr_audio
-                    generated_audios.append((updated_clip_url, updated_lyrics, updated_tags, section_name, "streaming"))
-
+                    generated_audios.append((updated_clip_url, updated_lyrics, updated_tags, "Full Song", "streaming"))
                     generated_audios = update_song_links(generated_audios)
                     clips_to_continue = gr.Dropdown(label='Clip to continue', value = "", choices=[x[3] for x in generated_audios]+[""], interactive=True)
 
+                    yield '', new_history, new_messages, "Full Song", full_lyrics, snippet_instrumental_tags, clips_to_continue, f'<audio controls><source src="{song_link}" type="audio/mp3"></audio>', generated_audios, []
 
-                    yield '', new_history, new_messages, section_name, updated_lyrics, updated_tags, clips_to_continue, f'<audio controls><source src="{updated_clip_url}" type="audio/mp3"></audio>', generated_audios, []
-            elif tool_function_name == 'finish_full_song':
-                ## args are sections_to_be_written, relevant_ideas, last_snippet_id, sni
-                
-                ## STEP 0: POP out instrumental args
-                snippet_instrumental_tags = tool_query_args.pop('snippet_instrumental_tags', None)
-                snippet_clip_to_continue_from = tool_query_args.pop('snippet_clip_to_continue_from', None)
-
-                if isinstance(tool_query_args['sections_written'], str):
-                    current_lyrics = tool_query_args['sections_written']
-                elif isinstance(tool_query_args['sections_written'], list):
-                    current_lyrics = "\n".join(tool_query_args['sections_written'])
-                else:
-                    current_lyrics = ""
+                elif tool_function_name == 'get_audio_snippet':
+                    #detangle tool_query_args dict
+                    snippet_lyrics = tool_query_args['snippet_lyrics'] + "\n[End]"
+                    snippet_instrumental_tags = tool_query_args['snippet_instrumental_tags']
 
-                ## STEP 1: WRITE ALL LYRICS using songwriterAssistant
-                remaining_lyrics = songwriterAssistant.write_all_lyrics(**tool_query_args)
-                full_lyrics = current_lyrics + remaining_lyrics + "\n[End]"
+                    snippet_clip_to_continue_from = tool_query_args.get('snippet_clip_to_continue_from', None)
+                    song_link = make_song(snippet_lyrics, snippet_instrumental_tags, snippet_clip_to_continue_from)
 
+                    if song_link == "Timeout":
+                        tool_message = {
+                            'role': 'tool',
+                            'tool_call_id': tool_call_id,
+                            'name': tool_function_name,
+                            'content': 'The request to generate the audio snippet timed out. Please try again later.'
+                        }
 
-                # current_section, current_lyrics, curr_tags, clip_to_continue, curr_audio
-                yield '', history, messages, "Full Song", full_lyrics, snippet_instrumental_tags, snippet_clip_to_continue_from, None, generated_audios, []
+                        new_messages = messages + [tool_message]
 
-                ## STEP 2: MAKE SONG FOR REMAINING LYRICS
-                song_link = make_song(remaining_lyrics, snippet_instrumental_tags, snippet_clip_to_continue_from)
+                        model_response_with_function_call = oai_client.chat.completions.create(
+                            model="gpt-4o",
+                            messages=new_messages,
+                        )  # get a new response from the model where it can see the function response
+                        current_response = model_response_with_function_call.choices[0].message.content
 
-                #tool and assistant message
-                tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': f'{full_lyrics}'}
-                audio_message = {'role': 'assistant', 'content': f'New snippet: \n <audio controls autoplay><source src="{song_link}" type="audio/mp3"></audio>'}
+                        role = "assistant"
+                        new_messages = new_messages + [{'role': role, 'content': current_response}]
+                        new_history = messages_to_history(new_messages)
 
-                new_messages = messages + [tool_message, audio_message]
-                new_history = messages_to_history(new_messages)
-                generated_audios.append((song_link, remaining_lyrics, snippet_instrumental_tags, "Rest of Song", "streaming"))
+                        generated_audios = update_song_links(generated_audios)
+                        clips_to_continue = gr.Dropdown(label='Clip to continue', value = "", choices=[x[3] for x in generated_audios]+[""], interactive=True)
+                        buttons = []
 
-                yield '', new_history, new_messages, "Rest of Song", full_lyrics, snippet_instrumental_tags, snippet_clip_to_continue_from, song_link, generated_audios, []
+                        yield '', new_history, new_messages, snippet_lyrics.split("\n")[0], snippet_lyrics, snippet_instrumental_tags, clips_to_continue, None, generated_audios, buttons
+                        
+                        return
 
-                ## STEP 3: MERGE FULL SONG
-                if snippet_clip_to_continue_from not in [None, ""]:
-                    updated_clip_url = "still streaming"
-                    while "still streaming" in updated_clip_url:
-                        if "https://audiopipe.suno.ai/?item_id=" in song_link:
-                            clip_id = song_link.split("https://audiopipe.suno.ai/?item_id=")[1]
-                        else:
-                            clip_id = updated_clip_url.split("https://cdn1.suno.ai/")[1].split(".mp3")[0]
-                        updated_clip_url, updated_lyrics, updated_tags, clips_list = concat_snippets(clip_id)
-                else:
-                    updated_clip_url, updated_lyrics, clips_list = song_link, remaining_lyrics, []
-                ## YIELD UPDATED CLIP URL, LYRICS, AND CLIPS LIST
-                if "https://audiopipe.suno.ai/?item_id=" in updated_clip_url:
-                    updated_clip_id = updated_clip_url.split("https://audiopipe.suno.ai/?item_id=")[1]
-                elif "https://cdn1.suno.ai/" in updated_clip_url:
-                    updated_clip_id = updated_clip_url.split("https://cdn1.suno.ai/")[1].split(".mp3")[0]
-                else:
-                    updated_clip_id = "unknown"
-
-                #tool and assistant message
-                tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': f'updated clip id: {updated_clip_id}\nupdated lyrics: {updated_lyrics}\nupdated clips path: {clips_list}'}
-                audio_message = {'role': 'assistant', 'content': f'All done! Thank you for participating :) \nFinal Lyrics: {full_lyrics} \nFinal song: <audio controls autoplay><source src="{updated_clip_url}" type="audio/mp3"></audio>'}
-
-                new_messages = messages + [tool_message, audio_message]
-                new_history = messages_to_history(new_messages)
-                generated_audios.append((updated_clip_url, updated_lyrics, updated_tags, "Full Song", "streaming"))
-                generated_audios = update_song_links(generated_audios)
-                clips_to_continue = gr.Dropdown(label='Clip to continue', value = "", choices=[x[3] for x in generated_audios]+[""], interactive=True)
-
-                yield '', new_history, new_messages, "Full Song", full_lyrics, snippet_instrumental_tags, clips_to_continue, f'<audio controls><source src="{song_link}" type="audio/mp3"></audio>', generated_audios, []
-
-            elif tool_function_name == 'get_audio_snippet':
-                #detangle tool_query_args dict
-                snippet_lyrics = tool_query_args['snippet_lyrics'] + "\n[End]"
-                snippet_instrumental_tags = tool_query_args['snippet_instrumental_tags']
-
-                snippet_clip_to_continue_from = tool_query_args.get('snippet_clip_to_continue_from', None)
-                song_link = make_song(snippet_lyrics, snippet_instrumental_tags, snippet_clip_to_continue_from)
-
-
-                if "still streaming" in song_link:
-                    tool_message = {
-                        'role': 'tool',
-                        'tool_call_id': tool_call_id,
-                        'name': tool_function_name,
-                        'content': 'The snippet to extend is still streaming. Please try generating this audio snippet in a little bit.'
-                    }
+                    if "still streaming" in song_link:
+                        tool_message = {
+                            'role': 'tool',
+                            'tool_call_id': tool_call_id,
+                            'name': tool_function_name,
+                            'content': 'The snippet to extend is still streaming. Please try generating this audio snippet in a little bit.'
+                        }
 
-                    new_messages = messages + [tool_message]
+                        new_messages = messages + [tool_message]
 
-                    model_response_with_function_call = oai_client.chat.completions.create(
-                        model="gpt-4o",
-                        messages=new_messages,
-                    )  # get a new response from the model where it can see the function response
-                    current_response = model_response_with_function_call.choices[0].message.content
+                        model_response_with_function_call = oai_client.chat.completions.create(
+                            model="gpt-4o",
+                            messages=new_messages,
+                        )  # get a new response from the model where it can see the function response
+                        current_response = model_response_with_function_call.choices[0].message.content
 
-                    role = "assistant"
-                    new_messages = new_messages + [{'role': role, 'content': current_response}]
-                    new_history = messages_to_history(new_messages)
+                        role = "assistant"
+                        new_messages = new_messages + [{'role': role, 'content': current_response}]
+                        new_history = messages_to_history(new_messages)
 
-                    generated_audios = update_song_links(generated_audios)
-                    clips_to_continue = gr.Dropdown(label='Clip to continue', value = "", choices=[x[3] for x in generated_audios]+[""], interactive=True)
-                    buttons = []
+                        generated_audios = update_song_links(generated_audios)
+                        clips_to_continue = gr.Dropdown(label='Clip to continue', value = "", choices=[x[3] for x in generated_audios]+[""], interactive=True)
+                        buttons = []
 
-                    yield '', new_history, new_messages, snippet_lyrics.split("\n")[0], snippet_lyrics, snippet_instrumental_tags, clips_to_continue, None, generated_audios, buttons
+                        yield '', new_history, new_messages, snippet_lyrics.split("\n")[0], snippet_lyrics, snippet_instrumental_tags, clips_to_continue, None, generated_audios, buttons
+                        
+                        return
                     
-                    return
-                
 
-                if "no clip with that ID" in song_link:
-                    tool_message = {
-                        'role': 'tool',
-                        'tool_call_id': tool_call_id,
-                        'name': tool_function_name,
-                        'content': 'The clip ID was incorrect, maybe clarify with the user.'
-                    }
+                    if "no clip with that ID" in song_link:
+                        tool_message = {
+                            'role': 'tool',
+                            'tool_call_id': tool_call_id,
+                            'name': tool_function_name,
+                            'content': 'The clip ID was incorrect, maybe clarify with the user.'
+                        }
 
-                    new_messages = messages + [tool_message]
+                        new_messages = messages + [tool_message]
 
-                    model_response_with_function_call = oai_client.chat.completions.create(
-                        model="gpt-4o",
-                        messages=new_messages,
-                    )  # get a new response from the model where it can see the function response
-                    current_response = model_response_with_function_call.choices[0].message.content
+                        model_response_with_function_call = oai_client.chat.completions.create(
+                            model="gpt-4o",
+                            messages=new_messages,
+                        )  # get a new response from the model where it can see the function response
+                        current_response = model_response_with_function_call.choices[0].message.content
 
-                    role = "assistant"
-                    new_messages = new_messages + [{'role': role, 'content': current_response}]
-                    new_history = messages_to_history(new_messages)
+                        role = "assistant"
+                        new_messages = new_messages + [{'role': role, 'content': current_response}]
+                        new_history = messages_to_history(new_messages)
 
-                    generated_audios = update_song_links(generated_audios)
-                    clips_to_continue = gr.Dropdown(label='Clip to continue', value = "", choices=[x[3] for x in generated_audios]+[""], interactive=True)
-                    buttons = []
+                        generated_audios = update_song_links(generated_audios)
+                        clips_to_continue = gr.Dropdown(label='Clip to continue', value = "", choices=[x[3] for x in generated_audios]+[""], interactive=True)
+                        buttons = []
 
-                    yield '', new_history, new_messages, snippet_lyrics.split("\n")[0], snippet_lyrics, snippet_instrumental_tags, clips_to_continue, None, generated_audios, buttons
+                        yield '', new_history, new_messages, snippet_lyrics.split("\n")[0], snippet_lyrics, snippet_instrumental_tags, clips_to_continue, None, generated_audios, buttons
+                        
+                        return
+                    print("MAKE SONG IS DONE")
+                    ## filter out suno link from tool query arg
+                    clip_id = song_link.split("https://audiopipe.suno.ai/?item_id=")[1]
+
+                    tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': f'snippet lyrics: {snippet_lyrics}\ninstrumental tags: {tool_query_args["snippet_instrumental_tags"]}, clip id: {clip_id}'}
+                    audio_message_content = "Here's what I've come up with:\n" + snippet_lyrics + '\n\n' + f'<audio controls autoplay><source src="{song_link}" type="audio/mp3"></audio><p>instrumental tags: {tool_query_args["snippet_instrumental_tags"]}</p>'
+                    audio_message_content += f'<p>continued from clip: {snippet_clip_to_continue_from}</p>'
+                    audio_message_content += "What do you think?"
+                    audio_message = {'role': 'assistant', 'content': audio_message_content}
                     
-                    return
-                print("MAKE SONG IS DONE")
-                ## filter out suno link from tool query arg
-                clip_id = song_link.split("https://audiopipe.suno.ai/?item_id=")[1]
-
-                tool_message = {'role': 'tool', 'tool_call_id': tool_call_id, 'name': tool_function_name, 'content': f'snippet lyrics: {snippet_lyrics}\ninstrumental tags: {tool_query_args["snippet_instrumental_tags"]}, clip id: {clip_id}'}
-                audio_message_content = "Here's what I've come up with:\n" + snippet_lyrics + '\n\n' + f'<audio controls autoplay><source src="{song_link}" type="audio/mp3"></audio><p>instrumental tags: {tool_query_args["snippet_instrumental_tags"]}</p>'
-                audio_message_content += f'<p>continued from clip: {snippet_clip_to_continue_from}</p>'
-                audio_message_content += "What do you think?"
-                audio_message = {'role': 'assistant', 'content': audio_message_content}
-                
 
-                section_name = snippet_lyrics.split("\n")[0].strip('[]* ')
-                section_name = determine_title(section_name, generated_audios)
+                    section_name = snippet_lyrics.split("\n")[0].strip('[]* ')
+                    section_name = determine_title(section_name, generated_audios)
 
-                #audio_message = {'role': 'assistant', 'content': gr.Audio(value=song_link, label=section_name, interactive=False, show_label=False, waveform_options={"show_controls": False})}
-                new_messages = messages + [tool_message, audio_message]
-                new_history = messages_to_history(new_messages)
-                print("AUDIO MESSAGE DONE")
-                generated_audios.append((song_link, snippet_lyrics, snippet_instrumental_tags, section_name, "streaming"))
+                    #audio_message = {'role': 'assistant', 'content': gr.Audio(value=song_link, label=section_name, interactive=False, show_label=False, waveform_options={"show_controls": False})}
+                    new_messages = messages + [tool_message, audio_message]
+                    new_history = messages_to_history(new_messages)
+                    print("AUDIO MESSAGE DONE")
+                    generated_audios.append((song_link, snippet_lyrics, snippet_instrumental_tags, section_name, "streaming"))
 
-                generated_audios = update_song_links(generated_audios)
+                    generated_audios = update_song_links(generated_audios)
 
-                buttons = ["revise lyrics", "revise genre", "merge snippets", "continue to next section"]
-                clips_to_continue = gr.Dropdown(label='Clip to continue', value = snippet_clip_to_continue_from, choices=[x[3] for x in generated_audios]+[""], interactive=True)
+                    buttons = ["revise lyrics", "edit lyrics directly", "revise genre", "edit genre directly", "merge snippets", "continue to next section"]
+                    clips_to_continue = gr.Dropdown(label='Clip to continue', value = snippet_clip_to_continue_from, choices=[x[3] for x in generated_audios]+[""], interactive=True)
 
-                yield '', new_history, new_messages, snippet_lyrics.split("\n")[0], snippet_lyrics, snippet_instrumental_tags, clips_to_continue, f'<audio controls><source src="{song_link}" type="audio/mp3"></audio>', generated_audios, buttons
+                    yield '', new_history, new_messages, snippet_lyrics.split("\n")[0], snippet_lyrics, snippet_instrumental_tags, clips_to_continue, f'<audio controls><source src="{song_link}" type="audio/mp3"></audio>', generated_audios, buttons
 
 
-            else: 
-                print(f"Error: function {tool_function_name} does not exist")
-        
+                else: 
+                    print(f"Error: function {tool_function_name} does not exist")
+        except Exception as e:
+            error_message = {
+                'role': 'assistant',
+                'content': f"An error occurred while processing your request: {str(e)}. Please re-phrase your request and try again."
+            }
+            messages_filtered.append(error_message)
+            yield '', messages_to_history(messages), messages_filtered, '', '', '', '', None, generated_audios, []
+
+            
     else: 
         # Model did not identify a function to call, result can be returned to the user 
         current_response = response_message.choices[0].message.content