Feat/parameter scan

README.md

@@ -25,6 +25,7 @@ Our toolkit currently consists of three intelligent agents, each designed to sim
 - Forward simulation of both internal and open-source models (BioModels).
 - Adjust parameters within the model to simulate different conditions.
 - Query simulation results.
+- Extract model information such as species, parameters, units and description.
 
 ### 2. Talk2Cells *(Work in Progress)*
 

aiagents4pharma/talk2biomodels/agents/t2b_agent.py

@@ -16,6 +16,7 @@
 from ..tools.simulate_model import SimulateModelTool
 from ..tools.custom_plotter import CustomPlotterTool
 from ..tools.ask_question import AskQuestionTool
+from ..tools.parameter_scan import ParameterScanTool
 from ..states.state_talk2biomodels import Talk2Biomodels
 
 # Initialize logger
@@ -35,17 +36,13 @@ def agent_t2b_node(state: Annotated[dict, InjectedState]):
         return response
 
     # Define the tools
-    simulate_model = SimulateModelTool()
-    custom_plotter = CustomPlotterTool()
-    ask_question = AskQuestionTool()
-    search_model = SearchModelsTool()
-    get_modelinfo = GetModelInfoTool()
     tools = ToolNode([
-                    simulate_model,
-                    ask_question,
-                    custom_plotter,
-                    search_model,
-                    get_modelinfo
+                    SimulateModelTool(),
+                    AskQuestionTool(),
+                    CustomPlotterTool(),
+                    SearchModelsTool(),
+                    GetModelInfoTool(),
+                    ParameterScanTool()
                     ])
 
     # Define the model

aiagents4pharma/talk2biomodels/models/basico_model.py

@@ -48,52 +48,49 @@ def check_biomodel_id_or_sbml_file_path(self):
             self.name = basico.model_info.get_model_name(model=self.copasi_model)
         return self
 
-    def simulate(self,
-                 parameters: Optional[Dict[str, Union[float, int]]] = None,
-                 duration: Union[int, float] = 10,
-                 interval: int = 10
-                 ) -> pd.DataFrame:
+    def update_parameters(self, parameters: Dict[str, Union[float, int]]) -> None:
+        """
+        Update model parameters with new values.
+        """
+        # Update parameters in the model
+        for param_name, param_value in parameters.items():
+            # check if the param_name is not None
+            if param_name is None:
+                continue
+            # if param is a kinetic parameter
+            df_all_params = basico.model_info.get_parameters(model=self.copasi_model)
+            if param_name in df_all_params.index.tolist():
+                basico.model_info.set_parameters(name=param_name,
+                                            exact=True,
+                                            initial_value=param_value,
+                                            model=self.copasi_model)
+            # if param is a species
+            else:
+                basico.model_info.set_species(name=param_name,
+                                            exact=True,
+                                            initial_concentration=param_value,
+                                            model=self.copasi_model)
+
+    def simulate(self, duration: Union[int, float] = 10, interval: int = 10) -> pd.DataFrame:
         """
         Simulate the COPASI model over a specified range of time points.
 
         Args:
-            parameters: Dictionary of model parameters to update before simulation.
             duration: Duration of the simulation in time units.
             interval: Interval between time points in the simulation.
 
         Returns:
             Pandas DataFrame with time-course simulation results.
         """
-
-        # Update parameters in the model
-        if parameters:
-            for param_name, param_value in parameters.items():
-                # check if the param_name is not None
-                if param_name is None:
-                    continue
-                # if param is a kinectic parameter
-                df_all_params = basico.model_info.get_parameters(model=self.copasi_model)
-                if param_name in df_all_params.index.tolist():
-                    basico.model_info.set_parameters(name=param_name,
-                                                exact=True,
-                                                initial_value=param_value,
-                                                model=self.copasi_model)
-                # if param is a species
-                else:
-                    basico.model_info.set_species(name=param_name,
-                                                exact=True,
-                                                initial_concentration=param_value,
-                                                model=self.copasi_model)
-
         # Run the simulation and return results
         df_result = basico.run_time_course(model=self.copasi_model,
                                         intervals=interval,
                                         duration=duration)
-        # Replace curly braces in column headers with square brackets
-        # Because curly braces in the world of LLMS are used for
-        # structured output
-        df_result.columns = df_result.columns.str.replace('{', '[', regex=False).\
-                    str.replace('}', ']', regex=False)
+        # # Replace curly braces in column headers with square brackets
+        # # Because curly braces in the world of LLMS are used for
+        # # structured output
+        # df_result.columns = df_result.columns.str.replace('{', '[', regex=False).\
+        #             str.replace('}', ']', regex=False)
         # Reset the index
         df_result.reset_index(inplace=True)
         # Store the simulation results

aiagents4pharma/talk2biomodels/models/sys_bio_model.py

