combined the two visualization tools into one (#57)

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",

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",
 ]

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")

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