aiagents4pharma 1.14.0__py3-none-any.whl → 1.14.1__py3-none-any.whl

aiagents4pharma/talk2biomodels/tests/test_steady_state.py

@@ -0,0 +1,90 @@
+'''
+Test cases for Talk2Biomodels steady state tool.
+'''
+
+from langchain_core.messages import HumanMessage, ToolMessage
+from ..agents.t2b_agent import get_app
+
+def test_steady_state_tool():
+    '''
+    Test the steady_state tool.
+    '''
+    unique_id = 123
+    app = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    #########################################################
+    # In this case, we will test if the tool returns an error
+    # when the model does not achieve a steady state. The tool
+    # status should be "error".
+    prompt = """Run a steady state analysis of model 537."""
+    # Invoke the agent
+    app.invoke(
+        {"messages": [HumanMessage(content=prompt)]},
+        config=config
+    )
+    current_state = app.get_state(config)
+    reversed_messages = current_state.values["messages"][::-1]
+    tool_msg_status = None
+    for msg in reversed_messages:
+        # Assert that the status of the
+        # ToolMessage is "error"
+        if isinstance(msg, ToolMessage):
+            # print (msg)
+            tool_msg_status = msg.status
+            break
+    assert tool_msg_status == "error"
+    #########################################################
+    # In this case, we will test if the tool is indeed invoked
+    # successfully
+    prompt = """Run a steady state analysis of model 64.
+    Set the initial concentration of `Pyruvate` to 0.2. The
+    concentration of `NAD` resets to 100 every 2 time units."""
+    # Invoke the agent
+    app.invoke(
+        {"messages": [HumanMessage(content=prompt)]},
+        config=config
+    )
+    # Loop through the reversed messages until a
+    # ToolMessage is found.
+    current_state = app.get_state(config)
+    reversed_messages = current_state.values["messages"][::-1]
+    steady_state_invoked = False
+    for msg in reversed_messages:
+        # Assert that the message is a ToolMessage
+        # and its status is "error"
+        if isinstance(msg, ToolMessage):
+            print (msg)
+            if msg.name == "steady_state" and msg.status != "error":
+                steady_state_invoked = True
+                break
+    assert steady_state_invoked
+    #########################################################
+    # In this case, we will test if the `ask_question` tool is
+    # invoked upon asking a question about the already generated
+    # steady state results
+    prompt = """What is the Phosphoenolpyruvate concentration
+        at the steady state? Show only the concentration, rounded
+        to 2 decimal places. For example, if the concentration is
+        0.123456, your response should be `0.12`. Do not return
+        any other information."""
+    # Invoke the agent
+    response = app.invoke(
+        {"messages": [HumanMessage(content=prompt)]},
+        config=config
+    )
+    assistant_msg = response["messages"][-1].content
+    current_state = app.get_state(config)
+    reversed_messages = current_state.values["messages"][::-1]
+    # Loop through the reversed messages until a
+    # ToolMessage is found.
+    ask_questool_invoked = False
+    for msg in reversed_messages:
+        # Assert that the message is a ToolMessage
+        # and its status is "error"
+        if isinstance(msg, ToolMessage):
+            if msg.name == "ask_question":
+                ask_questool_invoked = True
+                break
+    assert ask_questool_invoked
+    assert "0.06" in assistant_msg

aiagents4pharma/talk2biomodels/tools/ask_question.py

@@ -6,10 +6,11 @@ Tool for asking a question about the simulation results.
 
 import logging
 from typing import Type, Annotated, Literal
+import hydra
+import basico
 import pandas as pd
 from pydantic import BaseModel, Field
 from langchain_core.tools.base import BaseTool
-from langchain.agents.agent_types import AgentType
 from langchain_experimental.agents import create_pandas_dataframe_agent
 from langchain_openai import ChatOpenAI
 from langgraph.prebuilt import InjectedState
@@ -64,31 +65,51 @@ class AskQuestionTool(BaseTool):
                    question,
                    question_context,
                    experiment_name)
