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

aiagents4pharma/configs/config.yaml

@@ -1,3 +1,4 @@
 defaults:
   - _self_
-  - talk2biomodels/agents/t2b_agent: default
+  - talk2biomodels/agents/t2b_agent: default
+  - talk2biomodels/tools/ask_question: default

aiagents4pharma/configs/talk2biomodels/__init__.py

@@ -3,3 +3,4 @@ Import all the modules in the package
 '''
 
 from . import agents
+from . import tools

aiagents4pharma/configs/talk2biomodels/agents/t2b_agent/default.yaml

@@ -4,6 +4,5 @@ state_modifier: >
   If the user asks for the uploaded model, 
   then pass the use_uploaded_model argument
   as True. If the user asks for simulation
-  or steady state, suggest a value for the
-  `simulation_name` or `steadystate_name`
-  argument.
+  or param_scan or steady state, suggest a
+  value for the `experiment_name` argument.

aiagents4pharma/configs/talk2biomodels/tools/__init__.py

@@ -0,0 +1,4 @@
+'''
+Import all the modules in the package
+'''
+from . import ask_question

aiagents4pharma/configs/talk2biomodels/tools/ask_question/__init__.py

@@ -0,0 +1,3 @@
+'''
+Import all the modules in the package
+'''

aiagents4pharma/talk2biomodels/states/state_talk2biomodels.py

@@ -8,19 +8,33 @@ from typing import Annotated
 import operator
 from langgraph.prebuilt.chat_agent_executor import AgentState
 
+def add_data(data1: dict, data2: dict) -> dict:
+    """
+    A reducer function to merge two dictionaries.
+    """
+    left_idx_by_name = {data['name']: idx for idx, data in enumerate(data1)}
+    merged = data1.copy()
+    for data in data2:
+        idx = left_idx_by_name.get(data['name'])
+        if idx is not None:
+            merged[idx] = data
+        else:
+            merged.append(data)
+    return merged
+
 class Talk2Biomodels(AgentState):
     """
     The state for the Talk2BioModels agent.
     """
     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.
+    # which is not supported by the StateGraph. Hence,
+    # we need to add a reducer function to handle the
+    # concurrent updates.
     # 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]
-    dic_scanned_data: Annotated[list[dict], operator.add]
-    dic_steady_state_data: Annotated[list[dict], operator.add]
-    dic_annotations_data : Annotated[list[dict], operator.add]
+    dic_simulated_data: Annotated[list[dict], add_data]
+    dic_scanned_data: Annotated[list[dict], add_data]
+    dic_steady_state_data: Annotated[list[dict], add_data]
+    dic_annotations_data : Annotated[list[dict], add_data]

aiagents4pharma/talk2biomodels/tests/test_ask_question.py

@@ -0,0 +1,44 @@
+'''
+Test cases for Talk2Biomodels.
+'''
+
+from langchain_core.messages import HumanMessage, ToolMessage
+from ..agents.t2b_agent import get_app
+
+def test_ask_question_tool():
+    '''
+    Test the ask_question tool without the simulation results.
+    '''
+    unique_id = 12345
+    app = get_app(unique_id, llm_model='gpt-4o-mini')
+    config = {"configurable": {"thread_id": unique_id}}
+
+    ##########################################
+    # Test ask_question tool when simulation
+    # results are not available i.e. the
+    # simulation has not been run. In this
+    # case, the tool should return an error
+    ##########################################
+    # Update state
+    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    # Define the prompt
+    prompt = "Call the ask_question tool to answer the "
+    prompt += "question: What is the concentration of CRP "
+    prompt += "in serum at 1000 hours? The simulation name "
+    prompt += "is `simulation_name`."
+    # Invoke the tool
+    app.invoke(
+            {"messages": [HumanMessage(content=prompt)]},
+            config=config
+        )
+    # Get the messages from the current state
+    # and reverse the order
+    current_state = app.get_state(config)
+    reversed_messages = current_state.values["messages"][::-1]
+    # Loop through the reversed messages until a
+    # ToolMessage is found.
+    for msg in reversed_messages:
+        # Assert that the message is a ToolMessage
+        # and its status is "error"
+        if isinstance(msg, ToolMessage):
+            assert msg.status == "error"

aiagents4pharma/talk2biomodels/tests/test_get_annotation.py

@@ -1,6 +1,7 @@
 '''
-Test cases for Talk2Biomodels.
+Test cases for Talk2Biomodels get_annotation tool.
 '''
+
 import random
 import pytest
 from langchain_core.messages import HumanMessage, ToolMessage
@@ -17,9 +18,27 @@ def make_graph_fixture():
     config = {"configurable": {"thread_id": unique_id}}
     return graph, config
 
-def test_species_list(make_graph):
+def test_no_model_provided(make_graph):
+    '''
+    Test the tool by not specifying any model.
+    We are testing a condition where the user
+    asks for annotations of all species without
+    specifying a model.
+    '''
+    app, config = make_graph
+    prompt = "Extract annotations of all species. Call the tool get_annotation."
+    app.invoke({"messages": [HumanMessage(content=prompt)]},
+                        config=config
+                    )
+    current_state = app.get_state(config)
+    # Assert that the state key model_id is empty.
+    assert current_state.values["model_id"] == []
+
+def test_specific_species_provided(make_graph):
     '''
-    Test the tool by passing species names.
+    Test the tool by providing a specific species name.
+    We are testing a condition where the user asks for annotations
+    of a specific species in a specific model.
     '''
     # Test with a valid species name
     app, config = make_graph
@@ -87,15 +106,19 @@ def test_species_list(make_graph):
                 break
     assert artifact_was_none
 
-def test_all_species(make_graph):
+def test_all_species_annotations(make_graph):
     '''
     Test the tool by asking for annotations of all species is specific models.
