Update app.py

app.py

@@ -1,7 +1,7 @@
 from crewai import Agent, Task, Crew
 import gradio as gr
 import asyncio
-from typing import List, Generator, Any, Dict, Union
+from typing import List, Dict, Any, Generator
 from langchain_openai import ChatOpenAI
 import queue
 import threading
@@ -10,22 +10,17 @@ import os
 class AgentMessageQueue:
     def __init__(self):
         self.message_queue = queue.Queue()
-        self.final_output = None
-        
+        self.last_agent = None
+
     def add_message(self, message: Dict):
+        print(f"Adding message to queue: {message}")  # Debug print
         self.message_queue.put(message)
-        
+
     def get_messages(self) -> List[Dict]:
         messages = []
         while not self.message_queue.empty():
             messages.append(self.message_queue.get())
         return messages
-    
-    def set_final_output(self, output: str):
-        self.final_output = output
-    
-    def get_final_output(self) -> str:
-        return self.final_output
 
 class ArticleCrew:
     def __init__(self, api_key: str = None):
@@ -34,264 +29,283 @@ class ArticleCrew:
         self.planner = None
         self.writer = None
         self.editor = None
-    
+        self.current_agent = None
+        self.final_article = None
+
     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(
-            temperature=0.7,
-            model="gpt-4"
-        )
-        
+        llm = ChatOpenAI(temperature=0.7, model="gpt-4")
+
         self.planner = Agent(
             role="Content Planner",
             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.",
+            backstory="Expert content planner with focus on creating engaging outlines",
             allow_delegation=False,
             verbose=True,
             llm=llm
         )
-        
+
         self.writer = Agent(
             role="Content Writer",
-            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.",
+            goal=f"Write insightful and factually accurate piece about {topic}",
+            backstory="Expert content writer with focus on engaging articles",
             allow_delegation=False,
             verbose=True,
             llm=llm
         )
-        
+
         self.editor = Agent(
             role="Editor",
-            goal="Edit a given blog post to align with the writing style",
-            backstory="You are an editor who receives a blog post from the Content Writer.",
+            goal="Polish and refine the article",
+            backstory="Expert editor with eye for detail and clarity",
             allow_delegation=False,
             verbose=True,
             llm=llm
         )
 
-    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"
-                f"2. Identify the target audience, considering their interests and pain points.\n"
-                f"3. Develop a detailed content outline including introduction, key points, and call to action.\n"
-                f"4. Include SEO keywords and relevant data or sources."
-            ),
-            expected_output="A comprehensive content plan document with an outline, audience analysis, SEO keywords, and resources.",
+    def create_tasks(self, topic: str) -> List[Task]:
+        planner_task = Task(
+            description=f"""Create a detailed content plan for an article about {topic} by:
+1. Prioritizing the latest trends, key players, and noteworthy news
+2. Identifying the target audience, considering their interests and pain points
+3. Developing a detailed content outline including introduction, key points, and call to action
+4. Including SEO keywords and relevant data or sources""",
+            expected_output="A comprehensive content plan with outline, keywords, and target audience analysis",
             agent=self.planner
         )
 
