From 644d246088f6b024827241a8bacb4e86c79a19d6 Mon Sep 17 00:00:00 2001 From: Sam Cox Date: Thu, 21 Dec 2023 08:56:14 -0800 Subject: [PATCH] combined the two visualization tools into one (#57) --- mdagent/tools/base_tools/__init__.py | 6 +- .../base_tools/analysis_tools/__init__.py | 10 +-- .../base_tools/analysis_tools/vis_tools.py | 83 +++++++------------ mdagent/tools/maketools.py | 6 +- 4 files changed, 34 insertions(+), 71 deletions(-) diff --git a/mdagent/tools/base_tools/__init__.py b/mdagent/tools/base_tools/__init__.py index 4209d61a..b1e99c06 100644 --- a/mdagent/tools/base_tools/__init__.py +++ b/mdagent/tools/base_tools/__init__.py @@ -3,9 +3,8 @@ from .analysis_tools.rmsd_tools import RMSDCalculator from .analysis_tools.vis_tools import ( CheckDirectoryFiles, - PlanBVisualizationTool, VisFunctions, - VisualizationToolRender, + VisualizeProtein, ) from .preprocess_tools.clean_tools import ( AddHydrogensCleaningTool, @@ -36,7 +35,7 @@ "Name2PDBTool", "PackMolTool", "PPIDistance", - "PlanBVisualizationTool", + "VisualizeProtein", "RMSDCalculator", "RemoveWaterCleaningTool", "Scholar2ResultLLM", @@ -46,7 +45,6 @@ "SimulationOutputFigures", "SpecializedCleanTool", "VisFunctions", - "VisualizationToolRender", "get_pdb", "CleaningToolFunction", "SetUpandRunFunction", diff --git a/mdagent/tools/base_tools/analysis_tools/__init__.py b/mdagent/tools/base_tools/analysis_tools/__init__.py index 3797718a..2243f0d2 100644 --- a/mdagent/tools/base_tools/analysis_tools/__init__.py +++ b/mdagent/tools/base_tools/analysis_tools/__init__.py @@ -1,19 +1,13 @@ from .plot_tools import SimulationOutputFigures from .ppi_tools import PPIDistance from .rmsd_tools import RMSDCalculator -from .vis_tools import ( - CheckDirectoryFiles, - PlanBVisualizationTool, - VisFunctions, - VisualizationToolRender, -) +from .vis_tools import CheckDirectoryFiles, VisFunctions, VisualizeProtein __all__ = [ "PPIDistance", "RMSDCalculator", "SimulationOutputFigures", "CheckDirectoryFiles", - "PlanBVisualizationTool", + "VisualizeProtein", "VisFunctions", - "VisualizationToolRender", ] diff --git a/mdagent/tools/base_tools/analysis_tools/vis_tools.py b/mdagent/tools/base_tools/analysis_tools/vis_tools.py index 680fa42d..3d9ec17b 100644 --- a/mdagent/tools/base_tools/analysis_tools/vis_tools.py +++ b/mdagent/tools/base_tools/analysis_tools/vis_tools.py @@ -72,29 +72,44 @@ def create_notebook(self, query, PathRegistry): return "Visualization Complete" -class VisualizationToolRender(BaseTool): - """For this tool - to work you need - to instal molrender - https://github.com/molstar/molrender/tree/master""" +class VisualizeProtein(BaseTool): + """To get a png, you must install molrender + https://github.com/molstar/molrender/tree/master + Otherwise, you will get a notebook where you + can visualize the protein.""" name = "PDBVisualization" description = """This tool will create a visualization of a cif file as a png file in - the same directory. if - cif file doesnt exist - you should look for - alternatives in the directory""" + the same directory OR + it will create + a .ipynb file with the + visualization of the + file, depending on the + packages available. + If a notebook is created, + the user can open the + notebook and visualize the + system.""" + path_registry: Optional[PathRegistry] + + def __init__(self, path_registry: Optional[PathRegistry]): + super().__init__() + self.path_registry = path_registry def _run(self, query: str) -> str: """use the tool.""" + vis = VisFunctions() try: - vis = VisFunctions() vis.run_molrender(query) - return "Visualization created" - except Exception as e: - return f"An error occurred while running molrender: {str(e)}" + return "Visualization created as png" + except Exception: + try: + vis.create_notebook(query, self.path_registry) + return "Visualization created as notebook" + except Exception as e: + return f"An error occurred while running molrender: {str(e)}" async def _arun(self, query: str) -> str: """Use the tool asynchronously.""" @@ -119,45 +134,3 @@ def _run(self, query: str) -> str: async def _arun(self, query: str) -> str: """Use the tool asynchronously.""" raise NotImplementedError("custom_search does not support async") - - -class PlanBVisualizationTool(BaseTool): - """This tool will create - a .ipynb file with the - visualization of the - file. It is intended - to be used only - if VisualizationToolRender fails""" - - name = "PlanBVisualizationTool" - description = """This tool will create a .ipynb - file with the visualization - of the file. It is intended - to be used only if - VisualizationToolRender fails. - Give this tool the saved - name of the file - and the output - will be a notebook the - user can use to visualize - the file.""" - path_registry: Optional[PathRegistry] - - def __init__(self, path_registry: Optional[PathRegistry]): - super().__init__() - self.path_registry = path_registry - - def _run(self, query: str) -> str: - """use the tool.""" - try: - if self.path_registry is None: # this should not happen - return "Path registry not initialized" - vis = VisFunctions() - vis.create_notebook(query, self.path_registry) - return "Visualization Complete" - except Exception: - return "An error occurred while creating the notebook" - - async def _arun(self, query: str) -> str: - """Use the tool asynchronously.""" - raise NotImplementedError("custom_search does not support async") diff --git a/mdagent/tools/maketools.py b/mdagent/tools/maketools.py index b50adf52..c9f2891d 100644 --- a/mdagent/tools/maketools.py +++ b/mdagent/tools/maketools.py @@ -22,14 +22,13 @@ ModifyBaseSimulationScriptTool, Name2PDBTool, PackMolTool, - PlanBVisualizationTool, PPIDistance, RMSDCalculator, Scholar2ResultLLM, SerpGitTool, SetUpandRunFunction, SimulationOutputFigures, - VisualizationToolRender, + VisualizeProtein, ) from .subagent_tools import ExecuteSkill, SkillRetrieval, WorkflowPlan @@ -84,13 +83,12 @@ def make_all_tools( MapPath2Name(path_registry=path_instance), Name2PDBTool(path_registry=path_instance), PackMolTool(path_registry=path_instance), - PlanBVisualizationTool(path_registry=path_instance), + VisualizeProtein(path_registry=path_instance), PPIDistance(), RMSDCalculator(), SetUpandRunFunction(path_registry=path_instance), ModifyBaseSimulationScriptTool(path_registry=path_instance, llm=llm), SimulationOutputFigures(), - VisualizationToolRender(), ] # tools using subagents