fixing the submit answer tool and the search tool schema

Author: Jgmedina95

mdcrow/ldp_env/__init__.py

@@ -1,3 +1,8 @@
 from .state import MDCrowState
+from .environment import MDCrowEnv,MySimpleAgent
+from .utils import PathRegistry
+__all__ = ["MDCrowState",
+           "MDCrowEnv",
+           "PathRegistry",
+           "MySimpleAgent"]
 
-__all__ = ["MDCrowState"]

mdcrow/ldp_env/environment.py

@@ -6,7 +6,24 @@
 import os
 from typing import Any, Literal
 
-from analysis_tools import (
+from aviary.core import (
+    Environment,
+    Message,
+    Messages,
+    Tool,
+    ToolRequestMessage,
+    ToolResponseMessage,
+)
+
+# get secrets from environment variables
+from dotenv import load_dotenv
+from pydantic import BaseModel, ConfigDict, Field
+
+from ldp.agent import Agent
+from ldp.graph import LLMCallOp, OpResult, compute_graph
+from mdcrow.ldp_env.state import MDCrowState
+
+from .analysis_tools import (
     compute_bond_angles,
     compute_contacts,
     compute_distance,
@@ -24,18 +41,7 @@
     perform_pca_analysis,
     summarize_protein_structure,
 )
-from aviary.core import (
-    Environment,
-    Message,
-    Messages,
-    Tool,
-    ToolRequestMessage,
-    ToolResponseMessage,
-)
-
-# get secrets from environment variables
-from dotenv import load_dotenv
-from preprocess_tools import (
+from .preprocess_tools import (
     GetActiveSites,
     GetAllKnownSites,
     GetAllSequences,
@@ -57,13 +63,8 @@
     get_small_molecule_PDB,
     pack_molecules,
 )
-from pydantic import BaseModel, ConfigDict, Field
-from simulation_tools import modify_simulation_script, setup_and_run_simulation
-from util_tools import ListRegistryPaths, MapPath2Name, scholar2result_llm
-
-from ldp.agent import Agent
-from ldp.graph import LLMCallOp, OpResult, compute_graph
-from mdcrow.ldp_env.state import MDCrowState
+from .simulation_tools import modify_simulation_script, setup_and_run_simulation
+from .util_tools import ListRegistryPaths, MapPath2Name, scholar2result_llm
 
 load_dotenv()
 
@@ -312,30 +313,29 @@ async def step(
             False,
         )
 
-    def submit_answer(self, answer: str) -> tuple[bool, float, Literal[True]]:
+    def submit_answer(
+        self, answer: str, finished: bool
+    ) -> tuple[str, float, Literal[True]]:
         """Submit the proposed answer and check if it is correct. This action is terminal.
 
         Args:
             answer: Proposed answer.
+            finished: Whether the task is finished.
 
         Returns:
             Three-tuple of if correct, associated reward (correct_reward if correct,
                 tool_failure_reward if tool failure, otherwise incorrect_reward), and
                 True indicating done.
         """
         try:
-            correct: bool = (
-                abs(float(answer) - self.answer)
-                / (abs(self.answer) + self.config.rel_tol)
-                < self.config.rel_tol
-            )
+            answer = answer.strip()
             reward = (
-                self.config.correct_reward if correct else self.config.incorrect_reward
+                self.config.correct_reward if finished else self.config.incorrect_reward
             )
         except ValueError:
-            return False, self.config.tool_failure_reward, True
+            return answer, self.config.tool_failure_reward, True
         else:
-            return correct, reward, True
+            return answer, reward, True
 
     def calculator(self, expr: str) -> tuple[float | str, float, bool]:
         """Calculate a mathematical expression.

mdcrow/ldp_env/util_tools/search_tools.py

@@ -1,15 +1,19 @@
 import paperqa
 
+from mdcrow.ldp_env.state import MDCrowState
 llm_model_args = {
     "name": "gpt-4o-2024-08-06",
     "temperature": 0.5,
 }
 
 
-def scholar2result_llm(state, query):
+def scholar2result_llm(
+        state:MDCrowState, 
+        query:str
+        ):
     """
-     Useful to answer questions that may be found in literature.
-    Ask a specific question as the input.
+    Useful to answer questions that may be found in literature. Ask a 
+    specific question as the input.
 
     Args:
         state (MDCrowState): The state of the MDCrow environment.