Co-authored-by: Quinny Campbell <qcampbel@users.noreply.github.com>

ur-whitelab · Jun 24, 2024 · 7de6997 · 7de6997
2 parents 7ebb3a4 + 9beeb0f
commit 7de6997
Show file tree

Hide file tree

Showing 6 changed files with 1,914 additions and 3 deletions.
diff --git a/mdagent/agent/agent.py b/mdagent/agent/agent.py
@@ -40,7 +40,6 @@ def __init__(
         model="gpt-4-1106-preview",  # current name for gpt-4 turbo
         tools_model="gpt-4-1106-preview",
         temp=0.1,
-        max_iterations=40,
         verbose=True,
         ckpt_dir="ckpt",
         top_k_tools=20,  # set "all" if you want to use all tools

diff --git a/mdagent/tools/base_tools/__init__.py b/mdagent/tools/base_tools/__init__.py
@@ -27,6 +27,26 @@
 from .preprocess_tools.clean_tools import CleaningToolFunction
 from .preprocess_tools.packing import PackMolTool
 from .preprocess_tools.pdb_get import ProteinName2PDBTool, SmallMolPDB, get_pdb
+from .preprocess_tools.uniprot import (
+    GetActiveSites,
+    GetAllKnownSites,
+    GetAllSequences,
+    GetBindingSites,
+    GetGeneNames,
+    GetInteractions,
+    GetKineticProperties,
+    GetPDB3DInfo,
+    GetPDBProcessingInfo,
+    GetProteinAssociatedKeywords,
+    GetProteinFunction,
+    GetRelevantSites,
+    GetSequenceInfo,
+    GetSubunitStructure,
+    GetTurnsBetaSheetsHelices,
+    GetUniprotID,
+    MapProteinRepresentation,
+    UniprotID2Name,
+)
 from .simulation_tools.create_simulation import ModifyBaseSimulationScriptTool
 from .simulation_tools.setup_and_run import (
     SetUpandRunFunction,
@@ -39,6 +59,22 @@
 
 __all__ = [
     "CleaningToolFunction",
+    "GetActiveSites",
+    "GetAllKnownSites",
+    "GetAllSequences",
+    "GetBindingSites",
+    "GetGeneNames",
+    "GetInteractions",
+    "GetKineticProperties",
+    "GetPDB3DInfo",
+    "GetPDBProcessingInfo",
+    "GetProteinAssociatedKeywords",
+    "GetProteinFunction",
+    "GetRelevantSites",
+    "GetSequenceInfo",
+    "GetSubunitStructure",
+    "GetTurnsBetaSheetsHelices",
+    "GetUniprotID",
     "ComputeAngles",
     "ComputeChi1",
     "ComputeChi2",
@@ -58,6 +94,7 @@
     "DistanceMatrixTool",
     "ListRegistryPaths",
     "MapPath2Name",
+    "MapProteinRepresentation",
     "ModifyBaseSimulationScriptTool",
     "MomentOfInertia",
     "PackMolTool",
@@ -72,11 +109,12 @@
     "RMSDCalculator",
     "Scholar2ResultLLM",
     "SerpGitTool",
-    "SetUpandRunFunction",
     "SetUpAndRunTool",
+    "SetUpandRunFunction",
     "SimulationFunctions",
     "SimulationOutputFigures",
     "SmallMolPDB",
+    "UniprotID2Name",
     "SolventAccessibleSurfaceArea",
     "VisFunctions",
     "VisualizeProtein",

diff --git a/mdagent/tools/base_tools/preprocess_tools/pdb_get.py b/mdagent/tools/base_tools/preprocess_tools/pdb_get.py
@@ -33,7 +33,8 @@ def get_pdb(query_string: str, path_registry: PathRegistry):
     else:
         filetype = "pdb"
     if "result_set" in r.json() and len(r.json()["result_set"]) > 0:
-        pdbid = r.json()["result_set"][0]["identifier"]
+        results = r.json()["result_set"]
+        pdbid = max(results, key=lambda x: x["score"])["identifier"]
         print(f"PDB file found with this ID: {pdbid}")
         st.markdown(f"PDB file found with this ID: {pdbid}", unsafe_allow_html=True)
         url = f"https://files.rcsb.org/download/{pdbid}.{filetype}"