Update chat_handler.py

chat_handler.py

@@ -230,27 +230,41 @@ class ChatHandler:
         current_tool_id = None
         current_server_name = None
         tool_start_time = None
-        main_response_content = ""
+        text_segments = []  # Collect text segments separately
         
         # Process Claude's response
         for content in response.content:
             if content.type == "text":
-                main_response_content += content.text
+                # Collect text segments but don't combine them yet
+                text_content = content.text
                 # Check if Claude indicated media was generated
-                if "MEDIA_GENERATED:" in content.text:
-                    media_match = re.search(r"MEDIA_GENERATED:\s*([^\s]+)", content.text)
+                if "MEDIA_GENERATED:" in text_content:
+                    media_match = re.search(r"MEDIA_GENERATED:\s*([^\s]+)", text_content)
                     if media_match:
                         media_url = media_match.group(1)
                         # Clean up the response text
-                        main_response_content = re.sub(r"MEDIA_GENERATED:\s*[^\s]+", "", main_response_content).strip()
+                        text_content = re.sub(r"MEDIA_GENERATED:\s*[^\s]+", "", text_content).strip()
                         logger.info(f"🎯 Claude indicated media generated: {media_url}")
                         # Add media as separate message
                         chat_messages.append(ChatMessage(
                             role="assistant", 
                             content={"path": media_url}
                         ))
+                
+                if text_content.strip():
+                    text_segments.append(text_content.strip())
             
             elif hasattr(content, 'type') and content.type == "mcp_tool_use":
+                # Add any accumulated text before tool use
+                if text_segments:
+                    combined_text = " ".join(text_segments)
+                    if combined_text.strip():
+                        chat_messages.append(ChatMessage(
+                            role="assistant",
+                            content=combined_text.strip()
+                        ))
+                    text_segments = []  # Reset
+                
                 tool_name = content.name
                 server_name = content.server_name
                 current_tool_id = getattr(content, 'id', 'unknown')
@@ -287,27 +301,14 @@ class ChatHandler:
                             msg.metadata["duration"] = round(duration, 2)
                         break
                 
+                media_url = None
                 if content.content:
                     result_content = content.content[0]
                     result_text = result_content.text if hasattr(result_content, 'text') else str(result_content)
                     
                     logger.info(f"📝 MCP tool result: {result_text[:200]}...")
                     
-                    # Create a result message with metadata
-                    result_msg = ChatMessage(
-                        role="assistant",
-                        content=result_text,
-                        metadata={
-                            "title": "📋 Tool Result",
-                            "parent_id": current_tool_id,
-                            "status": "done"
-                        }
-                    )
-                    
-                    chat_messages.append(result_msg)
-                    
-                    # Try to extract media from the result
-                    media_url = None
+                    # Try to extract media from the result first
                     if current_server_name and current_server_name in self.mcp_client.servers:
                         config = self.mcp_client.servers[current_server_name]
                         extracted_media = self.mcp_client._extract_media_from_mcp_response(result_text, config)
@@ -323,8 +324,31 @@ class ChatHandler:
                                 logger.info(f"🎯 Extracted media from MCP result (fallback): {media_url}")
                                 break
                     
-                    # Add media as separate message if found
-                    if media_url:
+                    # Only show tool result if no media was found or if it's not just file data
+                    if not media_url:
+                        # No media found, show the result
+                        chat_messages.append(ChatMessage(
+                            role="assistant",
+                            content=result_text,
+                            metadata={
+                                "title": "📋 Tool Result",
+                                "parent_id": current_tool_id,
+                                "status": "done"
+                            }
+                        ))
+                    else:
+                        # Media found, show success message instead of raw result
+                        chat_messages.append(ChatMessage(
+                            role="assistant",
+                            content="✅ Successfully generated content",
+                            metadata={
+                                "title": "📋 Tool Result",
+                                "parent_id": current_tool_id,
+                                "status": "done"
+                            }
+                        ))
+                        
+                        # Add media as separate message
                         chat_messages.append(ChatMessage(
                             role="assistant", 
                             content={"path": media_url}
@@ -342,14 +366,17 @@ class ChatHandler:
                         }
                     ))
         
-        # Add the main response if there's any text content
-        if main_response_content.strip():
-            chat_messages.append(ChatMessage(
-                role="assistant",
-                content=main_response_content.strip()
-            ))
-        elif not chat_messages:
-            # Fallback if no content was processed
+        # Add any remaining text segments after all processing
+        if text_segments:
+            combined_text = " ".join(text_segments)
+            if combined_text.strip():
+                chat_messages.append(ChatMessage(
+                    role="assistant",
+                    content=combined_text.strip()
+                ))
+        
+        # Fallback if no content was processed
+        if not chat_messages:
             chat_messages.append(ChatMessage(
                 role="assistant",
                 content="I understand your request and I'm here to help."