Update app.py

app.py

@@ -1,140 +1,55 @@
 from crewai import Agent, Task, Crew
 import gradio as gr
-from gradio import ChatMessage
 import asyncio
-import re
-import sys
-from typing import List, Generator
-import os
-from dotenv import load_dotenv
-import threading
+from typing import List, Generator, Any, Dict, Union
 from langchain_openai import ChatOpenAI
+import queue
+import threading
+import os
 
-class OutputParser:
+class AgentMessageQueue:
     def __init__(self):
-        self.buffer = ""
-        self.current_agent = None
-        self.final_article_sent = False
-        self.message_queue = {
-            "Content Planner": [],
-            "Content Writer": [],
-            "Editor": []
-        }
-        self.agent_sequence = ["Content Planner", "Content Writer", "Editor"]
+        self.message_queue = queue.Queue()
+        self.final_output = None
         
-    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)
+    def add_message(self, message: Dict):
+        self.message_queue.put(message)
         
-        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]:
+    def get_messages(self) -> List[Dict]:
         messages = []
-        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)
-            self.message_queue[self.current_agent].append(ChatMessage(
-                role="assistant",
-                content=f"Starting work...",
-                metadata={"title": f"🤖 {self.current_agent}"}
-            ))
-        
-        # Look for task information
-        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()
-            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}"}
-                ))
-        
-        # 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] = []
-                
+        while not self.message_queue.empty():
+            messages.append(self.message_queue.get())
         return messages
-
-class StreamingCapture:
-    def __init__(self):
-        self.buffer = ""
     
-    def write(self, text):
-        self.buffer += text
-        return len(text)
+    def set_final_output(self, output: str):
+        self.final_output = output
     
-    def flush(self):
-        pass
+    def get_final_output(self) -> str:
+        return self.final_output
 
 class ArticleCrew:
     def __init__(self, api_key: str = None):
         self.api_key = api_key
-        self.initialize_agents()
+        self.message_queue = AgentMessageQueue()
+        self.planner = None
+        self.writer = None
+        self.editor = None
     
-    def initialize_agents(self):
-        # Create a ChatOpenAI instance with the API key
+    def initialize_agents(self, topic: str):
+        if not self.api_key:
+            raise ValueError("OpenAI API key is required")
+            
+        os.environ["OPENAI_API_KEY"] = self.api_key
+        
         llm = ChatOpenAI(
-            openai_api_key=self.api_key,
             temperature=0.7,
             model="gpt-4"
         )
         
