AutoGen SelectorGroupChat Not Maintaining Agent Continuity for Health-Related Chat Flow

[glarunsingh]

I'm building a simple chat application using AutoGen with Flask-SocketIO (async_mode='eventlet') to integrate real-time communication. The setup includes two agents: a SupervisorAgent (for routing and general queries) and a HealthcareAgent (for health-related queries). The goal is a seamless health-focused conversation where:

• SupervisorAgent only responds to the initial message (e.g., greeting) or explicit non-health topic shifts (e.g., "What’s the time?").
• HealthcareAgent handles all health-related messages and their follow-ups once a health topic is introduced, based on conversation history.

However, the SelectorGroupChat inconsistently routes follow-up messages. After an initial health-related message (e.g., "I have a headache"), subsequent responses like "It’s mild around my temples" or "Yes to light" are sometimes routed to SupervisorAgent instead of HealthcareAgent, disrupting the conversation flow. The SupervisorAgent shouldn’t comment during health discussions—it should silently route to HealthcareAgent.

Scenario

• Desired Flow:

User: I have a headache
Healthcare: How would you describe the headache—sharp, dull, or throbbing?
User: It’s mild around my temples
Healthcare: Thanks for the details. How long has it lasted?
User: Yes to light
Healthcare: Sensitivity to light could indicate a migraine. How long has it been?

• Actual Flow (Issue):
• SupervisorAgent interrupts health follow-ups with routing messages like "Please wait for the HealthcareAgent," even though the context is clearly health-related.
  Current Response
  Here’s the latest terminal output I’m getting:

Running updated script with improved selector logic

• Restarting with stat
• Debugger is active!
• Debugger PIN: 774-593-654
  (17019) wsgi starting up on http://127.0.0.1:5000
  ...
  Result: TaskResult(messages=[TextMessage(source='user', ..., content='i have head ache', ...), TextMessage(source='HealthcareAgent', ..., content="I'm sorry to hear that. How would you describe the headache—sharp, dull, or throbbing? And how long have you been experiencing it?", ...)], stop_reason='Maximum number of turns 1 reached.')
  user: i have head ache
  Healthcare: I'm sorry to hear that. How would you describe the headache—sharp, dull, or throbbing? And how long have you been experiencing it?
  Result: TaskResult(messages=[TextMessage(source='user', ..., content='its mild around my temples', ...), TextMessage(source='SupervisorAgent', ..., content='Please wait for the HealthcareAgent to assist you with that.', ...)], stop_reason='Maximum number of turns 1 reached.')
  user: its mild around my temples
  Supervisor: Please wait for the HealthcareAgent to assist you with that.
  Result: TaskResult(messages=[TextMessage(source='user', ..., content='ok', ...), TextMessage(source='HealthcareAgent', ..., content='Mild headaches around the temples can be due to tension or stress. Have you experienced any other symptoms, like sensitivity to light or nausea?', ...)], stop_reason='Maximum number of turns 1 reached.')
  user: ok
  Healthcare: Mild headaches around the temples can be due to tension or stress. Have you experienced any other symptoms, like sensitivity to light or nausea?
  Result: TaskResult(messages=[TextMessage(source='user', ..., content='yes to light', ...), TextMessage(source='SupervisorAgent', ..., content='Please continue discussing your symptoms with the HealthcareAgent for proper guidance.', ...)], stop_reason='Maximum number of turns 1 reached.')
  user: yes to light
  Supervisor: Please continue discussing your symptoms with the HealthcareAgent for proper guidance.

What’s Happening
Initial Health Query: "i have head ache" correctly triggers HealthcareAgent.
Follow-Up Issues: "its mild around my temples" and "yes to light" are routed to SupervisorAgent, which responds with routing messages instead of letting HealthcareAgent continue seamlessly.
Selector Failure: Despite a refined selector_prompt emphasizing {history} and agent stickiness, the selector still picks SupervisorAgent for health follow-ups, suggesting a deeper issue with how SelectorGroupChat interprets context or applies the prompt.
Possible Solutions to Triage
I’ve tried refining the selector_prompt multiple times, but the issue persists. Could the community suggest:

Prompt Optimization: Is there a more effective selector_prompt to ensure HealthcareAgent persists for health-related follow-ups without SupervisorAgent interrupting? Perhaps a stricter syntax or additional context cues?
Manual State Tracking: Should I bypass SelectorGroupChat’s built-in selector and implement a manual state tracker (e.g., a variable tracking the current topic) to force HealthcareAgent after a health topic starts? How would this integrate with team.run?
Alternative Approach: Would using team.run_stream with explicit agent selection (instead of team.run with a selector) provide better control over agent continuity? Any examples of this approach?
Debugging Selector Decisions: Is there a way to log the selector’s reasoning (e.g., why it chose SupervisorAgent) to diagnose its logic? This could help pinpoint where it’s misinterpreting {history}.
Any insights or code snippets to triage this would be greatly appreciated! I’d like to avoid SupervisorAgent responses entirely once a health topic is active, ensuring a smooth chat flow with HealthcareAgent.

Current Code

Here’s the simplified code I’m using:

import eventlet
eventlet.monkey_patch()

from flask import Flask, request
from flask_socketio import SocketIO, emit
from autogen_agentchat.agents import AssistantAgent, UserProxyAgent
from autogen_agentchat.conditions import TextMentionTermination, MaxMessageTermination
from autogen_agentchat.teams import SelectorGroupChat
from autogen_ext.models.openai import OpenAIChatCompletionClient
from dotenv import load_dotenv
import os
import asyncio

