added curriculum flag and unit test (#86)

* added curriculum flag and unit test * add true/false for current curriculum tool --------- Co-authored-by: qcampbel <[email protected]>
ur-whitelab · Feb 19, 2024 · 427624a · 427624a
1 parent dd0e5f0
commit 427624a
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 1 deletion.
diff --git a/mdagent/mainagent/agent.py b/mdagent/mainagent/agent.py
@@ -47,6 +47,7 @@ def __init__(
         resume=False,
         top_k_tools=20,  # set "all" if you want to use all tools (& skills if resume)
         use_human_tool=False,
+        curriculum=True,
         uploaded_files=[],  # user input files to add to path registry
     ):
         if path_registry is None:
@@ -83,6 +84,7 @@ def __init__(
             verbose=verbose,
             ckpt_dir=ckpt_dir,
             resume=resume,
+            curriculum=curriculum,
         )
 
     def _initialize_tools_and_agent(self, user_input=None):

diff --git a/mdagent/subagents/subagent_setup.py b/mdagent/subagents/subagent_setup.py
@@ -15,6 +15,7 @@ def __init__(
         ckpt_dir="ckpt",
         resume=False,
         retrieval_top_k=5,
+        curriculum=True,
     ):
         self.path_registry = path_registry
         self.subagents_model = subagents_model
@@ -24,6 +25,7 @@ def __init__(
         self.ckpt_dir = ckpt_dir
         self.resume = resume
         self.retrieval_top_k = retrieval_top_k
+        self.curriculum = curriculum
 
 
 class SubAgentInitializer:
@@ -40,6 +42,7 @@ def __init__(self, settings: Optional[SubAgentSettings] = None):
         self.ckpt_dir = settings.ckpt_dir
         self.resume = settings.resume
         self.retrieval_top_k = settings.retrieval_top_k
+        self.curriculum = settings.curriculum
 
     def create_action(self, **overrides):
         params = {
@@ -61,6 +64,8 @@ def create_critic(self, **overrides):
         return Critic(**params)
 
     def create_curriculum(self, **overrides):
+        if not self.curriculum:
+            return None
         params = {
             "model": self.subagents_model,
             "temp": self.temp,

diff --git a/mdagent/tools/maketools.py b/mdagent/tools/maketools.py
@@ -102,8 +102,9 @@ def make_all_tools(
             CreateNewTool(subagent_settings=subagent_settings),
             RetryExecuteSkill(subagent_settings=subagent_settings),
             SkillRetrieval(subagent_settings=subagent_settings),
-            WorkflowPlan(subagent_settings=subagent_settings),
         ]
+        if subagent_settings.curriculum:
+            WorkflowPlan(subagent_settings=subagent_settings)
 
     # add 'learned' tools here
     # disclaimer: assume they don't need path_registry

diff --git a/tests/test_agent.py b/tests/test_agent.py
@@ -3,6 +3,7 @@
 
 import pytest
 
+from mdagent.mainagent.agent import MDAgent
 from mdagent.subagents.agents.action import Action
 from mdagent.subagents.agents.skill import SkillManager
 from mdagent.subagents.subagent_fxns import Iterator
@@ -231,3 +232,11 @@ def test_update_skill_library(skill_manager):
             path="/mock_dir/code/test_function.py",
             description="Code for new tool test_function",
         )
+
+
+# test mdagent with and without curriculum
+def test_mdagent_curriculum():
+    mdagent_curr = MDAgent(curriculum=True)
+    mdagent_no_curr = MDAgent(curriculum=False)
+    assert mdagent_curr.subagents_settings.curriculum is True
+    assert mdagent_no_curr.subagents_settings.curriculum is False