-        write_task = Task(
-            description=(
-                "1. Use the content plan to craft a compelling blog post.\n"
-                "2. Incorporate SEO keywords naturally.\n"
-                "3. Sections/Subtitles are properly named in an engaging manner.\n"
-                "4. Ensure proper structure with introduction, body, and conclusion.\n"
-                "5. Proofread for grammatical errors."
-            ),
-            expected_output="A well-written blog post in markdown format, ready for publication.",
+        writer_task = Task(
+            description="""Based on the provided content plan:
+1. Use the content plan to craft a compelling blog post
+2. Incorporate SEO keywords naturally
+3. Ensure sections/subtitles are properly named in an engaging manner
+4. Create proper structure with introduction, body, and conclusion
+5. Proofread for grammatical errors""",
+            expected_output="A well-written article draft following the content plan",
             agent=self.writer
         )
 
-        edit_task = Task(
-            description="Proofread the given blog post for grammatical errors and alignment with the brand's voice.",
-            expected_output="A well-written blog post in markdown format, ready for publication.",
+        editor_task = Task(
+            description="""Review the written article by:
+1. Checking for clarity and coherence
+2. Correcting any grammatical errors and typos
+3. Ensuring consistent tone and style
+4. Verifying proper formatting and structure""",
+            expected_output="A polished, final version of the article ready for publication",
             agent=self.editor
         )
 
-        return [plan_task, write_task, edit_task]
+        return [planner_task, writer_task, editor_task]
 
     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 add_agent_messages(agent_name: str, tasks: str, emoji: str = "🤖"):
+            # Add agent header
+            self.message_queue.add_message({
+                "role": "assistant",
+                "content": agent_name,
+                "metadata": {"title": f"{emoji} {agent_name}"}
+            })
+            
+            # Add task description
+            self.message_queue.add_message({
+                "role": "assistant",
+                "content": tasks,
+                "metadata": {"title": f"📋 Task for {agent_name}"}
+            })
 
-        def task_callback(output: Any) -> None:
-            try:
-                content = str(output)
-                if hasattr(output, 'agent'):
-                    agent_name = str(output.agent)
-                else:
-                    agent_name = "Agent"
+        def setup_next_agent(current_agent: str) -> None:
+            agent_sequence = {
+                "Content Planner": ("Content Writer", """1. Use the content plan to craft a compelling blog post
+2. Incorporate SEO keywords naturally
+3. Ensure sections/subtitles are properly named in an engaging manner
+4. Create proper structure with introduction, body, and conclusion
+5. Proofread for grammatical errors"""),
                 
+                "Content Writer": ("Editor", """1. Review the article for clarity and coherence
+2. Check for grammatical errors and typos
+3. Ensure consistent tone and style
+4. Verify proper formatting and structure""")
+            }
+
+            if current_agent in agent_sequence:
+                next_agent, tasks = agent_sequence[current_agent]
+                self.current_agent = next_agent
+                add_agent_messages(next_agent, tasks)
+
+
+        def task_callback(task_output) -> None:
+            print(f"Task callback received: {task_output}")  # Debug print
+            
+            # Extract content from raw output
+            raw_output = task_output.raw
+            if "## Final Answer:" in raw_output:
+                content = raw_output.split("## Final Answer:")[1].strip()
+            else:
+                content = raw_output.strip()
+            
+            # Handle the output based on current agent
+            if self.current_agent == "Editor":
+                # Not going to show editor's output with metadata
+                # Instead, I am showing assistant message as the final article
                 self.message_queue.add_message({
                     "role": "assistant",
-                    "content": content.strip(),
-                    "metadata": {"title": f"✅ Task completed by {agent_name}"}
+                    "content": "Final article is ready!",
+                    "metadata": {"title": "📝 Final Article"}
                 })
                 
-                # 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"}
-                        })
+                # Convert common markdown patterns to Gradio-compatible markdown
+                formatted_content = content
+                # Ensure proper spacing for headers
+                formatted_content = formatted_content.replace("\n#", "\n\n#")
+                # Ensure proper spacing for lists
+                formatted_content = formatted_content.replace("\n-", "\n\n-")
+                formatted_content = formatted_content.replace("\n*", "\n\n*")
+                formatted_content = formatted_content.replace("\n1.", "\n\n1.")
+                # Ensure proper spacing for paragraphs
+                formatted_content = formatted_content.replace("\n\n\n", "\n\n")
                 
-            except Exception as e:
-                print(f"Error in task_callback: {str(e)}")
+                # Add the final article content without metadata
+                self.message_queue.add_message({
+                    "role": "assistant",
+                    "content": formatted_content
+                })
+            else:
+                # For other agents, show their output with metadata
+                self.message_queue.add_message({
+                    "role": "assistant",
+                    "content": content,
+                    "metadata": {"title": f"✨ Output from {self.current_agent}"}
+                })
+                # Setup next agent
+                setup_next_agent(self.current_agent)
 
-        self.initialize_agents(topic)
-        
-        crew = Crew(
-            agents=[self.planner, self.writer, self.editor],
-            tasks=self.create_tasks(topic),
-            verbose=True,
-            step_callback=step_callback,
-            task_callback=task_callback
-        )
-        
-        # 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()
-        
-        # 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 step_callback(output: Any) -> None:
+            print(f"Step callback received: {output}")  # Debug print
+            # We'll only use step_callback for logging purposes now
+            pass
+
+        try:
+            self.initialize_agents(topic)
+            self.current_agent = "Content Planner"
+
+            # Start process
+            yield [{
+                "role": "assistant",
+                "content": "Starting work on your article...",
+                "metadata": {"title": "🚀 Process Started"}
+            }]
+
+            # Initialize first agent
+            add_agent_messages("Content Planner", 
+                """1. Prioritize the latest trends, key players, and noteworthy news
+2. Identify the target audience, considering their interests and pain points
+3. Develop a detailed content outline including introduction, key points, and call to action
+4. Include SEO keywords and relevant data or sources""")
+
+            crew = Crew(
+                agents=[self.planner, self.writer, self.editor],
+                tasks=self.create_tasks(topic),
+                verbose=True,
+                step_callback=step_callback,
+                task_callback=task_callback
+            )
 