-        # print (f'Calling ask_question tool {question}, {question_context}, {experiment_name}')
+        # Load hydra configuration
+        with hydra.initialize(version_base=None, config_path="../../configs"):
+            cfg = hydra.compose(config_name='config',
+                                overrides=['talk2biomodels/tools/ask_question=default'])
+            cfg = cfg.talk2biomodels.tools.ask_question
+        # Get the context of the question
+        # and based on the context, get the data
+        # and prompt content to ask the question
         if question_context == "steady_state":
             dic_context = state["dic_steady_state_data"]
+            prompt_content = cfg.steady_state_prompt
         else:
             dic_context = state["dic_simulated_data"]
+            prompt_content = cfg.simulation_prompt
+        # Extract the
         dic_data = {}
         for data in dic_context:
             for key in data:
                 if key not in dic_data:
                     dic_data[key] = []
                 dic_data[key] += [data[key]]
-        # print (dic_data)
+        # Create a pandas dataframe of the data
         df_data = pd.DataFrame.from_dict(dic_data)
+        # Extract the data for the experiment
+        # matching the experiment name
         df = pd.DataFrame(
             df_data[df_data['name'] == experiment_name]['data'].iloc[0]
         )
-        prompt_content = None
-        # if run_manager and 'prompt' in run_manager.metadata:
-        #     prompt_content = run_manager.metadata['prompt']
-        # Create a pandas dataframe agent with OpenAI
+        logger.log(logging.INFO, "Shape of the dataframe: %s", df.shape)
+        # # Extract the model units
+        # model_units = basico.model_info.get_model_units()
+        # Update the prompt content with the model units
+        prompt_content += "Following are the model units:\n"
+        prompt_content += f"{basico.model_info.get_model_units()}\n\n"
+        # Create a pandas dataframe agent
         df_agent = create_pandas_dataframe_agent(
                         ChatOpenAI(model=state['llm_model']),
                         allow_dangerous_code=True,
-                        agent_type=AgentType.OPENAI_FUNCTIONS,
+                        agent_type='tool-calling',
                         df=df,
+                        max_iterations=5,
+                        include_df_in_prompt=True,
+                        number_of_head_rows=df.shape[0],
+                        verbose=True,
                         prefix=prompt_content)
+        # Invoke the agent with the question
         llm_result = df_agent.invoke(question)
         return llm_result["output"]

aiagents4pharma/talk2biomodels/tools/get_annotation.py

@@ -7,6 +7,7 @@ based on the provided model and species names.
 import math
 from typing import List, Annotated, Type
 import logging
+from dataclasses import dataclass
 from pydantic import BaseModel, Field
 import basico
 import pandas as pd
