Merge pull request #1479 from Steve-Dusty/fix/repeated-agent-flow-awareness

fix: correct sequential awareness for repeated agents in flow

Author: kyegomez

swarms/structs/agent_rearrange.py

@@ -327,35 +327,41 @@ def validate_flow(self):
                     )
                 agents_in_flow.append(agent_name)
 
-        # # If the length of the agents does not equal the length of the agents in flow
-        # if len(set(agents_in_flow)) != len(agents_in_flow):
-        #     raise ValueError(
-        #         "Duplicate agent names in the flow are not allowed."
-        #     )
-
         logger.info(f"Flow: {self.flow} is valid.")
         return True
 
     def _get_sequential_awareness(
-        self, agent_name: str, tasks: List[str]
+        self,
+        agent_name: str,
+        tasks: List[str],
+        task_idx: int = None,
     ) -> str:
         """
         Determines the sequential awareness information for an agent in a sequential flow.
 
         Args:
             agent_name (str): The name of the current agent.
             tasks (List[str]): The list of tasks in the flow.
+            task_idx (int, optional): The exact position index of this agent invocation
+                in the flow. When provided, uses this directly instead of searching by name.
+                This is essential for repeated agents (e.g., "writer -> reviewer -> writer")
+                so each occurrence gets correct positional awareness.
 
         Returns:
             str: A string describing the agents ahead and behind in the sequence.
         """
-        # Find the position of the current agent in the flow
-        agent_position = None
-        for i, task in enumerate(tasks):
-            agent_names = [name.strip() for name in task.split(",")]
-            if agent_name in agent_names:
-                agent_position = i
-                break
+        # Use provided position index if available, otherwise search by name
+        if task_idx is not None:
+            agent_position = task_idx
+        else:
+            agent_position = None
+            for i, task in enumerate(tasks):
+                agent_names = [
+                    name.strip() for name in task.split(",")
+                ]
+                if agent_name in agent_names:
+                    agent_position = i
+                    break
 
         if agent_position is None:
             return ""