-        # Initialize agents with the LLM
         self.planner = Agent(
             role="Content Planner",
-            goal="Plan engaging and factually accurate content on {topic}",
-            backstory="You're working on planning a blog article about the topic: {topic}. "
+            goal=f"Plan engaging and factually accurate content on {topic}",
+            backstory=f"You're working on planning a blog article about the topic: {topic}. "
                      "You collect information that helps the audience learn something "
                      "and make informed decisions.",
             allow_delegation=False,
@@ -144,8 +59,8 @@ class ArticleCrew:
         
         self.writer = Agent(
             role="Content Writer",
-            goal="Write insightful and factually accurate opinion piece about the topic: {topic}",
-            backstory="You're working on writing a new opinion piece about the topic: {topic}. "
+            goal=f"Write insightful and factually accurate opinion piece about the topic: {topic}",
+            backstory=f"You're working on writing a new opinion piece about the topic: {topic}. "
                      "You base your writing on the work of the Content Planner.",
             allow_delegation=False,
             verbose=True,
@@ -160,10 +75,11 @@ class ArticleCrew:
             verbose=True,
             llm=llm
         )
-        
-        self.output_parser = OutputParser()
 
     def create_tasks(self, topic: str):
+        if not self.planner or not self.writer or not self.editor:
+            self.initialize_agents(topic)
+            
         plan_task = Task(
             description=(
                 f"1. Prioritize the latest trends, key players, and noteworthy news on {topic}.\n"
@@ -195,68 +111,131 @@ class ArticleCrew:
 
         return [plan_task, write_task, edit_task]
 
-    async def process_article(self, topic: str) -> Generator[List[ChatMessage], None, None]:
+    async def process_article(self, topic: str) -> Generator[List[Dict], None, None]:
+        def step_callback(output: Any) -> None:
+            try:
+                output_str = str(output).strip()
+                
+                # Extract agent name
+                if "# Agent:" in output_str:
+                    agent_name = output_str.split("# Agent:")[1].split("\n")[0].strip()
+                else:
+                    agent_name = "Agent"
+                
+                # Extract task or final answer
+                if "## Task:" in output_str:
+                    content = output_str.split("## Task:")[1].split("\n#")[0].strip()
+                    self.message_queue.add_message({
+                        "role": "assistant",
+                        "content": content,
+                        "metadata": {"title": f"📋 {agent_name}'s Task"}
+                    })
+                elif "## Final Answer:" in output_str:
+                    content = output_str.split("## Final Answer:")[1].strip()
+                    if agent_name == "Editor":
+                        # For Editor's final answer, store it for later
+                        self.message_queue.set_final_output(content)
+                    self.message_queue.add_message({
+                        "role": "assistant",
+                        "content": content,
+                        "metadata": {"title": f"✅ {agent_name}'s Output"}
+                    })
+                else:
+                    self.message_queue.add_message({
+                        "role": "assistant",
+                        "content": output_str,
+                        "metadata": {"title": f"💭 {agent_name} thinking"}
+                    })
+                    
+            except Exception as e:
+                print(f"Error in step_callback: {str(e)}")
+
+        def task_callback(output: Any) -> None:
+            try:
+                content = str(output)
+                if hasattr(output, 'agent'):
+                    agent_name = str(output.agent)
+                else:
+                    agent_name = "Agent"
+                
+                self.message_queue.add_message({
+                    "role": "assistant",
+                    "content": content.strip(),
+                    "metadata": {"title": f"✅ Task completed by {agent_name}"}
+                })
+                
+                # If this is the Editor's task completion, add the final article
+                if agent_name == "Editor":
+                    final_content = self.message_queue.get_final_output()
+                    if final_content:
+                        self.message_queue.add_message({
+                            "role": "assistant",
+                            "content": "Here's your completed article:",
+                            "metadata": {"title": "📝 Final Article"}
+                        })
+                        self.message_queue.add_message({
+                            "role": "assistant",
+                            "content": final_content
+                        })
+                        self.message_queue.add_message({
+                            "role": "assistant",
+                            "content": "Article generation completed!",
+                            "metadata": {"title": "✨ Complete"}
+                        })
+                
+            except Exception as e:
+                print(f"Error in task_callback: {str(e)}")
+
+        self.initialize_agents(topic)
+        
         crew = Crew(
             agents=[self.planner, self.writer, self.editor],
             tasks=self.create_tasks(topic),
-            verbose=2
+            verbose=True,
+            step_callback=step_callback,
+            task_callback=task_callback
         )
-
-        capture = StreamingCapture()
-        original_stdout = sys.stdout
-        sys.stdout = capture
         
-        try:
-            # Start the crew task in a separate thread to not block streaming
-            result_container = []
-            
-            def run_crew():
-                try:
-                    result = crew.kickoff(inputs={"topic": topic})
-                    result_container.append(result)
-                except Exception as e:
-                    result_container.append(e)
-            
-            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"}
-                )]
+        # Start notification
+        yield [{
+            "role": "assistant",
+            "content": "Starting work on your article...",
+            "metadata": {"title": "🚀 Process Started"}
+        }]
+        
+        # Run crew in a separate thread
+        result_container = []
+        def run_crew():
+            try:
+                result = crew.kickoff(inputs={"topic": topic})
+                result_container.append(result)
+            except Exception as e:
+                result_container.append(e)
+                print(f"Error occurred: {str(e)}")
+        
+        thread = threading.Thread(target=run_crew)
+        thread.start()
         
-        finally:
-            sys.stdout = original_stdout
+        # Stream messages while crew is working
+        while thread.is_alive() or not self.message_queue.message_queue.empty():
+            messages = self.message_queue.get_messages()
+            if messages:
+                yield messages
+            await asyncio.sleep(0.1)
 
 def create_demo():
-    article_crew = None  # Initialize as None
+    article_crew = None
     
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown("# 📝 AI Article Writing Crew")
         gr.Markdown("Watch as this AI Crew collaborates to create your article! This application utilizes [CrewAI](https://www.crewai.com/) agents: Content Planner, Content Writer, and Content Editor, to write an article on any topic you choose. To get started, enter your OpenAI API Key below and press Enter!")
+        
         openai_api_key = gr.Textbox(
             label='OpenAI API Key', 
             type='password', 
             placeholder='Type your OpenAI API key and press Enter!', 
-            interactive=True)
+            interactive=True
+        )
 
         chatbot = gr.Chatbot(
             label="Writing Process",
@@ -264,7 +243,8 @@ def create_demo():
             height=700,
             type="messages",
             show_label=True,
-            visible=False
+            visible=False,
+            value=[]
         )
         
         with gr.Row(equal_height=True):
@@ -275,17 +255,40 @@ def create_demo():
                 visible=False
             )
             
-            async def process_input(topic, history, openai_api_key):
+            async def process_input(topic, history, api_key):
                 nonlocal article_crew
-                # Initialize ArticleCrew with the API key if not already initialized
+                if not api_key:
+                    history.append({
+                        "role": "assistant",
+                        "content": "Please provide an OpenAI API key first.",
+                        "metadata": {"title": "❌ Error"}
+                    })
+                    yield history  # Changed from return to yield
+                    return  # Early return without value
+                
+                # Initialize or update ArticleCrew with API key
                 if article_crew is None:
-                    article_crew = ArticleCrew(api_key=openai_api_key)
+                    article_crew = ArticleCrew(api_key=api_key)
+                else:
+                    article_crew.api_key = api_key
                 
-                history.append(ChatMessage(role="user", content=f"Write an article about: {topic}"))
+                # Add user message
+                history.append({
+                    "role": "user", 
+                    "content": f"Write an article about: {topic}"
+                })
                 yield history
                 
-                async for messages in article_crew.process_article(topic):
-                    history.extend(messages)
+                try:
+                    async for messages in article_crew.process_article(topic):
+                        history.extend(messages)
+                        yield history
+                except Exception as e:
+                    history.append({
+                        "role": "assistant",
+                        "content": f"An error occurred: {str(e)}",
+                        "metadata": {"title": "❌ Error"}
+                    })
                     yield history
             
             btn = gr.Button("Write Article", variant="primary", scale=1, visible=False)
@@ -306,7 +309,7 @@ def create_demo():
 
         btn.click(
             process_input,
-            inputs=[topic, chatbot, openai_api_key],
+            inputs=[topic, chatbot, openai_api_key],  # Added openai_api_key back as input
             outputs=[chatbot]
         )
 
@@ -315,4 +318,4 @@ def create_demo():
 if __name__ == "__main__":
     demo = create_demo()
     demo.queue()
-    demo.launch()
+    demo.launch(debug=True)