@@ -42,18 +43,29 @@ def prepare_content_msg(species_not_found: List[str],
                         {", ".join(species_without_description)}.'''
     return content
 
-class GetAnnotationInput(BaseModel):
+@dataclass
+class ArgumentData:
     """
-    Input schema for annotation tool.
+    Dataclass for storing the argument data.
     """
-    sys_bio_model: ModelData = Field(description="model data")
-    tool_call_id: Annotated[str, InjectedToolCallId]
+    experiment_name: Annotated[str, "An AI assigned _ separated name of"
+                                    " the experiment based on human query"
+                                    " and the context of the experiment."
+                                    " This must be set before the experiment is run."]
     list_species_names: List[str] = Field(
-        default=[],
+        default=None,
         description='''List of species names to fetch annotations for.
                       If not provided, annotations for all
                       species in the model will be fetched.'''
     )
+
+class GetAnnotationInput(BaseModel):
+    """
+    Input schema for annotation tool.
+    """
+    arg_data: ArgumentData = Field(description="argument data")
+    sys_bio_model: ModelData = Field(description="model data")
+    tool_call_id: Annotated[str, InjectedToolCallId]
     state: Annotated[dict, InjectedState]
 
 class GetAnnotationTool(BaseTool):
@@ -70,14 +82,16 @@ class GetAnnotationTool(BaseTool):
     return_direct: bool = False
 
     def _run(self,
+             arg_data: ArgumentData,
              tool_call_id: Annotated[str, InjectedToolCallId],
              state: Annotated[dict, InjectedState],
-             list_species_names: List[str] = None,
              sys_bio_model: ModelData = None) -> str:
         """
         Run the tool.
         """
-        logger.info("Running the GetAnnotationTool tool for species %s", list_species_names)
+        logger.info("Running the GetAnnotationTool tool for species %s, %s",
+                    arg_data.list_species_names,
+                    arg_data.experiment_name)
 
         # Prepare the model object
         sbml_file_path = state['sbml_file_path'][-1] if state['sbml_file_path'] else None
@@ -90,11 +104,11 @@ class GetAnnotationTool(BaseTool):
             # for example this may happen with model 20
             raise ValueError("Unable to extract species from the model.")
         # Fetch annotations for the species names
-        list_species_names = list_species_names or df_species.index.tolist()
+        arg_data.list_species_names = arg_data.list_species_names or df_species.index.tolist()
 
         (annotations_df,
          species_not_found,
-         species_without_description) = self._fetch_annotations(list_species_names)
+         species_without_description) = self._fetch_annotations(arg_data.list_species_names)
 
         # Check if annotations are empty
         # If empty, return a message
@@ -107,6 +121,7 @@ class GetAnnotationTool(BaseTool):
 
         # Prepare the simulated data
         dic_annotations_data = {
+            'name': arg_data.experiment_name,
             'source': sys_bio_model.biomodel_id if sys_bio_model.biomodel_id else 'upload',
             'tool_call_id': tool_call_id,
             'data': annotations_df.to_dict()

aiagents4pharma/talk2biomodels/tools/load_arguments.py

@@ -0,0 +1,114 @@
+#!/usr/bin/env python3
+
+"""
+A utility module for defining the dataclasses
+for the arguments to set up initial settings
+before the experiment is run.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Union, List, Optional, Annotated
+from pydantic import Field
+import basico
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+@dataclass
+class TimeData:
+    """
+    Dataclass for storing the time data.
+    """
+    duration: Union[int, float] = Field(
+        description="Duration of the simulation",
+        default=100)
+    interval: Union[int, float] = Field(
+        description="The interval is the time step or"
+        " the step size of the simulation. It is unrelated"
+        " to the step size of species concentration and parameter values.",
+        default=100)
+
+@dataclass
+class SpeciesInitialData:
+    """
+    Dataclass for storing the species initial data.
+    """
+    species_name: List[str] = Field(
+        description="List of species whose initial concentration is to be set."
+        " This does not include species that reoccur or the species whose"
+        " concentration is to be determined/observed at the end of the experiment."
+        " Do not hallucinate the species name.",
+        default=[])
+    species_concentration: List[Union[int, float]] = Field(
+        description="List of initial concentrations of species."
+        " This does not include species that reoccur or the species whose"
+        " concentration is to be determined/observed at the end of the experiment."
+        " Do not hallucinate the species concentration.",
+        default=[])
+
+@dataclass
+class TimeSpeciesNameConcentration:
+    """
+    Dataclass for storing the time,
+    species name, and concentration data.
+    """
+    time: Union[int, float] = Field(description="time point where the event occurs")
+    species_name: str = Field(description="species name")
+    species_concentration: Union[int, float] = Field(
+        description="species concentration at the time point")
+
+@dataclass
+class ReocurringData:
+    """
+    Dataclass for species that reoccur. In other words,
+    the concentration of the species resets to a certain
+    value after a certain time interval.
+    """
+    data: List[TimeSpeciesNameConcentration] = Field(
+        description="List of time, name, and concentration data"
+                    " of species or parameters that reoccur",
+                    default=[])
+
+@dataclass
+class ArgumentData:
+    """
+    Dataclass for storing the argument data.
+    """
+    experiment_name: Annotated[str, "An AI assigned _ separated name of"
+                                    " the experiment based on human query"
+                                    " and the context of the experiment."
+                                    " This must be set before the experiment is run."]
+    time_data: Optional[TimeData] = Field(
+        description="time data",
+        default=None)
+    species_to_be_analyzed_before_experiment: Optional[SpeciesInitialData] = Field(
+        description="Data of species whose initial concentration"
+        " is to be set before the experiment. This does not include"
+        " species that reoccur or the species whose concentration"
+        " is to be determined at the end of the experiment.",
+        default=None)
+    reocurring_data: Optional[ReocurringData] = Field(
+        description="List of concentration and time data of species that"
+        " reoccur. For example, a species whose concentration resets"
+        " to a certain value after a certain time interval.",
+        default=None)
+
+def add_rec_events(model_object, reocurring_data):
+    """
+    Add reocurring events to the model.
+
+    Args:
+        model_object: The model object.
+        reocurring_data: The reocurring data.
+
+    Returns:
+        None
+    """
+    for row in reocurring_data.data:
+        tp, sn, sc = row.time, row.species_name, row.species_concentration
+        basico.add_event(f'{sn}_{tp}',
+                            f'Time > {tp}',
+                            [[sn, str(sc)]],
+                            model=model_object.copasi_model)

aiagents4pharma/talk2biomodels/tools/parameter_scan.py

@@ -6,7 +6,7 @@ Tool for parameter scan.
 
 import logging
 from dataclasses import dataclass
-from typing import Type, Union, List, Annotated
+from typing import Type, Union, List, Annotated, Optional
 import pandas as pd
 import basico
 from pydantic import BaseModel, Field
@@ -16,61 +16,37 @@ from langchain_core.tools import BaseTool
 from langchain_core.messages import ToolMessage
 from langchain_core.tools.base import InjectedToolCallId
 from .load_biomodel import ModelData, load_biomodel
+from .load_arguments import TimeData, SpeciesInitialData
 
 # Initialize logger
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-@dataclass
-class TimeData:
-    """
-    Dataclass for storing the time data.
-    """
-    duration: Union[int, float] = 100
-    interval: Union[int, float] = 10
-
-@dataclass
-class SpeciesData:
-    """
-    Dataclass for storing the species data.
-    """
-    species_name: List[str] = Field(description="species name", default=[])
-    species_concentration: List[Union[int, float]] = Field(
-        description="initial species concentration",
-        default=[])
-
-@dataclass
-class TimeSpeciesNameConcentration:
-    """
-    Dataclass for storing the time, species name, and concentration data.
-    """
-    time: Union[int, float] = Field(description="time point where the event occurs")
-    species_name: str = Field(description="species name")
-    species_concentration: Union[int, float] = Field(
-        description="species concentration at the time point")
-
-@dataclass
-class ReocurringData:
-    """
-    Dataclass for species that reoccur. In other words, the concentration
-    of the species resets to a certain value after a certain time interval.
-    """
-    data: List[TimeSpeciesNameConcentration] = Field(
-        description="time, name, and concentration data of species that reoccur",
-        default=[])
-
 @dataclass
 class ParameterScanData(BaseModel):
     """
     Dataclass for storing the parameter scan data.
     """
-    species_names: List[str] = Field(description="species names to scan",
-                              default=[])
-    parameter_name: str = Field(description="Parameter name to scan",
-                                 default_factory=None)
-    parameter_values: List[Union[int, float]] = Field(
-        description="Parameter values to scan",
-        default_factory=None)
+    species_names: List[str] = Field(
+                    description="species to be observed after each scan."
+                    " These are the species whose concentration"
+                    " will be observed after the parameter scan."
+                    " Do not make up this data.",
+                    default=[])
+    species_parameter_name: str = Field(
+                    description="Species or parameter name to be scanned."
+                    " This is the species or parameter whose value will be scanned"
+                    " over a range of values. This does not include the species"
+                    " that are to be observed after the scan."
+                    "Do not make up this data.",
+                    default=None)
+    species_parameter_values: List[Union[int, float]] = Field(
+                    description="Species or parameter values to be scanned."
+                    " These are the values of the species or parameters that will be"
+                    " scanned over a range of values. This does not include the "
+                    "species that are to be observed after the scan."
+                    "Do not make up this data.",
+                    default=None)
 
 @dataclass
 class ArgumentData:
@@ -78,30 +54,20 @@ class ArgumentData:
     Dataclass for storing the argument data.
     """
     time_data: TimeData = Field(description="time data", default=None)
-    species_data: SpeciesData = Field(
-        description="species name and initial concentration data",
-        default=None)
-    reocurring_data: ReocurringData = Field(
-        description="""Concentration and time data of species that reoccur
-            For example, a species whose concentration resets to a certain value
-            after a certain time interval""")
+    species_to_be_analyzed_before_experiment: Optional[SpeciesInitialData] = Field(
+                    description=" This is the initial condition of the model."
+                    " This does not include species that reoccur or the species"
+                    " whose concentration is to be determined/observed at the end"
+                    " of the experiment. This also does not include the species"
+                    " or the parameter that is to be scanned. Do not make up this data.",
+                    default=None)
     parameter_scan_data: ParameterScanData = Field(
-        description="parameter scan data",
-        default=None)
-    scan_name: str = Field(
-        description="""An AI assigned `_` separated name of
-        the parameter scan experiment based on human query""")
-
-def add_rec_events(model_object, reocurring_data):
-    """
-    Add reocurring events to the model.
-    """
-    for row in reocurring_data.data:
-        tp, sn, sc = row.time, row.species_name, row.species_concentration
-        basico.add_event(f'{sn}_{tp}',
-                            f'Time > {tp}',
-                            [[sn, str(sc)]],
-                            model=model_object.copasi_model)
+                    description="parameter scan data",
+                    default=None)
+    experiment_name: str = Field(
+                    description="An AI assigned `_` separated unique name of"
+                    " the parameter scan experiment based on human query."
+                    " This must be unique for each experiment.")
 
 def make_list_dic_scanned_data(dic_param_scan, arg_data, sys_bio_model, tool_call_id):
     """
@@ -125,13 +91,18 @@ def make_list_dic_scanned_data(dic_param_scan, arg_data, sys_bio_model, tool_cal
         # Prepare the list dictionary of scanned data
         # that will be passed to the state of the graph
         list_dic_scanned_data.append({
-            'name': arg_data.scan_name+':'+species_name,
+            'name': arg_data.experiment_name+':'+species_name,
             'source': sys_bio_model.biomodel_id if sys_bio_model.biomodel_id else 'upload',
             'tool_call_id': tool_call_id,
             'data': df_param_scan.to_dict()
         })
     return list_dic_scanned_data
-def run_parameter_scan(model_object, arg_data, dic_species_data, duration, interval) -> dict:
+
+def run_parameter_scan(model_object,
+                       arg_data,
+                       dic_species_data,
+                       duration,
+                       interval) -> dict:
     """
     Run parameter scan on the model.
 
@@ -146,44 +117,61 @@ def run_parameter_scan(model_object, arg_data, dic_species_data, duration, inter
         dict: Dictionary of parameter scan results. Each key is a species name
         and each value is a DataFrame containing the results of the parameter scan.
     """
-    # Extract all parameter names from the model and verify if the given parameter name is valid
+    # Extract all parameter names from the model
     df_all_parameters = basico.model_info.get_parameters(model=model_object.copasi_model)
     all_parameters = df_all_parameters.index.tolist()
-    if arg_data.parameter_scan_data.parameter_name not in all_parameters:
-        logger.error(
-            "Invalid parameter name: %s", arg_data.parameter_scan_data.parameter_name)
-        raise ValueError(
-            f"Invalid parameter name: {arg_data.parameter_scan_data.parameter_name}")
-    # Extract all species name from the model and verify if the given species name is valid
+
+    # Extract all species name from the model
     df_all_species = basico.model_info.get_species(model=model_object.copasi_model)
     all_species = df_all_species['display_name'].tolist()
+
+    # Verify if the given species or parameter names to be scanned are valid
+    if arg_data.parameter_scan_data.species_parameter_name not in all_parameters + all_species:
+        logger.error(
+            "Invalid species or parameter name: %s",
+            arg_data.parameter_scan_data.species_parameter_name)
+        raise ValueError(
+            "Invalid species or parameter name: "
+            f"{arg_data.parameter_scan_data.species_parameter_name}.")
+
     # Dictionary to store the parameter scan results
     dic_param_scan_results = {}
+
+    # Loop through the species names that are to be observed
     for species_name in arg_data.parameter_scan_data.species_names:
+        # Verify if the given species name to be observed is valid
         if species_name not in all_species:
             logger.error("Invalid species name: %s", species_name)
-            raise ValueError(f"Invalid species name: {species_name}")
+            raise ValueError(f"Invalid species name: {species_name}.")
+
+        # Copy the model object to avoid modifying the original model
+        model_object_copy = model_object.model_copy()
+
         # Update the fixed model species and parameters
         # These are the initial conditions of the model
         # set by the user
-        model_object.update_parameters(dic_species_data)
+        model_object_copy.update_parameters(dic_species_data)
+
         # Initialize empty DataFrame to store results
         # of the parameter scan
         df_param_scan = pd.DataFrame()
-        for param_value in arg_data.parameter_scan_data.parameter_values:
+
+        # Loop through the parameter that are to be scanned
+        for param_value in arg_data.parameter_scan_data.species_parameter_values:
             # Update the parameter value in the model
-            model_object.update_parameters(
-                {arg_data.parameter_scan_data.parameter_name: param_value})
+            model_object_copy.update_parameters(
+                {arg_data.parameter_scan_data.species_parameter_name: param_value})
             # Simulate the model
-            model_object.simulate(duration=duration, interval=interval)
+            model_object_copy.simulate(duration=duration, interval=interval)
             # If the column name 'Time' is not present in the results DataFrame
             if 'Time' not in df_param_scan.columns:
-                df_param_scan['Time'] = model_object.simulation_results['Time']
+                df_param_scan['Time'] = model_object_copy.simulation_results['Time']
             # Add the simulation results to the results DataFrame
-            col_name = f"{arg_data.parameter_scan_data.parameter_name}_{param_value}"
-            df_param_scan[col_name] = model_object.simulation_results[species_name]
+            col_name = f"{arg_data.parameter_scan_data.species_parameter_name}_{param_value}"
+            df_param_scan[col_name] = model_object_copy.simulation_results[species_name]
 
         logger.log(logging.INFO, "Parameter scan results with shape %s", df_param_scan.shape)
+
         # Add the results of the parameter scan to the dictionary
         dic_param_scan_results[species_name] = df_param_scan
     # return df_param_scan
@@ -210,8 +198,9 @@ class ParameterScanTool(BaseTool):
     Tool for parameter scan.
     """
     name: str = "parameter_scan"
-    description: str = """A tool to perform parameter scan
-        of a list of parameter values for a given species."""
+    description: str = """A tool to perform scanning of a given
+    parameter over a range of values and observe the effect on
+    the concentration of a given species"""
     args_schema: Type[BaseModel] = ParameterScanInput
 
     def _run(self,
@@ -245,12 +234,18 @@ class ParameterScanTool(BaseTool):
         dic_species_data = {}
         if arg_data:
             # Prepare the dictionary of species data
-            if arg_data.species_data is not None:
-                dic_species_data = dict(zip(arg_data.species_data.species_name,
-                                            arg_data.species_data.species_concentration))
-            # Add reocurring events (if any) to the model
-            if arg_data.reocurring_data is not None:
-                add_rec_events(model_object, arg_data.reocurring_data)
+            if arg_data.species_to_be_analyzed_before_experiment is not None:
+                dic_species_data = dict(
+                    zip(
+                        arg_data.species_to_be_analyzed_before_experiment.species_name,
+                        arg_data.species_to_be_analyzed_before_experiment.species_concentration
+                        )
+                    )
+
+            # # Add reocurring events (if any) to the model
+            # if arg_data.reocurring_data is not None:
+            #     add_rec_events(model_object, arg_data.reocurring_data)
+
             # Set the duration and interval
             if arg_data.time_data is not None:
                 duration = arg_data.time_data.duration
@@ -284,7 +279,7 @@ class ParameterScanTool(BaseTool):
                 # update the message history
                 "messages": [
                     ToolMessage(
-                        content=f"Parameter scan results of {arg_data.scan_name}",
+                        content=f"Parameter scan results of {arg_data.experiment_name}",
                         tool_call_id=tool_call_id
                         )
                     ],