improve agent

src/agent.py

@@ -61,7 +61,7 @@ Remember: Stay focused on audio-related assistance and use your specialized tool
 class AudioAgent:
     def __init__(
         self,
-        model_name: str = "gpt-4o",
+        model_name: str = "gpt-4.1",
         server_url: str = "https://agents-mcp-hackathon-audioeditor.hf.space/gradio_api/mcp/sse",
     ):
         load_dotenv()
@@ -96,6 +96,8 @@ class AudioAgent:
         Input Audio Files: {', '.join(input_audio_files) if input_audio_files else 'None'}
         """
 
-        return await self.agent.ainvoke(
+        res = await self.agent.ainvoke(
             {"messages": [{"role": "user", "content": input_context}]}
-        )
+        )
+
+        return res["structured_response"]

src/ui.py

@@ -33,63 +33,40 @@ def user_input(user_message, audio_files, history):
             
             audio_file_urls.append(get_share_url(file_path))
     
-    # For display purposes, show what audio files were uploaded
-    if audio_file_urls:
-        audio_list = "\n".join([f"🎵 Uploaded: {url.split('/')[-1]}" for url in audio_file_urls])
-        combined_message = f"{user_message}\n\n{audio_list}" if user_message.strip() else f"Process uploaded audio files:\n{audio_list}"
-    else:
-        combined_message = user_message
-    
-    history.append({"role": "user", "content": combined_message})
+    # Add user message to history (no uploaded file display)
+    history.append({"role": "user", "content": user_message})
     return "", [], history, audio_file_urls
 
 async def bot_response(history, audio_file_urls):
     """
-    Generate bot response using the test agent
+    Generate bot response using the agent
     """
     if not history or history[-1]["role"] != "user":
-        return history
+        return history, []
     
-    # Get the actual user message (without the audio file display text)
+    # Get the user message
     user_message = history[-1]["content"]
     
-    # Clean the user message by removing the uploaded file display text
-    if "🎵 Uploaded:" in user_message:
-        lines = user_message.split('\n')
-        clean_lines = []
-        for line in lines:
-            if not line.strip().startswith('🎵 Uploaded:'):
-                clean_lines.append(line)
-        user_message = '\n'.join(clean_lines).strip()
-        
-        # If message is empty after cleaning, provide default message
-        if not user_message:
-            user_message = "Please process these audio files"
+    # If message is empty but we have audio files, provide default message
+    if not user_message.strip() and audio_file_urls:
+        user_message = "Please process these audio files"
     
     try:
-        # Use the test agent's run_agent method with separate parameters
+        # Use the agent's run_agent method
         result = await agent.run_agent(user_message, audio_file_urls or [])
 
         # Extract the final response and audio files from the result
         final_response = result.get("final_response", "")
         output_audio_files = result.get("output_audio_files", [])
         
-        # Format the response
-        formatted_content = final_response
-        
-        # Add processed audio files section if any
-        if output_audio_files:
-            formatted_content += "\n\n## 🎵 Generated Audio Files\n\n"
-            for audio_file in output_audio_files:
-                formatted_content += f"Audio Ready: {audio_file}\n"
+        # Add assistant response to history (only final_response)
+        history.append({"role": "assistant", "content": final_response})
         
-        # Add assistant response to history
-        history.append({"role": "assistant", "content": formatted_content.rstrip()})
+        return history, output_audio_files
         
     except Exception as e:
         history.append({"role": "assistant", "content": f"❌ **Error**: {e}"})
-    
-    return history
+        return history, []
 
 def bot_response_sync(history, audio_file_urls):
     """
@@ -107,7 +84,7 @@ def create_interface():
         title="Audio Agent - Professional Audio Processing", 
         theme=gr.themes.Soft(),
         css="""
-        .processed-audio {
+        .output-audio {
             background: #f0f9ff;
             border: 1px solid #0891b2;
             border-radius: 8px;
@@ -153,6 +130,18 @@ def create_interface():
                 - "How does audio normalization work?"
                 """)
         
+        # Output audio files section
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### 🎵 Generated Audio Files")
+                output_audio_files = gr.File(
+                    file_count="multiple",
+                    label="Download Generated Audio",
+                    interactive=False,
+                    visible=True,
+                    elem_classes=["output-audio"]
+                )
+        
         with gr.Row():
             msg = gr.Textbox(
                 label="Describe what you want to do",
@@ -201,15 +190,19 @@ def create_interface():
             new_msg, new_files, updated_history, audio_urls = user_input(message, files, history)
             return new_msg, new_files, updated_history, audio_urls
         
+        def handle_bot_response(history, audio_urls):
+            updated_history, output_files = bot_response_sync(history, audio_urls)
+            return updated_history, output_files
+        
         msg.submit(
             handle_submit,
             [msg, audio_files, chatbot], 
             [msg, audio_files, chatbot, audio_urls_state], 
             queue=False
         ).then(
-            bot_response_sync, 
+            handle_bot_response, 
             [chatbot, audio_urls_state], 
-            chatbot
+            [chatbot, output_audio_files]
         )
         
         send_btn.click(
@@ -218,19 +211,19 @@ def create_interface():
             [msg, audio_files, chatbot, audio_urls_state], 
             queue=False
         ).then(
-            bot_response_sync, 
+            handle_bot_response, 
             [chatbot, audio_urls_state], 
-            chatbot
+            [chatbot, output_audio_files]
         )
         
         # Clear chat
         def clear_chat():
-            return [], [], []
+            return [], [], [], []
         
         clear_btn.click(
             clear_chat,
             None, 
-            [chatbot, audio_files, audio_urls_state], 
+            [chatbot, audio_files, audio_urls_state, output_audio_files], 
             queue=False
         )