+    Here, we test the tool with three models since they have different use cases:
+        - model 12 contains a species with no URL provided.
+        - model 20 contains a species without description.
+        - model 56 contains a species with database outside of KEGG, UniProt, and OLS.
 
-    model 12 contains species with no URL.
-    model 20 contains species without description.
-    model 56 contains species with database outside of KEGG, UniProt, and OLS.
+    We are testing a condition where the user asks for annotations
+    of all species in a specific model.
     '''
-    # Test valid models
+    # Loop through the models and test the tool
+    # for each model's unique use case.
     for model_id in [12, 20, 56]:
         app, config = make_graph
         prompt = f"Extract annotations of all species model {model_id}."
@@ -103,71 +126,46 @@ def test_all_species(make_graph):
         app.invoke({"messages": [HumanMessage(content=prompt)]},
                             config=config
                         )
-        #print(response["messages"])
-        # assistant_msg = response["messages"][-1].content
-
         current_state = app.get_state(config)
 
         reversed_messages = current_state.values["messages"][::-1]
         # Coveres all of the use cases for the expecetd sting on all the species
         test_condition = False
         for msg in reversed_messages:
-            if isinstance(msg, ToolMessage) and msg.name == "get_annotation":
-                if model_id == 12:
-                # For model 12:
-                # Expect a successful extraction (artifact is True) and that the content
-                # matches what is returned by prepare_content_msg for species ['lac'].
-                    if (msg.artifact is True and msg.content == prepare_content_msg(['lac'],[])
-                        and msg.status=="success"):
-                        test_condition = True
-                        break
-
-                if model_id == 20:
-                # For model 20:
-                # Expect an error message containing a note that species extraction failed.
-                    if ("Unable to extract species from the model"
-                        in msg.content and msg.status == "error"):
-                        test_condition = True
-                        break
-
-                if model_id == 56:
-                    # For model 56:
-                    # Expect a successful extraction (artifact is True) and that the content
-                    # matches for for missing description ['ORI'].
-                    if (msg.artifact is True and
-                    msg.content == prepare_content_msg([],['ORI'])
-                    and msg.status == "success"):
-                        test_condition = True
-                        break
-
-    # Retrieve the dictionary that holds all the annotation data from the app's state
-    dic_annotations_data = current_state.values["dic_annotations_data"]
-
-    assert isinstance(dic_annotations_data, list),\
-        f"Expected a list for model {model_id}, got {type(dic_annotations_data)}"
-    assert len(dic_annotations_data) > 0,\
-        f"Expected species data for model {model_id}, but got empty list"
-    assert test_condition # Expected output is validated
-
-    # Test case where no model is specified
-    app, config = make_graph
-    prompt = "Extract annotations of all species."
-    app.invoke({"messages": [HumanMessage(content=prompt)]},
-                        config=config
-                    )
-    current_state = app.get_state(config)
-    # dic_annotations_data = current_state.values["dic_annotations_data"]
-    reversed_messages = current_state.values["messages"][::-1]
-    print(reversed_messages)
+            # Skip messages that are not ToolMessages and those that are not
+            # from the get_annotation tool.
+            if not isinstance(msg, ToolMessage) or msg.name != "get_annotation":
+                continue
+            if model_id == 12:
+                # Extact the first and second description of the LacI protein
+                # We already know that the first or second description is missing ('-')
+                dic_annotations_data = current_state.values["dic_annotations_data"][0]
+                first_descp_laci_protein = dic_annotations_data['data']['Description'][0]
+                second_descp_laci_protein = dic_annotations_data['data']['Description'][1]
 
-    test_condition = False
-    for msg in reversed_messages:
-        # Assert that the one of the messages is a ToolMessage
-        if isinstance(msg, ToolMessage) and msg.name == "get_annotation":
-            if "Error:" in msg.content and msg.status == "error":
-                test_condition = True
-                break
-    # Loop through the reversed messages until a
-    # ToolMessage is found.
-    # Ensure the system correctly informs the user to specify a model
-    assert test_condition, "Expected error message when no model is specified was not found."
+                # Expect a successful extraction (artifact is True) and that the content
+                # matches what is returned by prepare_content_msg for species.
+                # And that the first or second description of the LacI protein is missing.
+                if (msg.artifact is True and msg.content == prepare_content_msg([],[])
+                    and msg.status=="success" and (first_descp_laci_protein == '-' or
+                                                    second_descp_laci_protein == '-')):
+                    test_condition = True
+                    break
+
+            if model_id == 20:
+                # Expect an error message containing a note
+                # that species extraction failed.
+                if ("Unable to extract species from the model"
+                    in msg.content and msg.status == "error"):
+                    test_condition = True
+                    break
+
+            if model_id == 56:
+                # Expect a successful extraction (artifact is True) and that the content
+                # matches for for missing description ['ORI'].
+                if (msg.artifact is True and
+                msg.content == prepare_content_msg([],['ORI'])
+                and msg.status == "success"):
+                    test_condition = True
+                    break
+        assert test_condition # Expected output is validated

aiagents4pharma/talk2biomodels/tests/test_getmodelinfo.py

@@ -0,0 +1,26 @@
+'''
+Test cases for Talk2Biomodels get_modelinfo tool.
+'''
+
+from langchain_core.messages import HumanMessage
+from ..agents.t2b_agent import get_app
+
+def test_get_modelinfo_tool():
+    '''
+    Test the get_modelinfo tool.
+    '''
+    unique_id = 12345
+    app = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    # Update state
+    app.update_state(config,
+      {"sbml_file_path": ["aiagents4pharma/talk2biomodels/tests/BIOMD0000000449_url.xml"]})
+    prompt = "Extract all relevant information from the uploaded model."
+    # 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)

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