Update app.py

app.py

@@ -656,24 +656,38 @@ def first_node(state: State) -> State:
         query_chembl = None
       state["query_chembl"] = query_chembl
 
-  prompt = f'For the QUERY_TASK given below, determine if one or two of the tools descibed below \
-can complete the task. If so, reply with only the tool names followed by "#". If two tools \
-are required, reply with both tool names separated by a comma and followed by "#". \
-If the tools cannot complete the task, reply with "None #".\n \
-QUERY_TASK: {query_task}.\n \
-Tools: \n \
-uniprot_tool: this tool takes in the user requested protein and searches UNIPROT for matches. \
-It returns a string containing the protein ID, gene name, organism, and protein name.\n \
-list_bioactives_tool: Accepts a given UNIPROT ID and searches for bioactive molecules \n \
-get_bioactives_tool: Accepts a Chembl ID and get all bioactives molecule SMILES and IC50s for that ID\n \
-pdb_tool: Accepts a PDB ID and queires the protein databank for the number of chains in and sequence of the \n \
-protein, as well as other information such as ligands in the structure. \
-find_tool: Accepts a protein name and seaches for PDB IDs that match, returning the PDB ID and the title. \
-predict_tool: Predicts the IC50 value for the molecule indicated by the SMILES string provided. \
-Uses the LightGBM model. \n \
-gpt_tool: Uses a machine-learning GPT model to generate novel molecules for a chembl dataset. It returns a list  \
-of novel molecules generated by the GPT. \
-'
+  prompt = f'''
+# For the QUERY_TASK given below, determine if one or two of the tools descibed below 
+can complete the task. If so, reply with only the tool names followed by "#". If two tools
+are required, reply with both tool names separated by a comma and followed by "#".
+If the tools cannot complete the task, reply with "None #".
+
+## QUERY_TASK: {query_task}.
+
+## Tools:
+
+- uniprot_tool: this tool takes in the user requested protein and searches UNIPROT for matches.
+
+- It returns a string containing the uniprot protein ID, gene name, organism, and protein name.
+
+- list_bioactives_tool: Accepts a given UNIPROT ID and searches for Chembl IDs and  bioactive molecules. 
+Returns Chembl IDs and numbers of bioactive molecules.
+
+- get_bioactives_tool: Accepts a Chembl ID and get all bioactives molecule SMILES and IC50s for that ID. Requires a 
+chembl ID, so the list_bioactives_tool should be called before this tool.
+
+- pdb_tool: Accepts a PDB ID and queires the protein databank for the number of chains in and sequence of the
+protein, as well as other information such as ligands in the structure.
+
+- find_tool: Accepts a protein name and seaches for PDB IDs that match, returning the PDB ID and the title.
+
+- predict_tool: Predicts the IC50 value for the molecule indicated by the SMILES string provided
+Uses the LightGBM model for prediction. Requires a Chembl dataset, so the get_bioactives_tool should be called before this tool.
+
+- gpt_tool: Uses a machine-learning GPT model to generate novel molecules for a chembl dataset. It returns a list
+of novel molecules generated by the GPT and an image of the molecules. Requires a Chembl dataset, so the get_bioactives_tool
+should be called before this tool.
+'''
   res = chat_model.invoke(prompt)
 
   tool_choices = str(res).split('<|assistant|>')[1].split('#')[0].strip()
@@ -719,33 +733,48 @@ def retry_node(state: State) -> State:
   query_pdb = state["query_pdb"]
   query_smiles = state["query_smiles"]
 
