Update app.py

app.py

@@ -1,3 +1,4 @@
+# This works even better 
 from crewai import Agent, Task, Crew
 import gradio as gr
 from gradio import ChatMessage
@@ -14,18 +15,48 @@ class OutputParser:
     def __init__(self):
         self.buffer = ""
         self.current_agent = None
-        self.final_article_sent = False  # Add this flag
+        self.final_article_sent = False
+        self.message_queue = {
+            "Content Planner": [],
+            "Content Writer": [],
+            "Editor": []
+        }
+        self.agent_sequence = ["Content Planner", "Content Writer", "Editor"]
         
+    def format_output(self, raw_content: str, agent_name: str) -> str:
+        """Format the output content based on agent type."""
+        if agent_name == "Content Planner":
+            # Clean up the planner's output to make it more readable
+            lines = raw_content.split('\n')
+            formatted_lines = []
+            for line in lines:
+                # Remove number prefixes and clean up
+                line = re.sub(r'^\d+\.\s*', '', line.strip())
+                # Make text size normal by removing markdown formatting
+                line = re.sub(r'^#+\s*', '', line)
+                if line:
+                    formatted_lines.append(line)
+            return '\n\n'.join(formatted_lines)
+        
+        elif agent_name == "Content Writer":
+            # Clean up writer's output to make it more readable
+            # Remove markdown headers but keep the text
+            content = re.sub(r'^#+\s*(.+)$', r'\1', raw_content, flags=re.MULTILINE)
+            # Remove multiple newlines
+            content = re.sub(r'\n{3,}', '\n\n', content)
+            return content.strip()
+            
+        return raw_content.strip()
+
     def parse_output(self, text: str) -> List[ChatMessage]:
         messages = []
-        # Clean ANSI codes
         cleaned_text = re.sub(r'\x1B\[[0-9;]*[mK]', '', text)
         
         # Look for working agent declarations
         agent_match = re.search(r'\[DEBUG\]: == Working Agent: (.*?)(?=\n|$)', cleaned_text)
         if agent_match:
             self.current_agent = agent_match.group(1)
-            messages.append(ChatMessage(
+            self.message_queue[self.current_agent].append(ChatMessage(
                 role="assistant",
                 content=f"Starting work...",
                 metadata={"title": f"🤖 {self.current_agent}"}
@@ -35,38 +66,45 @@ class OutputParser:
         task_match = re.search(r'\[INFO\]: == Starting Task: (.*?)(?=\n\n|\n> Entering|$)', cleaned_text, re.DOTALL)
         if task_match and self.current_agent:
             task_content = task_match.group(1).strip()
-            messages.append(ChatMessage(
+            self.message_queue[self.current_agent].append(ChatMessage(
                 role="assistant",
                 content=task_content,
                 metadata={"title": f"📋 Task for {self.current_agent}"}
             ))
+
+        # Look for agent outputs in debug messages
+        debug_match = re.search(r'\[DEBUG\]: == \[(.*?)\] Task output: (.*?)(?=\[DEBUG\]|$)', cleaned_text, re.DOTALL)
+        if debug_match:
+            agent_name = debug_match.group(1)
+            output_content = debug_match.group(2).strip()
+            
+            # Format the output content
+            formatted_content = self.format_output(output_content, agent_name)
+            
+            if agent_name == "Editor" and not self.final_article_sent:
+                self.message_queue[agent_name].append(ChatMessage(
+                    role="assistant",
+                    content="Final article is ready!",
+                    metadata={"title": "📝 Final Article"}
+                ))
+                self.message_queue[agent_name].append(ChatMessage(
+                    role="assistant",
+                    content=formatted_content
+                ))
+                self.final_article_sent = True
+            elif agent_name != "Editor":
+                self.message_queue[agent_name].append(ChatMessage(
+                    role="assistant",
+                    content=formatted_content,
+                    metadata={"title": f"💡 Output from {agent_name}"}
+                ))
         
-        # Look for final answers from all agents
-        if "Final Answer:" in cleaned_text:
-            answer_match = re.search(r'Final Answer:\s*(.*?)(?=\n> Finished chain|$)', cleaned_text, re.DOTALL)
-            if answer_match:
-                answer_content = answer_match.group(1).strip()
+        # Return messages in the correct sequence
+        for agent in self.agent_sequence:
+            if self.message_queue[agent]:
+                messages.extend(self.message_queue[agent])
+                self.message_queue[agent] = []
                 
-                if self.current_agent == "Editor" and not self.final_article_sent:
-                    # This is the final article - send only once
-                    messages.append(ChatMessage(
-                        role="assistant",
-                        content="Final Article",
-                        metadata={"title": "📝 Final Article"}
-                    ))
-                    messages.append(ChatMessage(
-                        role="assistant",
-                        content=answer_content
-                    ))
-                    self.final_article_sent = True
-                elif self.current_agent != "Editor":
-                    # This is an intermediate output (Planner or Writer)
-                    messages.append(ChatMessage(
-                        role="assistant",
-                        content=answer_content,
-                        metadata={"title": f"💡 Output from {self.current_agent}"}
-                    ))
-        
         return messages
 
 class StreamingCapture:
@@ -236,4 +274,4 @@ def create_demo():
 if __name__ == "__main__":
     demo = create_demo()
     demo.queue()
-    demo.launch(debug=True, share=True)
+    demo.launch(debug=True, share=True)