A More Controllable Way To Manage Autogen Agent Conversations

Autogen Orchestrator

Taken directly from https://www.youtube.com/watch?v=4o8tymMQ5GM

autogen_orchestrator.py

import autogen


class Orchestrator:
  def __init__(self, name: str, agents: list[autogen.ConversableAgent]):
    self.name = name
    self.agents = agents
    self.messages = []
    self.complete_keyword = "APPROVED"
    self.error_keyword = "ERROR"

    if len(self.agents) < 2:
      raise Exception("Orchestrator needs at least 2 agents")
    
  @property
  def total_agents(self):
    return len(self.agents)

  @property
  def latest_message_is_dict(self):
    return isinstance(self.latest_message, dict)
  
  @property
  def latest_message_is_string(self):
    return isinstance(self.latest_message, str)
  
  @property
  def latest_message_is_func_call(self):
    return self.latest_message_is_dict and self.latest_message.get("function_call")

  @property
  def latest_message_is_content(self):
    return self.latest_message_is_dict and self.latest_message.get("content")

  @property
  def latest_message(self) -> dict | str | None:
    if not self.messages:
      return None

    return self.messages[-1]

  def add_message(self, message):
    self.messages.append(message)
  
  def has_functions(self, agent: autogen.ConversableAgent):
    return agent._function_map is not None
  
  def basic_chat(self, agent_a: autogen.ConversableAgent, agent_b: autogen.ConversableAgent, message: str):
    print(f"basic_chat: {agent_a.name} -> {agent_b.name}")

    agent_a.send(message, agent_b)

    reply = agent_b.generate_reply

    self.add_message(message)

  def memory_chat(self, agent_a: autogen.ConversableAgent, agent_b: autogen.ConversableAgent, message: str):
    print(f"basic_chat: {agent_a.name} -> {agent_b.name}")

    agent_a.send(message, agent_b)

    reply = agent_b.generate_reply

    agent_b.send(reply, agent_b)

    self.add_message(message)

    print(f"basic_chat: replied with: {reply}")

  def function_chat(self, agent_a: autogen.ConversableAgent, agent_b: autogen.ConversableAgent, message: str):
    print(f"function_chat: {agent_a.name} -> {agent_b.name}")

    self.basic_chat(agent_a, agent_a, message)

    assert self.latest_message_is_content

    self.basic_chat(agent_a, agent_b, self.latest_message)

    self.add_message(message)

    print(f"function_chat: replied with: {reply}")

  def sequential_conversation(self, prompt: str) -> tuple[bool, list[str]]:
    """
    Runs a sequential conversation between agents

    For example
      "Agent A" -> "Agent B" -> "Agent C" -> "Agent D" -> "Agent E"
    """

    print(f"\n\n-------- {self.name} Orchestrator Starting --------\n\n")

    self.add_message(prompt)

    for idx, agent in enumerate(self.agents):
      agent_a = self.agents[idx]
      agent_b = self.agents[idx + 1]

      print(f"\n\n-------- Running iteration {idx} with (agent_a: {agent_a.name}, agent_b: {agent_b.name}) --------\n\n")

      # agent_a -> chat -> agent_b
      if self.latest_message_is_string:
        self.basic_chat(agent_a, agent_b, self.latest_message)

      # agent_a -> func() -> agent_b
      if self.latest_message_is_func_call is None and self.has_functions(agent_a):
        self.function_chat(agent_a, agent_b, self.latest_message)


      if idx == self.total_agents - 2:
        print("-------- Orchestrator Complete --------\n\n")

        was_successful = self.complete_keyword in self.latest_message

        if was_successful:
          print("✅ Orchestrator was successful")
        else:
          print("❌ Orchestrator failed")

        return was_successful, self.messages

  def broad_conversation(self, prompt: str) -> tuple[bool, list[str]]:
    """
    Broadcast a message from agent_a to alll agents

    For example
      "Agent A" -> "Agent B"
      "Agent A" -> "Agent C"
      "Agent A" -> "Agent D"
      "Agent A" -> "Agent E"
    """

    print(f"\n\n-------- {self.name} Orchestrator Starting --------\n\n")

    self.add_message(prompt)

    broadcast_agent = self.agents[0]

    for idx, agent_iterate in enumerate(self.agents):

      print(f"\n\n-------- Running iteration {idx} with (agent_broadcast: {broadcast_agent.name}, agent_iterate: {agent_iterate.name}) --------\n\n")

      # agent_a -> chat -> agent_b
      if self.latest_message_is_string:
        self.memory_chat(broadcast_agent, agent_iterate, prompt)

      # agent_a -> func() -> agent_b
      if self.latest_message_is_func_call is None and self.has_functions(agent_iterate):
        self.function_chat(agent_iterate, agent_iterate, self.latest_message)


      if idx == self.total_agents - 2:
        print("-------- Orchestrator Complete --------\n\n")

        was_successful = self.complete_keyword in self.latest_message

        if was_successful:
          print("✅ Orchestrator was successful")
        else:
          print("❌ Orchestrator failed")

        return was_successful, self.messages