-  prompt = f'You were previously given the QUERY_TASK below, and asked to determine if one \
-or two of the tools described below could complete the task. The tool choices did not succeed. \
-Please re-examine the tool choices and determine if one or two of the tools described below \
-can complete the task. If so, reply with only the tool names followed by "#". If two tools \
-are required, reply with both tool names separated by a comma and followed by "#". \
-If the tools cannot complete the task, reply with "None #".\n \
-The information provided by the user is:\n \
-QUERY_PROTEIN: {query_protein}.\n \
-QUERY_UP_ID: {query_up_id}.\n \
-QUERY_CHEMBL: {query_chembl}.\n \
-QUERY_PDB: {query_pdb}.\n \
-QUERY_SMILES: {query_smiles}.\n \
-The task is: \
-QUERY_TASK: {query_task}.\n \
-Tool options: \n \
-uniprot_tool: this tool takes in the user requested protein and searches UNIPROT for matches. \
-It returns a string containing the protein ID, gene name, organism, and protein name.\n \
-list_bioactives_tool: Accepts a given UNIPROT ID and searches for bioactive molecules \n \
-get_bioactives_tool: Accepts a Chembl ID and get all bioactives molecule SMILES and IC50s for that ID\n \
-pdb_tool: Accepts a PDB ID and queires the protein databank for the number of chains in and sequence of the \
-protein, as well as other information such as ligands in the structure. \n \
-find_tool: Accepts a protein name and seaches for PDB IDs that match, returning the PDB ID and the title. \
-predict_tool: Predicts the IC50 value for the molecule indicated by the SMILES string provided. \
-Uses the LightGBM model. \n \
-gpt_tool: Uses a machine-learning GPT model to generate novel molecules for a chembl dataset. It returns a list  \
-of novel molecules generated by the GPT. \
-'
+  prompt = f'''
+# You were previously given the QUERY_TASK below, and asked to determine if one 
+or two of the tools described below could complete the task. The tool choices did not succeed.
+Please re-examine the tool choices and determine if one or two of the tools described below
+can complete the task. If so, reply with only the tool names followed by "#". If two tools
+are required, reply with both tool names separated by a comma and followed by "#".
+If the tools cannot complete the task, reply with "None #".
+
+## The information provided by the user is:
+- QUERY_PROTEIN: {query_protein}.
+- QUERY_UP_ID: {query_up_id}.
+- QUERY_CHEMBL: {query_chembl}.
+- QUERY_PDB: {query_pdb}.
+- QUERY_SMILES: {query_smiles}.
+
+## The task is: 
+- QUERY_TASK: {query_task}.
+
+## Tools:
+
+- uniprot_tool: this tool takes in the user requested protein and searches UNIPROT for matches.
+
+- It returns a string containing the uniprot protein ID, gene name, organism, and protein name.
+
+- list_bioactives_tool: Accepts a given UNIPROT ID and searches for Chembl IDs and  bioactive molecules. 
+Returns Chembl IDs and numbers of bioactive molecules.
+
+- get_bioactives_tool: Accepts a Chembl ID and get all bioactives molecule SMILES and IC50s for that ID. Requires a 
+chembl ID, so the list_bioactives_tool should be called before this tool.
+
+- pdb_tool: Accepts a PDB ID and queires the protein databank for the number of chains in and sequence of the
+protein, as well as other information such as ligands in the structure.
+
+- find_tool: Accepts a protein name and seaches for PDB IDs that match, returning the PDB ID and the title.
+
+- predict_tool: Predicts the IC50 value for the molecule indicated by the SMILES string provided
+Uses the LightGBM model for prediction. Requires a Chembl dataset, so the get_bioactives_tool should be called before this tool.
+
+- gpt_tool: Uses a machine-learning GPT model to generate novel molecules for a chembl dataset. It returns a list
+of novel molecules generated by the GPT and an image of the molecules. Requires a Chembl dataset, so the get_bioactives_tool
+should be called before this tool.
+'''
 
   res = chat_model.invoke(prompt)
 
@@ -978,13 +1007,19 @@ graph = builder.compile()
 @spaces.GPU
 def ProteinAgent(task, protein, up_id, chembl_id, pdb_id, smiles):
   '''
-  This Agent can perform several protein-related tasks. It can find UNIPROT IDs for a protein, or, given
-  a UNIPROT ID it can find Chembl IDs that match. It can find numbers of and lists of bioactive molecules
-  based on a Chembl ID. It can query the protein databank to find PDB IDs matching a protein name and return
-  the IDs and titles. It can find a particular PDB ID and report information such as how many chains it contains, 
-  the sequence, and any small molecules or ligands bound in the structure. It can predict the IC50 value of a molecule
-  based on a Chembl dataset using the LightGBM model. It can generate novel molecules using a finetuned GPT based on
-  a Chembl dataset.
+  This Agent can perform several protein-related tasks. 
+  1. It can find UNIPROT IDs for a protein, or, 
+  2. given a UNIPROT ID it can find Chembl IDs that match. 
+  3. It can find numbers of and lists of bioactive molecules based on a Chembl ID. 
+  4. It can query the protein databank to find PDB IDs matching a protein name and return the IDs and titles. 
+  5. It can find a particular PDB ID and report information such as how many chains it contains, 
+  the sequence, and any small molecules or ligands bound in the structure. 
+  6. It can predict the IC50 value of a molecule based on a Chembl dataset using the LightGBM model. 
+  7. It can generate novel molecules using a finetuned GPT based on a Chembl dataset.
+
+  If Task 6 or 7 are to be called, a chembl dataset is needed. If a Chembl ID is not provided, then task 2 should be called 
+  first to find chembl IDs, then task 4 should be called to collect the dataset based on the ID. If a chembl ID is provided, 
+  then task 4 should be called to collect the chembl dataset. 
 
       Args:
           task: the task to carry out