@@ -35,18 +35,21 @@ def get_model_metadata(self) -> Dict[str, Union[str, int]]:
         Returns:
             dict: Dictionary with model metadata
         """
+    @abstractmethod
+    def update_parameters(self, parameters: Dict[str, Union[float, int]]) -> None:
+        """
+        Abstract method to update model parameters.
+
+        Args:
+            parameters: Dictionary of parameter values.
+        """
 
     @abstractmethod
-    def simulate(self,
-                 parameters: Dict[str, Union[float, int]],
-                 duration: Union[int, float]) -> List[float]:
+    def simulate(self, duration: Union[int, float]) -> List[float]:
         """
         Abstract method to run a simulation of the model.
-        This method should be implemented to simulate model 
-        behavior based on the provided parameters.
 
         Args:
-            parameters: Dictionary of parameter values.
             duration: Duration of the simulation.
 
         Returns:

aiagents4pharma/talk2biomodels/states/state_talk2biomodels.py

@@ -12,13 +12,13 @@ class Talk2Biomodels(AgentState):
     """
     The state for the Talk2BioModels agent.
     """
-    model_id: Annotated[list, operator.add]
-    # sbml_file_path: str
+    llm_model: str
     # A StateGraph may receive a concurrent updates
     # which is not supported by the StateGraph.
     # Therefore, we need to use Annotated to specify
     # the operator for the sbml_file_path field.
     # https://langchain-ai.github.io/langgraph/troubleshooting/errors/INVALID_CONCURRENT_GRAPH_UPDATE/
+    model_id: Annotated[list, operator.add]
     sbml_file_path: Annotated[list, operator.add]
     dic_simulated_data: Annotated[list[dict], operator.add]
-    llm_model: str
+    dic_scanned_data: Annotated[list[dict], operator.add]

aiagents4pharma/talk2biomodels/tests/test_basico_model.py

