app.py

from crewai import Agent, Task, Crew
import gradio as gr
import asyncio
from typing import List, Dict, Any, Generator
from langchain_openai import ChatOpenAI
import queue
import threading
import os

class AgentMessageQueue:
    def __init__(self):
        self.message_queue = queue.Queue()

    def add_message(self, message: Dict):
        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

class PressReleaseCrew:
    def __init__(self, api_key: str = None):
        self.api_key = api_key
        self.message_queue = AgentMessageQueue()
        self.researcher = None
        self.writer = None
        self.editor = None
        self.current_agent = 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")

        self.researcher = Agent(
            role="News Researcher",
            goal=f"Gather critical details and facts for a press release about {topic}",
            backstory="An experienced journalist who specializes in gathering news and structuring press releases.",
            allow_delegation=False,
            verbose=True,
            llm=llm
        )

        self.writer = Agent(
            role="Press Release Writer",
            goal=f"Draft a compelling and structured press release about {topic}",
            backstory="A seasoned news writer with expertise in crafting clear and engaging press releases.",
            allow_delegation=False,
            verbose=True,
            llm=llm
        )

        self.editor = Agent(
            role="News Editor",
            goal="Refine and finalize the press release for accuracy and professionalism.",
            backstory="A skilled editor with an eye for clarity, conciseness, and journalistic integrity.",
            allow_delegation=False,
            verbose=True,
            llm=llm
        )

    def create_tasks(self, topic: str) -> List[Task]:
        researcher_task = Task(
            description=f"""As a news researcher, compile essential details for a press release on {topic} by:
1. Identifying key facts, statistics, and industry trends
2. Structuring the information into a brief, clear outline
3. Suggesting a compelling press release headline""",
            expected_output="A structured summary with key facts, statistics, and a proposed headline.",
            agent=self.researcher
        )

        writer_task = Task(
            description="""Using the research provided:
1. Write a clear and engaging press release following a journalistic structure
2. Ensure it includes a compelling headline, subheading, lead paragraph, supporting details, and a boilerplate
3. Maintain a professional and neutral tone""",
            expected_output="A structured press release draft ready for editing.",
            agent=self.writer
        )

        editor_task = Task(
            description="""Review and refine the press release by:
1. Checking for clarity, conciseness, and accuracy
2. Ensuring proper journalistic tone and structure
3. Correcting any grammatical or formatting issues""",
            expected_output="A polished, publication-ready press release.",
            agent=self.editor
        )

        return [researcher_task, writer_task, editor_task]

    async def process_press_release(self, topic: str) -> Generator[List[Dict], None, None]:
        def add_agent_messages(agent_name: str, tasks: str, emoji: str = "📰"):
            self.message_queue.add_message({
                "role": "assistant",
                "content": agent_name,
                "metadata": {"title": f"{emoji} {agent_name}"}
            })
            self.message_queue.add_message({
                "role": "assistant",
                "content": tasks,
                "metadata": {"title": f"📄 Task for {agent_name}"}
            })

        def setup_next_agent(current_agent: str) -> None:
            agent_sequence = {
                "News Researcher": ("Press Release Writer", """Write a structured and engaging press release including:
1. A compelling headline and subheading
2. A strong lead paragraph
3. Supporting details and key statistics
4. A conclusion with a call to action or company statement"""),
                
                "Press Release Writer": ("News Editor", """Review and refine the press release for:
1. Clarity and conciseness
2. Proper journalistic tone and structure
3. Grammatical accuracy and formatting""")
            }

            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:
            raw_output = task_output.raw.strip()
            
            if self.current_agent == "News Editor":
                self.message_queue.add_message({
                    "role": "assistant",
                    "content": "Final press release is ready!",
                    "metadata": {"title": "📝 Final Press Release"}
                })
                self.message_queue.add_message({
                    "role": "assistant",
                    "content": raw_output
                })
            else:
                self.message_queue.add_message({
                    "role": "assistant",
                    "content": raw_output,
                    "metadata": {"title": f"📰 Output from {self.current_agent}"}
                })
                setup_next_agent(self.current_agent)

        def step_callback(output: Any) -> None:
            pass

        try:
            self.initialize_agents(topic)
            self.current_agent = "News Researcher"

            yield [{
                "role": "assistant",
                "content": "Starting press release preparation...",
                "metadata": {"title": "🚀 Press Release Process Started"}
            }]

            add_agent_messages("News Researcher", 
                """Gather essential details for a press release:
1. Identify key facts, statistics, and industry trends
2. Structure information into a brief outline
3. Suggest a compelling headline""")

            crew = Crew(
                agents=[self.researcher, 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:
                    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:
                    yield messages
                await asyncio.sleep(0.1)

        except Exception as e:
            yield [{
                "role": "assistant",
                "content": f"An error occurred: {str(e)}",
                "metadata": {"title": "❌ Error"}
            }]

def create_demo():
    press_release_crew = None

    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 📰 AI Space Launch Newsroom - Press Release Generator")
        
        openai_api_key = gr.Textbox(
            label='OpenAI API Key',
            type='password',
            placeholder='Enter your OpenAI API key...'
        )

        chatbot = gr.Chatbot(
            label="Press Release Process",
            height=700,
            type="messages",
            show_label=True
        )

        topic = gr.Textbox(
            label="Press Release Topic",
            placeholder="Enter topic..."
        )
        btn = gr.Button("Generate Press Release", variant="primary")

        async def process_input(topic, history, api_key):
            nonlocal press_release_crew
            if not api_key:
                yield history + [{"role": "assistant", "content": "Please provide an OpenAI API key."}]
                return

            if press_release_crew is None:
                press_release_crew = PressReleaseCrew(api_key=api_key)

            async for messages in press_release_crew.process_press_release(topic):
                yield messages

        btn.click(process_input, [topic, chatbot, openai_api_key], [chatbot])

    return demo

if __name__ == "__main__":
    demo = create_demo()
    demo.queue()
    demo.launch(debug=True)