+            def run_crew():
+                try:
+                    crew.kickoff()
+                except Exception as e:
+                    print(f"Error in crew execution: {str(e)}")  # Debug print
+                    self.message_queue.add_message({
+                        "role": "assistant",
+                        "content": f"An error occurred: {str(e)}",
+                        "metadata": {"title": "❌ Error"}
+                    })
+
+            thread = threading.Thread(target=run_crew)
+            thread.start()
+
+            while thread.is_alive() or not self.message_queue.message_queue.empty():
+                messages = self.message_queue.get_messages()
+                if messages:
+                    print(f"Yielding messages: {messages}")  # Debug print
+                    yield messages
+                await asyncio.sleep(0.1)
+
+        except Exception as e:
+            print(f"Error in process_article: {str(e)}")  # Debug print
+            yield [{
+                "role": "assistant",
+                "content": f"An error occurred: {str(e)}",
+                "metadata": {"title": "❌ Error"}
+            }]
+
+
+# Gradio code
 def create_demo():
     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!', 
+            label='OpenAI API Key',
+            type='password',
+            placeholder='Enter your OpenAI API key...',
             interactive=True
         )
 
         chatbot = gr.Chatbot(
             label="Writing Process",
-            avatar_images=(None, "https://avatars.githubusercontent.com/u/170677839?v=4"),
             height=700,
             type="messages",
             show_label=True,
             visible=False,
-            value=[]
+            bubble_full_width=False,  # Allow messages to wrap naturally
+            render_markdown=True  # Enable markdown rendering
         )
-        
+
         with gr.Row(equal_height=True):
             topic = gr.Textbox(
                 label="Article Topic",
-                placeholder="Enter the topic you want an article about...",
+                placeholder="Enter topic...",
                 scale=4,
                 visible=False
             )
-            
-            async def process_input(topic, history, api_key):
-                nonlocal article_crew
-                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=api_key)
-                else:
-                    article_crew.api_key = api_key
-                
-                # Add user message
+            btn = gr.Button("Write Article", variant="primary", scale=1, visible=False)
+
+        async def process_input(topic, history, api_key):
+            nonlocal article_crew
+            if not api_key:
+                history = history or []
                 history.append({
-                    "role": "user", 
-                    "content": f"Write an article about: {topic}"
+                    "role": "assistant",
+                    "content": "Please provide an OpenAI API key.",
+                    "metadata": {"title": "❌ Error"}
                 })
                 yield history
-                
-                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"}
-                    })
+                return
+
+            if article_crew is None:
+                article_crew = ArticleCrew(api_key=api_key)
+
+            history = history or []
+            history.append({"role": "user", "content": f"Write an article about: {topic}"})
+            yield history
+
+            try:
+                async for messages in article_crew.process_article(topic):
+                    history.extend(messages)
                     yield history
-            
-            btn = gr.Button("Write Article", variant="primary", scale=1, visible=False)
+            except Exception as e:
+                history.append({
+                    "role": "assistant",
+                    "content": f"An error occurred: {str(e)}",
+                    "metadata": {"title": "❌ Error"}
+                })
+                yield history
 
         def show_interface():
             return {
@@ -301,17 +315,8 @@ def create_demo():
                 btn: gr.Button(visible=True)
             }
 
-        openai_api_key.submit(
-            show_interface,
-            None,
-            [openai_api_key, chatbot, topic, btn]
-        )
-
-        btn.click(
-            process_input,
-            inputs=[topic, chatbot, openai_api_key],  # Added openai_api_key back as input
-            outputs=[chatbot]
-        )
+        openai_api_key.submit(show_interface, None, [openai_api_key, chatbot, topic, btn])
+        btn.click(process_input, [topic, chatbot, openai_api_key], [chatbot])
 
     return demo