@@ -526,6 +532,7 @@ def _run_sequential_workflow(
         self,
         agent_name: str,
         tasks: List[str],
+        task_idx: int = None,
         img: str = None,
         *args,
         **kwargs,
@@ -541,6 +548,8 @@ def _run_sequential_workflow(
             agent_name (str): Name of the agent to run sequentially.
             tasks (List[str]): List of all tasks in the flow for awareness context.
                 Used to determine the agent's position and provide awareness info.
+            task_idx (int, optional): The position index of this agent in the flow.
+                Essential for repeated agents so each occurrence gets correct awareness.
             img (str, optional): Image input for agents that support it.
                 Defaults to None.
             *args: Additional positional arguments passed to agent execution.
@@ -560,7 +569,7 @@ def _run_sequential_workflow(
 
         # Add sequential awareness information for the agent
         awareness_info = self._get_sequential_awareness(
-            agent_name, tasks
+            agent_name, tasks, task_idx=task_idx
         )
         if awareness_info:
             self.conversation.add("system", awareness_info)
@@ -676,11 +685,20 @@ def _run(
                     result = self._run_sequential_workflow(
                         agent_name=agent_name,
                         tasks=tasks,
+                        task_idx=task_idx,
                         img=img,
                         *args,
                         **kwargs,
                     )
-                    response_dict[agent_name] = result
+
+                    # Use indexed key to preserve all outputs
+                    # from repeated agents (e.g., "Writer_0", "Writer_2")
+                    if agent_name in response_dict:
+                        response_dict[
+                            f"{agent_name}_{task_idx}"
+                        ] = result
+                    else:
+                        response_dict[agent_name] = result
 
             loop_count += 1
 

tests/structs/test_agent_rearrange.py

@@ -597,6 +597,216 @@ def test_successful_callable_returns_result():
     assert result is not None
 
 
+# ============================================================================
+# Repeated Agent Flow Tests
+# ============================================================================
+
+
+def create_repeated_flow_agents():
+    """Create agents for repeated flow testing."""
+    return [
+        Agent(
+            agent_name="Writer",
+            agent_description="Expert in writing content",
+            system_prompt="You are a writer. Write one concise sentence.",
+            model_name="gpt-4o-mini",
+            max_loops=1,
+            verbose=False,
+        ),
+        Agent(
+            agent_name="Reviewer",
+            agent_description="Expert in reviewing content",
+            system_prompt="You are a reviewer. Give one critique.",
+            model_name="gpt-4o-mini",
+            max_loops=1,
+            verbose=False,
+        ),
+    ]
+
+
+def test_repeated_agent_flow_valid():
+    """Test that flows with repeated agents pass validation."""
+    agents = create_repeated_flow_agents()
+
+    agent_rearrange = AgentRearrange(
+        agents=agents,
+        flow="Writer -> Reviewer -> Writer",
+    )
+
+    assert agent_rearrange.validate_flow() is True
+    print("✓ test_repeated_agent_flow_valid passed")
+
+
+def test_repeated_agent_awareness_position_0():
+    """Test that the first occurrence of a repeated agent gets correct awareness."""
+    agents = create_repeated_flow_agents()
+
+    agent_rearrange = AgentRearrange(
+        agents=agents,
+        flow="Writer -> Reviewer -> Writer",
+    )
+
+    tasks = agent_rearrange.flow.split("->")
+    awareness = agent_rearrange._get_sequential_awareness(
+        "Writer", tasks, task_idx=0
+    )
+
+    # Writer at position 0: no agent ahead, Reviewer behind
+    assert "Agent behind" in awareness
+    assert "Reviewer" in awareness
+    assert "Agent ahead" not in awareness
+    print("✓ test_repeated_agent_awareness_position_0 passed")
+
+
+def test_repeated_agent_awareness_position_2():
+    """Test that the second occurrence of a repeated agent gets correct awareness."""
+    agents = create_repeated_flow_agents()
+
+    agent_rearrange = AgentRearrange(
+        agents=agents,
+        flow="Writer -> Reviewer -> Writer",
+    )
+
+    tasks = agent_rearrange.flow.split("->")
+    awareness = agent_rearrange._get_sequential_awareness(
+        "Writer", tasks, task_idx=2
+    )
+
+    # Writer at position 2: Reviewer ahead, no agent behind
+    assert "Agent ahead" in awareness
+    assert "Reviewer" in awareness
+    assert "Agent behind" not in awareness
+    print("✓ test_repeated_agent_awareness_position_2 passed")
+
+
+def test_repeated_agent_awareness_differs_per_position():
+    """Test that each occurrence of a repeated agent gets different awareness."""
+    agents = create_repeated_flow_agents()
+
+    agent_rearrange = AgentRearrange(
+        agents=agents,
+        flow="Writer -> Reviewer -> Writer",
+    )
+
+    tasks = agent_rearrange.flow.split("->")
+    awareness_0 = agent_rearrange._get_sequential_awareness(
+        "Writer", tasks, task_idx=0
+    )
+    awareness_2 = agent_rearrange._get_sequential_awareness(
+        "Writer", tasks, task_idx=2
+    )
+
+    assert awareness_0 != awareness_2
+    print("✓ test_repeated_agent_awareness_differs_per_position passed")
+
+
+def test_repeated_agent_awareness_fallback_without_idx():
+    """Test that awareness still works when task_idx is not provided (backward compat)."""
+    agents = create_repeated_flow_agents()
+
+    agent_rearrange = AgentRearrange(
+        agents=agents,
+        flow="Writer -> Reviewer -> Writer",
+    )
+
+    tasks = agent_rearrange.flow.split("->")
+    # Without task_idx, falls back to finding first occurrence
+    awareness = agent_rearrange._get_sequential_awareness(
+        "Writer", tasks
+    )
+
+    assert awareness is not None
+    assert isinstance(awareness, str)
+    assert "Sequential awareness" in awareness
+    print("✓ test_repeated_agent_awareness_fallback_without_idx passed")
+
+
+def test_repeated_agent_three_occurrences():
+    """Test awareness correctness with three occurrences of the same agent."""
+    agents = create_repeated_flow_agents()
+
+    agent_rearrange = AgentRearrange(
+        agents=agents,
+        flow="Writer -> Reviewer -> Writer -> Reviewer -> Writer",
+    )
+
+    tasks = agent_rearrange.flow.split("->")
+
+    # Writer at pos 0: no ahead, Reviewer behind
+    a0 = agent_rearrange._get_sequential_awareness("Writer", tasks, task_idx=0)
+    assert "Agent behind" in a0
+    assert "Agent ahead" not in a0
+
+    # Writer at pos 2: Reviewer ahead, Reviewer behind
+    a2 = agent_rearrange._get_sequential_awareness("Writer", tasks, task_idx=2)
+    assert "Agent ahead" in a2
+    assert "Agent behind" in a2
+
+    # Writer at pos 4: Reviewer ahead, no behind
+    a4 = agent_rearrange._get_sequential_awareness("Writer", tasks, task_idx=4)
+    assert "Agent ahead" in a4
+    assert "Agent behind" not in a4
+
+    print("✓ test_repeated_agent_three_occurrences passed")
+
+
+def test_repeated_agent_run():
+    """Test that a repeated agent flow runs end-to-end."""
+    agents = create_repeated_flow_agents()
+
+    agent_rearrange = AgentRearrange(
+        name="repeated-flow-test",
+        agents=agents,
+        flow="Writer -> Reviewer -> Writer",
+        max_loops=1,
+    )
+
+    result = agent_rearrange.run("Write about the moon.")
+    assert result is not None
+    assert len(str(result)) > 0
+
+    # Verify Writer appears twice in conversation
+    messages = agent_rearrange.conversation.to_dict()
+    writer_msgs = [m for m in messages if m.get("role") == "Writer"]
+    assert len(writer_msgs) == 2, (
+        f"Expected 2 Writer messages, got {len(writer_msgs)}"
+    )
+
+    print("✓ test_repeated_agent_run passed")
+
+
+def test_repeated_agent_awareness_in_conversation():
+    """Test that different awareness messages are injected for each occurrence."""
+    agents = create_repeated_flow_agents()
+
+    agent_rearrange = AgentRearrange(
+        name="awareness-conv-test",
+        agents=agents,
+        flow="Writer -> Reviewer -> Writer",
+        max_loops=1,
+    )
+
+    agent_rearrange.run("Write about rain.")
+
+    messages = agent_rearrange.conversation.to_dict()
+
+    # Find awareness messages that precede Writer messages
+    writer_awareness = []
+    for idx, msg in enumerate(messages):
+        if "Sequential awareness" in str(msg.get("content", "")):
+            if idx + 1 < len(messages) and messages[idx + 1].get("role") == "Writer":
+                writer_awareness.append(msg.get("content", ""))
+
+    assert len(writer_awareness) == 2, (
+        f"Expected 2 awareness messages before Writer, got {len(writer_awareness)}"
+    )
+    assert writer_awareness[0] != writer_awareness[1], (
+        "Both Writer invocations got identical awareness"
+    )
+
+    print("✓ test_repeated_agent_awareness_in_conversation passed")
+
+
 def main():
     """Run all tests."""
     tests = [
@@ -629,6 +839,14 @@ def main():
         test_error_logged_once,
         test_successful_run_returns_result,
         test_successful_callable_returns_result,
+        test_repeated_agent_flow_valid,
+        test_repeated_agent_awareness_position_0,
+        test_repeated_agent_awareness_position_2,
+        test_repeated_agent_awareness_differs_per_position,
+        test_repeated_agent_awareness_fallback_without_idx,
+        test_repeated_agent_three_occurrences,
+        test_repeated_agent_run,
+        test_repeated_agent_awareness_in_conversation,
     ]
 
     print("=" * 60)