main.py

# TODO

# This file will be updated to show how the orchestrator can be used

WIP refactor.py

import autogen

APPROVED = "APPROVED"
ERROR = "ERROR"

class GroupChatOrchestrator:
  def __init__(self, agents: list[autogen.ConversableAgent]):
    if len(agents) < 2:
      raise ValueError("Orchestrator needs at least 2 agents")

    self.agents = agents
    self.messages: list[str | dict] = []

  @property
  def latest_message(self) -> dict | str | None:
    return self.messages[-1] if self.messages else None

  def chat(self, agent_a: autogen.ConversableAgent, agent_b: autogen.ConversableAgent, message: str, with_memory: bool = False) -> str:
    """
    Executes a chat between two agents and returns the reply.
    If with_memory is True, the reply is sent back to agent_b.
    """
    try:
      agent_a.send(message, agent_b)

      reply = agent_b.generate_reply

      self.messages.append(message)

      if with_memory:
        agent_b.send(reply, agent_b)

      return reply
    except Exception as e:
      print(f"Error executing chat: {e}")
      return ERROR

  def function_chat(self, agent_a: autogen.ConversableAgent, agent_b: autogen.ConversableAgent, message: str) -> str:
    """
    Executes a function chat between two agents.
    """
    reply = self.chat(agent_a, agent_a, message)

    if isinstance(self.latest_message, dict) and self.latest_message.get("content"):
      reply = self.chat(agent_a, agent_b, self.latest_message)

    return reply

def sequential_conversation(self, prompt: str) -> tuple[bool, list[str | dict]]:
  """
  Runs a sequential conversation between agents.
  """
  self.messages.append(prompt)

  for idx, agent in enumerate(self.agents[:-1]):
    agent_a = self.agents[idx]
    agent_b = self.agents[idx + 1]

    if isinstance(self.latest_message, str):
      self.chat(agent_a, agent_b, self.latest_message)
    elif isinstance(self.latest_message, dict) and self.latest_message.get("function_call") is None and agent_a._function_map:
      self.function_chat(agent_a, agent_b, self.latest_message)

  was_successful = APPROVED in self.latest_message

  return was_successful, self.messages


def broadcast_conversation(self, prompt: str) -> tuple[bool, list[str | dict]]:
  """
  Broadcasts a message from the first agent to all other agents.
  """
  self.messages.append(prompt)

  broadcast_agent = self.agents[0]

  for agent_iterate in self.agents[1:]:
    if isinstance(self.latest_message, str):
      self.chat(broadcast_agent, agent_iterate, prompt, with_memory=True)
    elif isinstance(self.latest_message, dict) and self.latest_message.get("function_call") is None and agent_iterate._function_map:
      self.function_chat(agent_iterate, agent_iterate, self.latest_message)

  was_successful = APPROVED in self.latest_message

  return was_successful, self.messages