load_dotenv()

app = Flask(__name__)
socketio = SocketIO(app, async_mode='eventlet')

teams = {}
event_loops = {}

def create_team():
  model_client = OpenAIChatCompletionClient(model="gpt-4o-mini")
  supervisor_agent = AssistantAgent(
      name="SupervisorAgent",
      description="Routes to the appropriate agent based on context.",
      model_client=model_client,
      system_message="You are a routing agent. Your sole job is to direct the user to the correct agent based on their message and history. Do not provide advice, commentary, or engage in the conversation. If the topic is health-related (e.g., symptoms, doctors), route to 'HealthcareAgent' silently without responding. Respond only for the first message or if the topic shifts to a non-health general query (e.g., 'hi', 'what’s the time') with 'I’m here to assist with general questions. How can I help?'",
  )
  healthcare_agent = AssistantAgent(
      name="HealthcareAgent",
      description="Handles health-related queries.",
      model_client=model_client,
      system_message="You are a healthcare agent. Provide concise, helpful responses to health-related queries (e.g., symptoms, advice). Ask follow-up questions if needed, but keep it to one question at a time. Avoid repeating prior messages unless relevant.",
  )
  user_proxy = UserProxyAgent(name="user_proxy")
  termination = TextMentionTermination("TERMINATE") | MaxMessageTermination(max_messages=25)
  selector_prompt = """
  Select one agent to respond based on the user’s message and conversation history.

  {roles}

  Conversation history:
  {history}

  Rules:
  - If this is the first message and it’s a general greeting or non-health question (e.g., 'hi', 'what time is it'), select 'SupervisorAgent'.
  - If the message mentions health issues or symptoms (e.g., headache, pain, mild, temples, sick, light, nausea) OR if {history} includes any health-related discussion (e.g., prior mention of symptoms, health queries, or HealthcareAgent responses), select 'HealthcareAgent'.
  - Once 'HealthcareAgent' is involved in a health topic, always select 'HealthcareAgent' for subsequent messages unless the user explicitly shifts to a non-health topic (e.g., 'what’s the weather today').
  - Do not select 'SupervisorAgent' for health-related follow-ups, even if the message is vague (e.g., 'ok', 'yes')—assume it relates to the ongoing health context in {history}.
  - Return the selected agent’s name only (e.g., 'HealthcareAgent' or 'SupervisorAgent').

  Analyze the current message and {history} to decide.
  """
  team = SelectorGroupChat(
      [user_proxy, supervisor_agent, healthcare_agent],
      model_client=model_client,
      termination_condition=termination,
      selector_prompt=selector_prompt,
      max_turns=1
  )
  return team

@socketio.on('connect')
def handle_connect():
  socket_id = request.sid
  team = create_team()
  loop = asyncio.new_event_loop()
  teams[socket_id] = team
  event_loops[socket_id] = loop
  emit('response', {'message': "Hi, I'm Zoe, your health care assistant. How can I help you today?"})

@socketio.on('disconnect')
def handle_disconnect():
  socket_id = request.sid
  if socket_id in teams:
      event_loops[socket_id].close()
      del teams[socket_id]
      del event_loops[socket_id]

@socketio.on('message')
def handle_message(data):
  socket_id = request.sid
  message = data['message']
  team = teams[socket_id]
  loop = event_loops[socket_id]

  def run_team_task():
      asyncio.set_event_loop(loop)
      result = loop.run_until_complete(team.run(task=message))
      return result

  greenlet = eventlet.spawn(run_team_task)
  result = greenlet.wait()

  print(f"Result: {result}")
  latest_message = None
  for msg in result.messages:
      sender = "User" if msg.source == "user_proxy" else msg.source.replace("Agent", "")
      print(f"{sender}: {msg.content}")
      if msg.source == "HealthcareAgent":
          latest_message = msg.content
      elif msg.source == "SupervisorAgent" and not any("health" in str(m).lower() for m in result.messages[:-1]):
          latest_message = msg.content
  if latest_message:
      emit('response', {'message': latest_message})

@app.route('/')
def index():
  return '''
  <html>
  <head>
      <script src="https://cdnjs.cloudflare.com/ajax/libs/socket.io/4.0.1/socket.io.js"></script>
      <script>
          var socket = io();
          socket.on('connect', function() {});
          socket.on('response', function(data) {
              console.log('Received:', data.message);
              var chatWindow = document.getElementById('chat-window');
              chatWindow.innerHTML += '<p>' + data.message + '</p>';
          });
          function sendMessage() {
              var input = document.getElementById('message-input');
              var message = input.value.trim();
              if (message) {
                  socket.emit('message', {message: message});
                  var chatWindow = document.getElementById('chat-window');
                  chatWindow.innerHTML += '<p>You: ' + message + '</p>';
                  input.value = '';
              }
          }
      </script>
      <style>
          #chat-window { border: 1px solid #ccc; padding: 10px; height: 300px; overflow-y: auto; }
          #message-input { width: 70%; padding: 5px; }
          button { padding: 5px 10px; }
      </style>
  </head>
  <body>
      <div id="chat-window"></div>
      <input type="text" id="message-input" onkeypress="if(event.key === 'Enter') sendMessage()">
      <button onclick="sendMessage()">Send</button>
  </body>
  </html>
  '''

if __name__ == '__main__':
  print("Running updated script with improved selector logic")
  socketio.run(app, debug=True)

[glarunsingh]

Also created this - #5535