Multi-Agent Systems

"""Message passing infrastructure for multi-agent systems."""
from dataclasses import dataclass, field
from typing import Optional
from enum import Enum
import uuid
from datetime import datetime

class MessageType(Enum):
    TASK = "task"           # Assignment from manager
    RESULT = "result"       # Completed work
    FEEDBACK = "feedback"   # Revision request
    QUERY = "query"         # Question to another agent
    STATUS = "status"       # Progress update

@dataclass
class AgentMessage:
    """Structured message between agents."""
    id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
    sender: str = ""
    receiver: str = ""
    type: MessageType = MessageType.TASK
    content: str = ""
    metadata: dict = field(default_factory=dict)
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
    parent_id: Optional[str] = None  # For threading conversations

class MessageBus:
    """Central message router for agent communication."""
    
    def __init__(self):
        self.queues: dict[str, list[AgentMessage]] = {}
        self.history: list[AgentMessage] = []
    
    def register_agent(self, agent_name: str):
        self.queues[agent_name] = []
    
    def send(self, message: AgentMessage):
        """Route message to receiver's queue."""
        self.history.append(message)
        if message.receiver in self.queues:
            self.queues[message.receiver].append(message)
        else:
            raise ValueError(f"Unknown agent: {message.receiver}")
    
    def receive(self, agent_name: str) -> list[AgentMessage]:
        """Get and clear pending messages for an agent."""
        messages = self.queues.get(agent_name, [])
        self.queues[agent_name] = []
        return messages
    
    def get_conversation(self, agent_a: str, agent_b: str) -> list[AgentMessage]:
        """Get all messages between two agents."""
        return [m for m in self.history 
                if (m.sender == agent_a and m.receiver == agent_b) or
                   (m.sender == agent_b and m.receiver == agent_a)]

"""Multi-Agent Content Pipeline — Researcher → Writer → Editor."""
import openai

client = openai.OpenAI()

class Agent:
    """Base agent with role, tools, and message handling."""
    
    def __init__(self, name: str, system_prompt: str, model: str = "gpt-4o"):
        self.name = name
        self.system_prompt = system_prompt
        self.model = model
        self.conversation: list[dict] = []
    
    def process(self, task: str, context: str = "") -> str:
        """Process a task and return the result."""
        messages = [
            {"role": "system", "content": self.system_prompt},
        ]
        if context:
            messages.append({"role": "user", "content": f"Context from previous agents:\n{context}"})
        messages.append({"role": "user", "content": task})
        
        response = client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=0.7,
        )
        
        result = response.choices[0].message.content
        self.conversation.append({"task": task, "result": result})
        return result

# ─── Define Specialized Agents ────────────────────────────────

researcher = Agent(
    name="Researcher",
    system_prompt="""You are a research specialist. Your job is to:
1. Identify key topics and subtopics to investigate
2. Find relevant facts, statistics, and examples
3. Organize findings into structured notes
4. Flag areas that need more investigation

Output FORMAT: structured research notes with headers, bullet points, and source attribution.
Be thorough but concise. Focus on facts, not opinions."""
)

writer = Agent(
    name="Writer",
    system_prompt="""You are a professional technical writer. Your job is to:
1. Take research notes and transform them into clear, engaging prose
2. Structure content with logical flow (intro → body → conclusion)
3. Use concrete examples and analogies
4. Maintain a professional but accessible tone

Output FORMAT: well-structured article with headers, paragraphs, and code examples where relevant.
Never fabricate facts — only use what's in the research notes."""
)

editor = Agent(
    name="Editor",
    system_prompt="""You are a senior editor. Your job is to:
1. Check factual accuracy against the research notes
2. Improve clarity, conciseness, and flow
3. Fix grammatical issues and awkward phrasing
4. Ensure consistent tone and style
5. Rate the final quality (1-10) and list remaining issues

Output FORMAT: 
- EDITED_CONTENT: [the improved text]
- QUALITY_SCORE: [1-10]
- ISSUES: [any remaining problems]
- VERDICT: PUBLISH or NEEDS_REVISION"""
)

# ─── Pipeline Orchestrator ────────────────────────────────────

