implement multi-agent system to create & manage new tools (#46)

* first commit from cocoding aug31

* added skill and curriculum, did some clean up in tools dir

* del extra files, need to check inits

* removed duplicated tool files

* add scaffolding for subagent-based tools

* reorganized subagent-based tools

* changed and improved how subagents are initialized

* added remaining prompts

* completed add_new_tool fxn in skill agent and added new prompts

* updated and completed refining curriculum and iterator code

* adding 'Agent' suffix to the subagent classes (#28)

* reformat init files for circular import issue

* fixed init files

* fixed subagent_setup, passed all pre-commits

* updated iterator code and fixed few bugs

* fixed circular import by adding new utils dir

* test util functions for skill and action agent. Bugs: Agent-change sys to io Skill-add ['text'] at the output from llmchains

* fixed clean_tools so it can use mapped names correctly

* updated and tested skill agent

* debugged subagents & fixed iterator code (#38)

* debugged subagents & fixed iterator code

* tested and debugged agents using AddNewTool tool

* added ExecuteSkill tool to use new tools

* tested and debugged ExecuteSkill tool

* Implement VectorDB and simplify adding new tools (#43)

* implement vectordb for skills and skill retrieval tool

* fixed chroma version due to backward comptability issue

* updated and fixed SkillRetrieval tool

* add tool retrieval for MDAgent

* improve ExecuteSkill and all subagent tools with defined arg_schema

* revamp skill subagent, simplify adding tools

* changed action little to provide function name

* pass list of tools to CreateNewTool iterator code

* debugged, simplifed using skills for current & next prompts

* fixed multi-input tool issue

* fixed get_pdb unit test

* action agent tests failed - remove for now

* fixed and addressed reviewer comments

* tidy up and regroup tools

* update curriculum agent into its own separate tool

* fixed base_tool imports

* clean up notebooks and tests. debugged and fixed few tools

---------

Co-authored-by: Samantha Cox <samantha.wright822@gmail.com>
Co-authored-by: Jorge <97254349+Jgmedina95@users.noreply.github.com>
Co-authored-by: Jorge <jmedina9@ur.rochester.edu>

.env.example

@@ -6,3 +6,6 @@ OPENAI_API_KEY=YOUR_OPENAI_API_KEY_GOES_HERE  # pragma: allowlist secret
 
 # PQA API Key
 PQA_API_KEY=YOUR_PQA_API_KEY_GOES_HERE  # pragma: allowlist secret
+
+# Serp API key
+SERP_API_KEY=YOUR_SERP_API_KEY_GOES_HERE  # pragma: allowlist secret

.gitignore

@@ -136,6 +136,12 @@ dmypy.json
 *.cif
 !3pqr.cif
 
-# generated files for testing
+# output files
+*.csv
+*.png
+ckpt/
+Visualization.ipynb
+
+# path registry files
 *registry.json
 *paths_registry.json

.secrets.baseline

@@ -5,3 +5,6 @@ OpenAI API Key: \b[secrets]{3}_[a-zA-Z0-9]{32}\b
 
 # Rule for detecting pqa API keys
 PQA API Key: "pqa[a-zA-Z0-9-._]+"
+
+# Rule for detecting serp API keys
+# Serp API Key: "[a-zA-Z0-9]{64}"

CONTRIBUTING.md

@@ -11,7 +11,6 @@ git clone https://github.com/ur-whitelab/md-agent.git
 cd md-agent
 conda env create -n mdagent -f environment.yaml
 conda activate mdagent
-conda install -c conda-forge openmm pdbfixer
 ```
 
 ### Installing the Package and Dependencies and Configuring Pre-commit Hooks

mdagent/__init__.py

@@ -1,3 +1,3 @@
-from .agent import MDAgent
+from .mainagent import MDAgent
 
 __all__ = ["MDAgent"]

mdagent/agent_cheat_sheet.md

@@ -0,0 +1,151 @@
+# Organization Plan
+```
+├── .github
+├── mdagent
+│   ├── _init_.py
+│   ├── mainagent
+│   │   ├── __init__.py
+│   │   ├── agent.py
+│   │   ├── mdagent_prompt.py
+│   ├── subagents
+│   │   ├── __init__.py
+│   │   ├── subagent_fxns.py    # contains multiagent functions
+│   │   ├── subagent_setup.py   # contains SubAgentInitializer
+│   │   ├── agents
+│   │   │   ├── __init__.py
+│   │   │   ├── skill.py
+│   │   │   ├── currciulum.py
+│   │   │   ├── action.py
+│   │   │   ├── task_critic.py
+│   │   │   ├── code_critic.py
+│   │   ├── prompts
+│   │   │   ├── __init__.py
+│   │   │   ├── action_prompts.py
+│   │   │   ├── critic_prompts.py
+│   │   │   ├── curriculum_prompts.py
+│   │   │   ├── skill_prompts.py
+│   ├── tools
+│   │   ├── __init__.py
+│   │   ├── maketools.py
+│   │   ├── subagent_tools.py
+│   │   ├── base_tools
+│   │   │   └── ...
+└── notebooks
+│   │   ├── ...
+└── tests
+│   │   ├── ...
+├── ...
+```
+
+# Import Chain
+Top-level to lower-level dependencies to prevent circular imports
+```
+mainagent   - depends on everything below
+  ↑
+tools       - depends on subagents, utils
+  ↑
+subagents   - depends on utils
+  ↑
+utils       - depends on nothing
+```
+
+# Agent Cheat Sheet - for development
+
+We are working with multiple agents, who interact with each other in various ways. To make this process more seamless, here is an agent cheat sheet.
+
+For each agent, please use the following guide:
+
+AGENT NAME
+-purpose
+-input(s)
+-output(s)
+-place within agent framework
+-how to call
+
+For example, curriculum's "place within agent framework" might be first step in iteration, after 1st iteration and as a tool within mrkl. This is a rather complex example, but you get the point.
+
+## ActionAgent - 1st
+- runs code, executes code
+- inputs:
+    - files
+    - task
+    - context (user prompt)
+    - skills
+- outputs:
+    - success (boolean, did it execute)
+    - proposed code
+    - code output (or error)
+- first step in first iteration only
+- Action._run_code
+
+## ActionAgent - resume
+- runs code, executes code
+- inputs:
+    - recent history
+    - full_history
+    - skills
+- outputs:
+    - success (boolean, did it execute)
+    - proposed code
+    - code output (or error)
+- lives after curriculum, before code critic in iteration (then action <> critics)
+- Action._run_code
+
+## Code Critic Agent
+- critique code & determine if pass/fail and how to improve
+- inputs:
+    - code (from action)
+    - code output (from action)
+    - task (from curriculum or MRKL)
+    - context (user prompt)
+- outputs:
+    - critique includes:
+        - code quality, validity, success (boolean), critique, suggestions
+- Lives after action in iteration
+    - if success = True --> call task critic
+    - if success = False --> back to action or curriculum (if max iter)
+- CodeCritic._run
+
+## Task Critic Agent
+- critique whether valid code addresses prompt
+- inputs:
+    - files
+    - code (from action)
+    - code output (from action)
+    - task (from curriculum or MRKL)
+    - context (user prompt)
+    - additional information (optional)
+        - used only if invalid formatting on 1st attempt
+- outputs:
+    - success (boolean)
+    - task critique (if success = False)
+- Lives after code critic & is only envoked if code critic deems code successful.
+    - if successful task critic, add to skill library
+    - if unsuccessful, continue
+- TaskCritic._run_task_critic
+
+## Skill Manager
+- creates tool_name and description for new code, wraps it into Langchain tool
+- also stores everything in 'skill_library' directory
+- input:
+    - fxn_name
+    - code
+- output:
+    - langchain tool name
+- Lives whenever new code is successful and needs to store as a tool
+- SkillManager.add_new_tool(fxn_name, code)
+
+## Curriculum Agent
+- proposes a curriculum of subtasks to achieve the original prompt. Useful
+at the beginning of MDAgent's cycle or whenever MDAgent gets stuck. Also
+useful when the user wants to explore.
+- inputs:
+    - original_task (user prompt)
+    - current list of tools/skills (from MDAgent and skill library)
+    - files
+    - failed tasks, if any
+- output:
+    - the rationale for the plan of subtasks
+    - a list of subtasks
+- Lives outside iterator code, as a separate tool.
+- CurriculumAgent.run(original_task, tools, files, failed_tasks)

mdagent/mainagent/agent.py

@@ -0,0 +1,97 @@
+from dotenv import load_dotenv
+from langchain.agents import AgentExecutor
+from langchain.agents.openai_functions_agent.base import OpenAIFunctionsAgent
+from langchain.prompts import PromptTemplate
+
+from mdagent.subagents import SubAgentSettings
+from mdagent.tools import get_tools
+from mdagent.utils import PathRegistry, _make_llm
+
+load_dotenv()
+
+main_prompt = PromptTemplate(
+    input_variables=["input"],
+    template="""
+    You are an expert molecular dynamics scientist and your
+    task is to respond to the question or
+    solve the problem to the best of your ability using
+    the provided tools. Once you map a path to a short name,
+    you may only use that short name in future actions.
+    Here is the input:
+    input: {input}
+    """,
+)
+
+
+class MDAgent:
+    def __init__(
+        self,
+        tools=None,
+        model="gpt-4-1106-preview",  # current name for gpt-4 turbo
+        tools_model="gpt-4-1106-preview",
+        temp=0.1,
+        max_iterations=40,
+        api_key=None,
+        verbose=True,
+        path_registry=None,
+        subagents_model="gpt-4-1106-preview",
+        ckpt_dir="ckpt",
+        resume=False,
+        top_k_tools=10,
+        use_human_tool=False,
+    ):
+        self.llm = _make_llm(model, temp, verbose)
+        self.tools_llm = _make_llm(tools_model, temp, verbose)
+        self.tools = tools
+        self.top_k_tools = top_k_tools
+        self.ckpt_dir = ckpt_dir
+        self.human = use_human_tool
+        if path_registry is None:
+            path_registry = PathRegistry.get_instance()
+        self.subagents_settings = SubAgentSettings(
+            path_registry=path_registry,
+            subagents_model=subagents_model,
+            temp=temp,
+            max_iterations=max_iterations,
+            api_key=api_key,
+            verbose=verbose,
+            ckpt_dir=ckpt_dir,
+            resume=resume,
+            retrieval_top_k=top_k_tools,
+        )
+
+    def run(self, user_prompt):
+        # get tools relevant to user prompt
+        if self.tools is None:
+            tools = get_tools(
+                query=user_prompt,
+                llm=self.tools_llm,
+                subagent_settings=self.subagents_settings,
+                ckpt_dir=self.ckpt_dir,
+                retrieval_top_k=self.top_k_tools,
+                human=self.human,
+            )
+        else:
+            tools = self.tools
+
+        # initialize agent here with retrieved tools
+        self.agent_executor = AgentExecutor.from_agent_and_tools(
+            tools=tools,
+            agent=OpenAIFunctionsAgent.from_llm_and_tools(self.llm, tools),
+            # agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
+            handle_parsing_errors=True,
+            verbose=True,
+            return_intermediate_steps=True,
+        )
+        outputs = self.agent_executor(main_prompt.format(input=user_prompt))
+        # Parse long output (with intermediate steps)
+        outputs["intermediate_steps"]
+        final = ""
+        # for step in intermed:
+        #     final += (
+        #         f"Action: {step[0].tool}\n"
+        #         f"Action Input: {step[0].tool_input}\n"
+        #         f"Observation: {step[1]}\n"
+        #     )
+        final += f"Final Answer: {outputs['output']}"
+        return final

mdagent/subagents/__init__.py

@@ -0,0 +1,20 @@
+from .agents import (
+    ActionAgent,
+    CodeCriticAgent,
+    CurriculumAgent,
+    SkillManager,
+    TaskCriticAgent,
+)
+from .subagent_fxns import Iterator
+from .subagent_setup import SubAgentInitializer, SubAgentSettings
+
+__all__ = [
+    "ActionAgent",
+    "CodeCriticAgent",
+    "CurriculumAgent",
+    "SkillManager",
+    "TaskCriticAgent",
+    "Iterator",
+    "SubAgentInitializer",
+    "SubAgentSettings",
+]

mdagent/subagents/agents/__init__.py

@@ -0,0 +1,13 @@
+from .action import ActionAgent
+from .code_critic import CodeCriticAgent
+from .curriculum import CurriculumAgent
+from .skill import SkillManager
+from .task_critic import TaskCriticAgent
+
+__all__ = [
+    "ActionAgent",
+    "CodeCriticAgent",
+    "CurriculumAgent",
+    "SkillManager",
+    "TaskCriticAgent",
+]