@@ -19,13 +19,14 @@ def test_with_biomodel_id(model):
     Test initialization of BasicoModel with biomodel_id.
     """
     assert model.biomodel_id == 64
+    model.update_parameters(parameters={'Pyruvate': 0.5, 'KmPFKF6P': 1.5})
+    df_species = basico.model_info.get_species(model=model.copasi_model)
+    assert df_species.loc['Pyruvate', 'initial_concentration'] == 0.5
+    df_parameters = basico.model_info.get_parameters(model=model.copasi_model)
+    assert df_parameters.loc['KmPFKF6P', 'initial_value'] == 1.5
     # check if the simulation results are a pandas DataFrame object
-    assert isinstance(model.simulate(parameters={'Pyruvate': 0.5, 'KmPFKF6P': 1.5},
-                                     duration=2,
-                                     interval=2),
-                    pd.DataFrame)
-    assert isinstance(model.simulate(parameters={None: None}, duration=2, interval=2),
-                    pd.DataFrame)
+    assert isinstance(model.simulate(duration=2, interval=2), pd.DataFrame)
+    model.update_parameters(parameters={None: None})
     assert model.description == basico.biomodels.get_model_info(model.biomodel_id)["description"]
 
 def test_with_sbml_file():
@@ -35,8 +36,6 @@ def test_with_sbml_file():
     model_object = BasicoModel(sbml_file_path="./BIOMD0000000064_url.xml")
     assert model_object.sbml_file_path == "./BIOMD0000000064_url.xml"
     assert isinstance(model_object.simulate(duration=2, interval=2), pd.DataFrame)
-    assert isinstance(model_object.simulate(parameters={'NADH': 0.5}, duration=2, interval=2),
-                      pd.DataFrame)
 
 def test_check_biomodel_id_or_sbml_file_path():
     '''

aiagents4pharma/talk2biomodels/tests/test_langgraph.py

@@ -119,6 +119,68 @@ def test_simulate_model_tool():
     # Check if the data of the second model contains
     assert 'mTORC2' in dic_simulated_data[1]['data']
 
+def test_param_scan_tool():
+    '''
+    In this test, we will test the parameter_scan tool.
+    We will prompt it to scan the parameter `kIL6RBind`
+    from 1 to 100 in steps of 10, record the changes
+    in the concentration of the species `Ab{serum}` in
+    model 537.
+
+    We will pass the inaccuarate parameter (`KIL6Rbind`)
+    and species names (just `Ab`) to the tool to test
+    if it can deal with it.
+
+    We expect the agent to first invoke the parameter_scan
+    tool and raise an error. It will then invoke another
+    tool get_modelinfo to get the correct parameter
+    and species names. Finally, the agent will reinvoke
+    the parameter_scan tool with the correct parameter
+    and species names.
+
+    '''
+    unique_id = 123
+    app = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    prompt = """How will the value of Ab in model 537 change
+            if the param kIL6Rbind is varied from 1 to 100 in steps of 10?
+            Set the initial `DoseQ2W` concentration to 300.
+            Reset the IL6{serum} concentration to 100 every 500 hours.
+            Assume that the model is simulated for 2016 hours with
+            an interval of 2016."""
+    # Invoke the agent
+    app.invoke(
+        {"messages": [HumanMessage(content=prompt)]},
+        config=config
+    )
+    current_state = app.get_state(config)
+    reversed_messages = current_state.values["messages"][::-1]
+    # Loop through the reversed messages until a
+    # ToolMessage is found.
+    df = pd.DataFrame(columns=['name', 'status', 'content'])
+    names = []
+    statuses = []
+    contents = []
+    for msg in reversed_messages:
+        # Assert that the message is a ToolMessage
+        # and its status is "error"
+        if not isinstance(msg, ToolMessage):
+            continue
+        names.append(msg.name)
+        statuses.append(msg.status)
+        contents.append(msg.content)
+    df = pd.DataFrame({'name': names, 'status': statuses, 'content': contents})
+    # print (df)
+    assert any((df["status"] == "error") &
+               (df["name"] == "parameter_scan") &
+               (df["content"].str.startswith("Error: ValueError('Invalid parameter name:")))
+    assert any((df["status"] == "success") &
+               (df["name"] == "parameter_scan") &
+               (df["content"].str.startswith("Parameter scan results of")))
+    assert any((df["status"] == "success") &
+               (df["name"] == "get_modelinfo"))
+
 def test_integration():
     '''
     Test the integration of the tools.
@@ -184,9 +246,9 @@ def test_integration():
     reversed_messages = current_state.values["messages"][::-1]
     # Loop through the reversed messages
     # until a ToolMessage is found.
-    expected_header = ['Time', 'CRP[serum]', 'CRPExtracellular']
+    expected_header = ['Time', 'CRP{serum}', 'CRPExtracellular']
     expected_header += ['CRP Suppression (%)', 'CRP (% of baseline)']
-    expected_header += ['CRP[liver]']
+    expected_header += ['CRP{liver}']
     predicted_artifact = []
     for msg in reversed_messages:
         if isinstance(msg, ToolMessage):

aiagents4pharma/talk2biomodels/tests/test_sys_bio_model.py

@@ -16,22 +16,27 @@ class TestBioModel(SysBioModel):
     sbml_file_path: Optional[str] = Field(None, description="Path to an SBML file")
     name: Optional[str] = Field(..., description="Name of the model")
     description: Optional[str] = Field("", description="Description of the model")
+    param1: Optional[float] = Field(0.0, description="Parameter 1")
+    param2: Optional[float] = Field(0.0, description="Parameter 2")
 
     def get_model_metadata(self) -> Dict[str, Union[str, int]]:
         '''
         Get the metadata of the model.
         '''
         return self.biomodel_id
 
-    def simulate(self,
-                 parameters: Dict[str, Union[float, int]],
-                 duration: Union[int, float]) -> List[float]:
+    def update_parameters(self, parameters):
+        '''
+        Update the model parameters.
+        '''
+        self.param1 = parameters.get('param1', 0.0)
+        self.param2 = parameters.get('param2', 0.0)
+
+    def simulate(self, duration: Union[int, float]) -> List[float]:
         '''
         Simulate the model.
         '''
-        param1 = parameters.get('param1', 0.0)
-        param2 = parameters.get('param2', 0.0)
-        return [param1 + param2 * t for t in range(int(duration))]
+        return [self.param1 + self.param2 * t for t in range(int(duration))]
 
 def test_get_model_metadata():
     '''
@@ -53,5 +58,6 @@ def test_simulate():
     Test the simulate method of the BioModel class.
     '''
     model = TestBioModel(biomodel_id=123, name="Test Model", description="A test model")
-    results = model.simulate(parameters={'param1': 1.0, 'param2': 2.0}, duration=4.0)
+    model.update_parameters({'param1': 1.0, 'param2': 2.0})
+    results = model.simulate(duration=4.0)
     assert results == [1.0, 3.0, 5.0, 7.0]

aiagents4pharma/talk2biomodels/tools/__init__.py

@@ -6,4 +6,5 @@
 from . import ask_question
 from . import custom_plotter
 from . import get_modelinfo
+from . import parameter_scan
 from . import load_biomodel

aiagents4pharma/talk2biomodels/tools/get_modelinfo.py

@@ -47,8 +47,10 @@ class GetModelInfoTool(BaseTool):
     """
     This tool ise used extract model information.
     """
-    name: str = "get_parameters"
-    description: str = "A tool for extracting model information."
+    name: str = "get_modelinfo"
+    description: str = """A tool for extracting name,
+                    description, species, parameters,
+                    compartments, and units from a model."""
     args_schema: Type[BaseModel] = GetModelInfoInput
 
     def _run(self,
@@ -81,7 +83,7 @@ def _run(self,
         # Extract species from the model
         if requested_model_info.species:
             df_species = basico.model_info.get_species(model=model_obj.copasi_model)
-            dic_results['Species'] = df_species.index.tolist()
+            dic_results['Species'] = df_species['display_name'].tolist()
             dic_results['Species'] = ','.join(dic_results['Species'])
 
         # Extract parameters from the model