diff --git a/agents/aggregation_agent.py b/agents/aggregation_agent.py index 5d40724..22836be 100644 --- a/agents/aggregation_agent.py +++ b/agents/aggregation_agent.py @@ -13,6 +13,11 @@ logger = logging.getLogger("MLEvolve") +def _join_branch_summaries(summaries: List[str]) -> str: + separator = "\n" + "-" * 80 + "\n" + return separator.join(summaries) + + def _collect_branch_representatives(agent) -> List[SearchNode]: representatives = [] @@ -111,7 +116,7 @@ def run( ) reference_summaries.append(branch_info) - reference_experiences = "\n" + "-" * 80 + "\n".join(reference_summaries) + reference_experiences = _join_branch_summaries(reference_summaries) prompt: Any = { "Introduction": introduction, diff --git a/tests/test_aggregation_agent.py b/tests/test_aggregation_agent.py new file mode 100644 index 0000000..72b4bcf --- /dev/null +++ b/tests/test_aggregation_agent.py @@ -0,0 +1,22 @@ +import unittest + +from agents.aggregation_agent import _join_branch_summaries + + +class JoinBranchSummariesTest(unittest.TestCase): + def test_places_separator_between_branch_summaries(self): + separator = "\n" + "-" * 80 + "\n" + + result = _join_branch_summaries(["branch one", "branch two", "branch three"]) + + self.assertEqual( + result, + f"branch one{separator}branch two{separator}branch three", + ) + + def test_preserves_a_single_branch_summary(self): + self.assertEqual(_join_branch_summaries(["branch one"]), "branch one") + + +if __name__ == "__main__": + unittest.main()