fix test_setup + fix rdg tool init

mdagent/tools/base_tools/analysis_tools/rdf_tool.py

@@ -20,9 +20,7 @@ class RDFToolInput(BaseModel):
         description="Selections for RDF. Do not use for now. As "
         "it will only calculate RDF for protein and water molecules.",
     )
-    # atom_indices: Optional[List[int]] = Field(
-    #    None, description="Atom indices to load in the trajectory"
-    # )
+
     # TODO: Add pairs of atoms to calculate RDF within the tool
     ##pairs: Optional[str] = Field(None, description="Pairs of atoms to calculate RDF ")
 
@@ -45,8 +43,9 @@ class RDFTool(BaseTool):
     args_schema = RDFToolInput
     path_registry: Optional[PathRegistry]
 
-    # def __init__(self, path_registry: PathRegistry):
-    #    self.path_registry = path_registry
+    def __init__(self, path_registry: PathRegistry):
+        super().__init__()
+        self.path_registry = path_registry
 
     def _run(self, input):
         try:

mdagent/tools/base_tools/preprocess_tools/clean_tools.py

@@ -1,10 +1,10 @@
 import os
-from typing import Dict, Optional, Type
+from typing import Optional, Type
 
 from langchain.tools import BaseTool
 from openmm.app import PDBFile, PDBxFile
 from pdbfixer import PDBFixer
-from pydantic import BaseModel, Field, root_validator
+from pydantic import BaseModel, Field
 
 from mdagent.utils import FileType, PathRegistry
 
@@ -227,12 +227,6 @@ class CleaningToolFunctionInput(BaseModel):
     )
     add_hydrogens_ph: int = Field(7.0, description="pH at which hydrogens are added.")
 
-    @root_validator(skip_on_failure=True)
-    def validate_query(cls, values) -> Dict:
-        """Check that the input is valid."""
-
-        return values
-
 
 class CleaningToolFunction(BaseTool):
     name = "CleaningToolFunction"

mdagent/tools/base_tools/preprocess_tools/pdb_fix.py

@@ -663,7 +663,6 @@ class PDBFilesFixInp(BaseModel):
     @root_validator(skip_on_failure=True)
     def validate_input(cls, values: Union[str, Dict[str, Any]]) -> Dict:
         if isinstance(values, str):
-            print("values is a string", values)
             raise ValidationError("Input must be a dictionary")
 
         pdbfile = values.get("pdbfiles", "")

tests/test_setup.py

@@ -7,7 +7,7 @@
 
 @pytest.fixture
 def setupandrun(get_registry):
-    return SetUpandRunFunction(get_registry("raw", False))
+    return SetUpandRunFunction(path_registry=get_registry("raw", False))
 
 
 def test_parse_cutoff(setupandrun):