class ContentPipeline:
    """Orchestrates the Researcher → Writer → Editor flow."""
    
    def __init__(self, max_revisions: int = 2):
        self.max_revisions = max_revisions
        self.bus = MessageBus()
        self.bus.register_agent("Researcher")
        self.bus.register_agent("Writer")
        self.bus.register_agent("Editor")
    
    def run(self, topic: str) -> dict:
        """Run the full pipeline."""
        print(f"📋 Topic: {topic}\n")
        
        # Step 1: Research
        print("🔍 Researcher working...")
        research_notes = researcher.process(
            f"Research the following topic thoroughly: {topic}"
        )
        print(f"   ✓ Research complete ({len(research_notes)} chars)\n")
        
        # Step 2: Writing (with possible revisions)
        draft = None
        for revision in range(self.max_revisions + 1):
            if revision == 0:
                print("✍️  Writer working on first draft...")
                draft = writer.process(
                    f"Write a comprehensive article about: {topic}",
                    context=research_notes
                )
            else:
                print(f"✍️  Writer working on revision {revision}...")
                draft = writer.process(
                    f"Revise the article based on this editorial feedback:\n{feedback}",
                    context=f"Original research:\n{research_notes}\n\nCurrent draft:\n{draft}"
                )
            
            print(f"   ✓ Draft complete ({len(draft)} chars)\n")
            
            # Step 3: Editing
            print("📝 Editor reviewing...")
            edit_result = editor.process(
                "Review this article for quality, accuracy, and clarity.",
                context=f"Research notes:\n{research_notes}\n\nDraft:\n{draft}"
            )
            
            # Parse editor verdict
            if "PUBLISH" in edit_result or revision == self.max_revisions:
                print("   ✓ Editor approved!\n")
                break
            else:
                feedback = edit_result
                print(f"   ↩ Editor requested revision\n")
        
        return {
            "topic": topic,
            "research": research_notes,
            "final_draft": draft,
            "editor_notes": edit_result,
            "revisions": revision,
        }

# ─── Run the Pipeline ────────────────────────────────────────

pipeline = ContentPipeline(max_revisions=2)
result = pipeline.run("The impact of transformer architecture on modern NLP")
print(f"\n{'='*60}")
print(f"Final article ({len(result['final_draft'])} chars, {result['revisions']} revisions)")
print(result['final_draft'][:500] + "...")

"""Framework comparison — same task in CrewAI vs custom."""

# ─── CrewAI Version (for reference) ──────────────────────────
# pip install crewai

from crewai import Agent as CrewAgent, Task, Crew

crew_researcher = CrewAgent(
    role="Research Analyst",
    goal="Find comprehensive information about {topic}",
    backstory="You are an expert researcher with access to the internet.",
    tools=[search_tool, read_tool],
)

crew_writer = CrewAgent(
    role="Technical Writer",
    goal="Write clear, engaging articles from research notes",
    backstory="You are a seasoned tech writer for a top publication.",
)

research_task = Task(
    description="Research {topic} thoroughly",
    agent=crew_researcher,
    expected_output="Structured research notes",
)

writing_task = Task(
    description="Write a comprehensive article from the research",
    agent=crew_writer,
    expected_output="Polished article",
    context=[research_task],  # Depends on research
)

crew = Crew(agents=[crew_researcher, crew_writer], tasks=[research_task, writing_task])
result = crew.kickoff(inputs={"topic": "quantum computing advances 2024"})

"""Consensus: Multiple agents vote on an answer."""

def consensus_answer(question: str, num_agents: int = 3) -> str:
    """Get multiple perspectives and find consensus."""
    perspectives = [
        "You are a conservative analyst who values caution and proven approaches.",
        "You are an innovative thinker who values novel solutions and experimentation.",
        "You are a pragmatic engineer who values simplicity and practicality.",
    ]
    
    answers = []
    for i in range(num_agents):
        agent = Agent(f"Agent_{i}", perspectives[i % len(perspectives)])
        answer = agent.process(question)
        answers.append(answer)
    
    # Synthesizer combines the answers
    synthesizer = Agent("Synthesizer", 
        "You synthesize multiple expert opinions into a balanced conclusion. "
        "Note where experts agree and where they disagree.")
    
    combined = "\n\n".join(f"Expert {i+1}: {a}" for i, a in enumerate(answers))
    return synthesizer.process(
        f"Question: {question}\n\nSynthesize these expert opinions:\n{combined}"
    )

"""Manager that dynamically creates specialized agents."""

class DynamicManager:
    """Creates agents on-the-fly based on task requirements."""
    
    def __init__(self):
        self.planner = Agent("Planner",
            "You are a project planner. Given a complex task, break it into subtasks "
            "and specify what kind of specialist is needed for each. Output JSON:\n"
            '[{"task": "...", "specialist": "...", "skills": "..."}]')
    
    def execute(self, goal: str) -> str:
        # Step 1: Plan and identify needed specialists
        plan_json = self.planner.process(f"Break this into subtasks: {goal}")
        subtasks = json.loads(plan_json)
        
        # Step 2: Dynamically create agents for each subtask
        results = []
        for subtask in subtasks:
            specialist = Agent(
                name=subtask["specialist"],
                system_prompt=f"You are a {subtask['specialist']} with expertise in: "
                             f"{subtask['skills']}. Complete the assigned task thoroughly."
            )
            result = specialist.process(subtask["task"])
            results.append({"task": subtask["task"], "result": result})
        
        # Step 3: Synthesize results
        synthesizer = Agent("Synthesizer", "Combine these results into a coherent final output.")
        context = "\n\n".join(f"Task: {r['task']}\nResult: {r['result']}" for r in results)
        return synthesizer.process(f"Goal: {goal}\n\nResults:\n{context}")