Update app.py

app.py

@@ -6,14 +6,95 @@ import re
 import sys
 from typing import List, Generator
 import os
-from langchain_openai import ChatOpenAI
 from dotenv import load_dotenv
+import threading
 
 load_dotenv()
 
+class OutputParser:
+    def __init__(self):
+        self.buffer = ""
+        self.current_agent = None
+        
+    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(
+                role="assistant",
+                content=f"Starting work...",
+                metadata={"title": f"🤖 {self.current_agent}"}
+            ))
+        
+        # Look for task information with full task list
+        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(
+                role="assistant",
+                content=task_content,
+                metadata={"title": f"📋 Task for {self.current_agent}"}
+            ))
+        
+        # Look for thought processes
+        thought_match = re.search(r'Thought: (.*?)(?=\nAction:|$)', cleaned_text, re.DOTALL)
+        if thought_match and self.current_agent:
+            thought_content = thought_match.group(1).strip()
+            messages.append(ChatMessage(
+                role="assistant",
+                content=thought_content,
+                metadata={"title": f"💭 {self.current_agent}'s Thoughts"}
+            ))
+        
+        # Look for final answers from non-Editor agents
+        if "Final Answer:" in cleaned_text and self.current_agent != "Editor":
+            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()
+                messages.append(ChatMessage(
+                    role="assistant",
+                    content=answer_content,
+                    metadata={"title": f"💡 Output from {self.current_agent}"}
+                ))
+        
+        # Special handling for Editor's final answer (the final article)
+        elif "Final Answer:" in cleaned_text and self.current_agent == "Editor":
+            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()
+                # First send the metadata marker
+                messages.append(ChatMessage(
+                    role="assistant",
+                    content="Final article is ready!",
+                    metadata={"title": "📝 Final Article"}
+                ))
+                # Then send the actual content without metadata
+                messages.append(ChatMessage(
+                    role="assistant",
+                    content=answer_content
+                ))
+        
+        return messages
+
+class StreamingCapture:
+    def __init__(self):
+        self.buffer = ""
+    
+    def write(self, text):
+        self.buffer += text
+        return len(text)
+    
+    def flush(self):
+        pass
+
 class ArticleCrew:
     def __init__(self):
-        # Agent definitions remain the same
+        # Initialize agents
         self.planner = Agent(
             role="Content Planner",
             goal="Plan engaging and factually accurate content on {topic}",
@@ -44,7 +125,6 @@ class ArticleCrew:
         self.output_parser = OutputParser()
 
     def create_tasks(self, topic: str):
-        # Task definitions remain the same
         plan_task = Task(
             description=(
                 f"1. Prioritize the latest trends, key players, and noteworthy news on {topic}.\n"
@@ -83,38 +163,46 @@ class ArticleCrew:
             verbose=2
         )
 
-        class StreamCapture:
-            def __init__(self):
-                self.data = []
-                self.current_chunk = ""
-            
-            def write(self, text):
-                self.current_chunk += text
-                if "\n" in text:
-                    self.data.append(self.current_chunk)
-                    self.current_chunk = ""
-                return len(text)
-
-        stream = StreamCapture()
+        capture = StreamingCapture()
         original_stdout = sys.stdout
-        sys.stdout = stream
+        sys.stdout = capture
         
         try:
-            result = crew.kickoff(inputs={"topic": topic})
+            # Start the crew task in a separate thread to not block streaming
+            result_container = []
             
-            # Process intermediate outputs
-            for chunk in stream.data:
-                messages = self.output_parser.parse_output(chunk)
-                if messages:
-                    for msg in messages:
-                        yield [msg]
+            def run_crew():
+                try:
+                    result = crew.kickoff(inputs={"topic": topic})
+                    result_container.append(result)
+                except Exception as e:
+                    result_container.append(e)
             
-            # Send final result
-            yield [ChatMessage(
-                role="assistant",
-                content=result,
-                metadata={"title": "📝 Final Article"}
-            )]
+            thread = threading.Thread(target=run_crew)
+            thread.start()
+            
+            # Stream output while the crew is working
+            last_processed = 0
+            while thread.is_alive() or last_processed < len(capture.buffer):
+                if len(capture.buffer) > last_processed:
+                    new_content = capture.buffer[last_processed:]
+                    messages = self.output_parser.parse_output(new_content)
+                    if messages:
+                        for msg in messages:
+                            yield [msg]
+                    last_processed = len(capture.buffer)
+                await asyncio.sleep(0.1)
+            
+            # Check if we got a result or an error
+            if result_container and not isinstance(result_container[0], Exception):
+                # Final messages already sent by the parser
+                pass
+            else:
+                yield [ChatMessage(
+                    role="assistant",
+                    content="An error occurred while generating the article.",
+                    metadata={"title": "❌ Error"}
+                )]
         
         finally:
             sys.stdout = original_stdout
@@ -141,16 +229,14 @@ def create_demo():
         )
         
         async def process_input(topic, history):
-            # Add user message as ChatMessage
             history.append(ChatMessage(role="user", content=f"Write an article about: {topic}"))
             yield history
             
-            # Process and add agent messages
             async for messages in article_crew.process_article(topic):
                 history.extend(messages)
                 yield history
         
-        btn = gr.Button("Write Article")
+        btn = gr.Button("Write Article", variant="primary")
         btn.click(
             process_input,
             inputs=[topic, chatbot],
@@ -161,4 +247,5 @@ def create_demo():
 
 if __name__ == "__main__":
     demo = create_demo()
+    demo.queue()
     demo.launch(debug=True, share=True)