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

aiagents4pharma/talk2biomodels/tests/test_integration.py

@@ -0,0 +1,126 @@
+'''
+Test cases for Talk2Biomodels.
+'''
+
+import pandas as pd
+from langchain_core.messages import HumanMessage, ToolMessage
+from ..agents.t2b_agent import get_app
+
+def test_integration():
+    '''
+    Test the integration of the tools.
+    '''
+    unique_id = 1234567
+    app = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    # ##########################################
+    # ## Test simulate_model tool
+    # ##########################################
+    prompt = '''Simulate the model 537 for 100 hours and intervals
+    100 with an initial concentration of `DoseQ2W`
+    set to 300 and `Dose` set to 0. Reset the concentration
+    of `Ab{serum}` to 100 every 25 hours.'''
+    # Test the tool get_modelinfo
+    response = app.invoke(
+                        {"messages": [HumanMessage(content=prompt)]},
+                        config=config
+                    )
+    assistant_msg = response["messages"][-1].content
+    print (assistant_msg)
+    # Check if the assistant message is a string
+    assert isinstance(assistant_msg, str)
+    ##########################################
+    # Test ask_question tool when simulation
+    # results are available
+    ##########################################
+    # Update state
+    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    prompt = """What is the concentration of CRP in serum after 100 hours?
+    Round off the value to 2 decimal places."""
+    # Test the tool get_modelinfo
+    response = app.invoke(
+                        {"messages": [HumanMessage(content=prompt)]},
+                        config=config
+                    )
+    assistant_msg = response["messages"][-1].content
+    # print (assistant_msg)
+    # Check if the assistant message is a string
+    assert '211' in assistant_msg
+
+    ##########################################
+    # Test custom_plotter tool when the
+    # simulation results are available
+    ##########################################
+    prompt = "Plot only CRP related species."
+
+    # Update state
+    app.update_state(config, {"llm_model": "gpt-4o-mini"}
+                    )
+    # Test the tool get_modelinfo
+    response = app.invoke(
+                        {"messages": [HumanMessage(content=prompt)]},
+                        config=config
+                    )
+    assistant_msg = response["messages"][-1].content
+    current_state = app.get_state(config)
+    # Get the messages from the current state
+    # and reverse the order
+    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 += ['CRP Suppression (%)', 'CRP (% of baseline)']
+    expected_header += ['CRP{liver}']
+    predicted_artifact = []
+    for msg in reversed_messages:
+        if isinstance(msg, ToolMessage):
+            # Work on the message if it is a ToolMessage
+            # These may contain additional visuals that
+            # need to be displayed to the user.
+            if msg.name == "custom_plotter":
+                predicted_artifact = msg.artifact
+                break
+    # Convert the artifact into a pandas dataframe
+    # for easy comparison
+    df = pd.DataFrame(predicted_artifact)
+    # Extract the headers from the dataframe
+    predicted_header = df.columns.tolist()
+    # Check if the header is in the expected_header
+    # assert expected_header in predicted_artifact
+    assert set(expected_header).issubset(set(predicted_header))
+    ##########################################
+    # Test custom_plotter tool when the
+    # simulation results are available but
+    # the species is not available
+    ##########################################
+    prompt = """Make a custom plot showing the
+        concentration of the species `TP53` over
+        time. Do not show any other species."""
+    # Update state
+    app.update_state(config, {"llm_model": "gpt-4o-mini"}
+                    )
+    # Test the tool get_modelinfo
+    response = app.invoke(
+                        {"messages": [HumanMessage(content=prompt)]},
+                        config=config
+                    )
+    assistant_msg = response["messages"][-1].content
+    # print (response["messages"])
+    current_state = app.get_state(config)
+    # Get the messages from the current state
+    # and reverse the order
+    reversed_messages = current_state.values["messages"][::-1]
+    # Loop through the reversed messages until a
+    # ToolMessage is found.
+    predicted_artifact = []
+    for msg in reversed_messages:
+        if isinstance(msg, ToolMessage):
+            # Work on the message if it is a ToolMessage
+            # These may contain additional visuals that
+            # need to be displayed to the user.
+            if msg.name == "custom_plotter":
+                predicted_artifact = msg.artifact
+                break
+    # Check if the the predicted artifact is `None`
+    assert predicted_artifact is None

aiagents4pharma/talk2biomodels/tests/test_param_scan.py

@@ -0,0 +1,68 @@
+'''
+Test cases for Talk2Biomodels parameter scan tool.
+'''
+
+import pandas as pd
+from langchain_core.messages import HumanMessage, ToolMessage
+from ..agents.t2b_agent import get_app
+
+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 = 1234
+    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 serum 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. Assume
+            that the model is simulated for 2016 hours with an interval of 50."""
+    # 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 species or 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"))

aiagents4pharma/talk2biomodels/tests/test_search_models.py

@@ -0,0 +1,28 @@
+'''
+Test cases for Talk2Biomodels search models tool.
+'''
+
+from langchain_core.messages import HumanMessage
+from ..agents.t2b_agent import get_app
+
+def test_search_models_tool():
+    '''
+    Test the search_models tool.
+    '''
+    unique_id = 12345
+    app = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    # Update state
+    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    prompt = "Search for models on Crohn's disease."
+    # Test the tool get_modelinfo
+    response = app.invoke(
+                        {"messages": [HumanMessage(content=prompt)]},
+                        config=config
+                    )
+    assistant_msg = response["messages"][-1].content
+    # Check if the assistant message is a string
+    assert isinstance(assistant_msg, str)
+    # Check if the assistant message contains the
+    # biomodel id BIO0000000537
+    assert "BIOMD0000000537" in assistant_msg

aiagents4pharma/talk2biomodels/tests/test_simulate_model.py

@@ -0,0 +1,39 @@
+'''
+Test cases for Talk2Biomodels.
+'''
+
+from langchain_core.messages import HumanMessage
+from ..agents.t2b_agent import get_app
+
+def test_simulate_model_tool():
+    '''
+    Test the simulate_model tool when simulating
+    multiple models.
+    '''
+    unique_id = 123
+    app = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    # Upload a model to the state
+    app.update_state(config,
+        {"sbml_file_path": ["aiagents4pharma/talk2biomodels/tests/BIOMD0000000449_url.xml"]})
+    prompt = "Simulate model 64 and the uploaded model"
+    # Invoke the agent
+    app.invoke(
+        {"messages": [HumanMessage(content=prompt)]},
+        config=config
+    )
+    current_state = app.get_state(config)
+    dic_simulated_data = current_state.values["dic_simulated_data"]
+    # Check if the dic_simulated_data is a list
+    assert isinstance(dic_simulated_data, list)
+    # Check if the length of the dic_simulated_data is 2
+    assert len(dic_simulated_data) == 2
+    # Check if the source of the first model is 64
+    assert dic_simulated_data[0]['source'] == 64
+    # Check if the source of the second model is upload
+    assert dic_simulated_data[1]['source'] == "upload"
+    # Check if the data of the first model contains
+    assert '1,3-bisphosphoglycerate' in dic_simulated_data[0]['data']
+    # Check if the data of the second model contains
+    assert 'mTORC2' in dic_simulated_data[1]['data']

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/__init__.py

@@ -9,3 +9,4 @@ from . import get_modelinfo
 from . import parameter_scan
 from . import steady_state
 from . import load_biomodel
+from . import get_annotation

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