Spaces:
Sleeping
Sleeping
| 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 | |
| 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): | |
| # Initialize agents | |
| 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}. " | |
| "You collect information that helps the audience learn something " | |
| "and make informed decisions.", | |
| allow_delegation=False, | |
| verbose=True | |
| ) | |
| 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}. " | |
| "You base your writing on the work of the Content Planner.", | |
| allow_delegation=False, | |
| verbose=True | |
| ) | |
| 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.", | |
| allow_delegation=False, | |
| verbose=True | |
| ) | |
| self.output_parser = OutputParser() | |
| def create_tasks(self, topic: str): | |
| 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.", | |
| 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.", | |
| 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.", | |
| agent=self.editor | |
| ) | |
| return [plan_task, write_task, edit_task] | |
| async def process_article(self, topic: str) -> Generator[List[ChatMessage], None, None]: | |
| crew = Crew( | |
| agents=[self.planner, self.writer, self.editor], | |
| tasks=self.create_tasks(topic), | |
| verbose=2 | |
| ) | |
| 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"} | |
| )] | |
| finally: | |
| sys.stdout = original_stdout | |
| def create_demo(): | |
| article_crew = ArticleCrew() | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| gr.Markdown("# 📝 AI Article Writing Crew") | |
| gr.Markdown("Watch as our AI crew collaborates to create your article!") | |
| chatbot = gr.Chatbot( | |
| label="Writing Process", | |
| avatar_images=(None, "🤖"), | |
| height=700, | |
| type="messages", | |
| show_label=True | |
| ) | |
| topic = gr.Textbox( | |
| label="Article Topic", | |
| placeholder="Enter the topic you want an article about...", | |
| lines=2 | |
| ) | |
| async def process_input(topic, history): | |
| history.append(ChatMessage(role="user", content=f"Write an article about: {topic}")) | |
| yield history | |
| async for messages in article_crew.process_article(topic): | |
| history.extend(messages) | |
| yield history | |
| btn = gr.Button("Write Article", variant="primary") | |
| btn.click( | |
| process_input, | |
| inputs=[topic, chatbot], | |
| outputs=[chatbot] | |
| ) | |
| return demo | |
| if __name__ == "__main__": | |
| demo = create_demo() | |
| demo.queue() | |
| demo.launch(debug=True, share=True) |