Merge branch 'main' of github.com:ur-whitelab/md-agent into conftest

mdagent/mainagent/agent.py

@@ -8,7 +8,7 @@
 from mdagent.utils import PathRegistry, _make_llm
 
 from ..tools import get_tools, make_all_tools
-from .prompt import openaifxn_prompt, structured_prompt
+from .query_filter import make_prompt
 
 load_dotenv()
 
@@ -76,11 +76,6 @@ def __init__(
         else:
             self.skip_subagents = True
 
-        if agent_type == "Structured":
-            self.prompt = structured_prompt
-        elif agent_type == "OpenAIFunctionsAgent":
-            self.prompt = openaifxn_prompt
-
         self.subagents_settings = SubAgentSettings(
             path_registry=path_registry,
             subagents_model=subagents_model,
@@ -124,5 +119,6 @@ def _initialize_tools_and_agent(self, user_input=None):
         )
 
     def run(self, user_input, callbacks=None):
+        self.prompt = make_prompt(user_input, self.agent_type, model="gpt-3.5-turbo")
         self.agent = self._initialize_tools_and_agent(user_input)
-        return self.agent.run(self.prompt.format(input=user_input), callbacks=callbacks)
+        return self.agent.run(self.prompt, callbacks=callbacks)

mdagent/mainagent/prompt.py

@@ -30,7 +30,65 @@
         and you must explicitly state the actions you are taking.
         Question: {input} """,
 )
+modular_analysis_prompt = PromptTemplate(
+    input_variables=[
+        "Main_Task",
+        "Subtask_types",
+        "Proteins",
+        "Parameters",
+        "UserProposedPlan",
+    ],
+    template="""
+        Approach the molecular dynamics inquiry by dissecting it into its modular
+        components:
+        Main Task: {Main_Task}
+        Subtasks: {Subtask_types}
+        Target Proteins: {Proteins}
+        Parameters: {Parameters}
+        Initial Plan Proposed by User: {UserProposedPlan}
+
+        The Main Task is the user's request.
+
+        The Subtasks are (some of/all) the individual steps that may need to be taken
+        to complete the Main Task; Preprocessing/Preparation usually involves
+        cleaning the initial pdb file (adding hydrogens, removing/adding water, etc.)
+        or making the required box for the simulation, Simulation involves running the
+        simulation and/or modifying the simulation script, Postprocessing involves
+        analyzing the results of the simulation (either using provided tools or figuring
+        it out on your own). Finally, Question is used if the user query is more
+        of a question than a request for a specific task.
+
+        the Target Proteins are the protein(s) that the user wants to focus on,
+        the Parameters are the 'special' conditions that the user wants to set and use
+        for the simulation, preprocessing and or analysis.
+
+        Sometimes users already have an idea of what is needed to be done.
+        Initial Plan Proposed by User is the user's initial plan for the simulation. You
+        can use this as a guide to understand what the user wants to do. You can also
+        modify it if you think is necessary.
+
+        You can only respond with a single complete
+        'Thought, Action, Action Input' format
+        OR a single 'Final Answer' format.
+
+        Complete format:
+        Thought: (reflect on your progress and decide what " "to do next)
+        Action: (the action name, should be the name of a tool)
+        Action Input: (the input string to the action)
 
+        OR
+
+        Final Answer: (the final answer to the original input
+        question)
+
+        Use the tools provided, using the most specific tool
+        available for each action.
+        Your final answer should contain all information
+        necessary to answer the question and subquestions.
+        Your thought process should be clean and clear,
+        and you must explicitly state the actions you are taking.
+    """,
+)
 
 openaifxn_prompt = PromptTemplate(
     input_variables=["input"],