aiagents4pharma 0.0.0__py3-none-any.whl

aiagents4pharma/talk2biomodels/tests/__init__.py

@@ -0,0 +1,3 @@
+"""
+This module contains the test cases.
+"""

aiagents4pharma/talk2biomodels/tests/test_api.py

@@ -0,0 +1,31 @@
+"""
+Test cases for Talk2Biomodels.
+"""
+
+from ..api.ols import fetch_from_ols
+from ..api.uniprot import search_uniprot_labels
+
+
+def test_search_uniprot_labels():
+    """
+    Test the search_uniprot_labels function.
+    """
+    # "P61764" = Positive result, "P0000Q" = negative result
+    identifiers = ["P61764", "P0000Q"]
+    results = search_uniprot_labels(identifiers)
+    assert results["P61764"] == "Syntaxin-binding protein 1"
+    assert results["P0000Q"].startswith("Error: 400")
+
+
+def test_fetch_from_ols():
+    """
+    Test the fetch_from_ols function.
+    """
+    term_1 = "GO:0005886"  # Positive result
+    term_2 = "GO:ABC123"  # Negative result
+    label_1 = fetch_from_ols(term_1)
+    label_2 = fetch_from_ols(term_2)
+    assert isinstance(label_1, str), f"Expected string, got {type(label_1)}"
+    assert isinstance(label_2, str), f"Expected string, got {type(label_2)}"
+    assert label_1 == "plasma membrane"
+    assert label_2.startswith("Error: 404")

aiagents4pharma/talk2biomodels/tests/test_ask_question.py

@@ -0,0 +1,42 @@
+"""
+Test cases for Talk2Biomodels.
+"""
+
+from langchain_core.messages import HumanMessage, ToolMessage
+from langchain_openai import ChatOpenAI
+
+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=ChatOpenAI(model="gpt-4o-mini", temperature=0))
+    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
+    ##########################################
+    # 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_basico_model.py

@@ -0,0 +1,67 @@
+"""
+A test BasicoModel class for pytest unit testing.
+"""
+
+import basico
+import pandas as pd
+import pytest
+
+from ..models.basico_model import BasicoModel
+
+
+@pytest.fixture(name="model")
+def model_fixture():
+    """
+    A fixture for the BasicoModel class.
+    """
+    return BasicoModel(biomodel_id=64, species={"Pyruvate": 100}, duration=2, interval=2)
+
+
+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(duration=2, interval=2), pd.DataFrame)
+    # Pass a None value to the update_parameters method
+    # and it should not raise an error
+    model.update_parameters(parameters={None: None})
+    # check if the model description is updated
+    assert model.description == basico.biomodels.get_model_info(model.biomodel_id)["description"]
+    # check if an error is raised if an invalid species/parameter (`Pyruv`)
+    # is passed and it should raise a ValueError
+    with pytest.raises(ValueError):
+        model.update_parameters(parameters={"Pyruv": 0.5})
+
+
+def test_with_sbml_file():
+    """
+    Test initialization of BasicoModel with sbml_file_path.
+    """
+    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)
+
+
+def test_check_biomodel_id_or_sbml_file_path():
+    """
+    Test the check_biomodel_id_or_sbml_file_path method of the BioModel class.
+    """
+    with pytest.raises(ValueError):
+        BasicoModel(species={"Pyruvate": 100}, duration=2, interval=2)
+
+
+def test_get_model_metadata():
+    """
+    Test the get_model_metadata method of the BasicoModel class.
+    """
+    model = BasicoModel(biomodel_id=64)
+    metadata = model.get_model_metadata()
+    assert metadata["Model Type"] == "SBML Model (COPASI)"
+    assert metadata["Parameter Count"] == len(basico.get_parameters())

aiagents4pharma/talk2biomodels/tests/test_get_annotation.py

@@ -0,0 +1,190 @@
+"""
+Test cases for Talk2Biomodels get_annotation tool.
+"""
+
+import random
+
+import pytest
+from langchain_core.messages import HumanMessage, ToolMessage
+from langchain_openai import ChatOpenAI
+
+from ..agents.t2b_agent import get_app
+from ..tools.get_annotation import prepare_content_msg
+
+LLM_MODEL = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
+
+@pytest.fixture(name="make_graph")
+def make_graph_fixture():
+    """
+    Create an instance of the talk2biomodels agent.
+    """
+    unique_id = random.randint(1000, 9999)
+    graph = get_app(unique_id, llm_model=LLM_MODEL)
+    config = {"configurable": {"thread_id": unique_id}}
+    graph.update_state(config, {"llm_model": LLM_MODEL})
+    return graph, config
+
+
+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_valid_species_provided(make_graph):
+    """
+    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
+    prompt = "Extract annotations of species IL6 in model 537."
+    app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
+    current_state = app.get_state(config)
+    # print (current_state.values["dic_annotations_data"])
+    dic_annotations_data = current_state.values["dic_annotations_data"]
+
+    # The assert statement checks if IL6 is present in the returned annotations.
+    assert dic_annotations_data[0]["data"]["Species Name"][0] == "IL6"
+
+
+def test_invalid_species_provided(make_graph):
+    """
+    Test the tool by providing an invalid species name.
+    We are testing a condition where the user asks for annotations
+    of an invalid species in a specific model.
+    """
+    # Test with an invalid species name
+    app, config = make_graph
+    prompt = "Extract annotations of only species NADH in model 537."
+    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.
+
+    test_condition = False
+    for msg in reversed_messages:
+        # Assert that the one of the messages is a ToolMessage
+        # and its artifact is None.
+        if isinstance(msg, ToolMessage) and msg.name == "get_annotation":
+            # If a ToolMessage exists and artifact is None (meaning no valid annotation was found)
+            # and the rejected species (NADH) is mentioned, the test passes.
+            if msg.artifact is None and msg.status == "error":
+                # If artifact is None, it means no annotation was found
+                # (likely due to an invalid species).
+                test_condition = True
+                break
+    assert test_condition
+
+
+def test_invalid_and_valid_species_provided(make_graph):
+    """
+    Test the tool by providing an invalid species name and a valid species name.
+    We are testing a condition where the user asks for annotations
+    of an invalid species and a valid species in a specific model.
+    """
+    # Test with an invalid species name and a valid species name
+    app, config = make_graph
+    prompt = "Extract annotations of species NADH, NAD, and IL7 in model 64."
+    app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
+    current_state = app.get_state(config)
+    dic_annotations_data = current_state.values["dic_annotations_data"]
+    # List of species that are expected to be found in the annotations
+    extracted_species = []
+    for idx in dic_annotations_data[0]["data"]["Species Name"]:
+        extracted_species.append(dic_annotations_data[0]["data"]["Species Name"][idx])
+    reversed_messages = current_state.values["messages"][::-1]
+    # Loop through the reversed messages until a
+    # ToolMessage is found.
+    tool_status_success = False
+    for msg in reversed_messages:
+        # Assert that the one of the messages is a ToolMessage
+        # and its artifact is None.
+        if isinstance(msg, ToolMessage) and msg.name == "get_annotation":
+            if msg.artifact is True and msg.status == "success":
+                tool_status_success = True
+                break
+    assert tool_status_success
+    assert set(extracted_species) == {"NADH", "NAD"}
+
+
+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 UniProt, and OLS.
+
+    We are testing a condition where the user asks for annotations
+    of all species in a specific model.
+    """
+    # 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}."
+        # Test the tool get_modelinfo
+        app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
+        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:
+            # 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]
+
+                # 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,92 @@
+"""
+Test cases for Talk2Biomodels get_modelinfo tool.
+"""
+
+from langchain_core.messages import HumanMessage, ToolMessage
+from langchain_openai import ChatOpenAI
+
+from ..agents.t2b_agent import get_app
+
+LLM_MODEL = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
+
+def test_get_modelinfo_tool():
+    """
+    Test the get_modelinfo tool.
+    """
+    unique_id = 12345
+    app = get_app(unique_id, LLM_MODEL)
+    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)
+
+
+def test_model_with_no_species():
+    """
+    Test the get_modelinfo tool with a model that does not
+    return any species.
+
+    This should raise a tool error.
+    """
+    unique_id = 12345
+    app = get_app(unique_id, LLM_MODEL)
+    config = {"configurable": {"thread_id": unique_id}}
+    prompt = "Extract all species from model 20"
+    # Test the tool get_modelinfo
+    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.
+    test_condition = False
+    for msg in reversed_messages:
+        # Check if the message is a ToolMessage from the get_modelinfo tool
+        if isinstance(msg, ToolMessage) and msg.name == "get_modelinfo":
+            # Check if the message is an error message
+            if (
+                msg.status == "error"
+                and "ValueError('Unable to extract species from the model.')" in msg.content
+            ):
+                test_condition = True
+                break
+    assert test_condition
+
+
+def test_model_with_no_parameters():
+    """
+    Test the get_modelinfo tool with a model that does not
+    return any parameters.
+
+    This should raise a tool error.
+    """
+    unique_id = 12345
+    app = get_app(unique_id, LLM_MODEL)
+    config = {"configurable": {"thread_id": unique_id}}
+    prompt = "Extract all parameters from model 10"
+    # Test the tool get_modelinfo
+    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.
+    test_condition = False
+    for msg in reversed_messages:
+        # Check if the message is a ToolMessage from the get_modelinfo tool
+        if isinstance(msg, ToolMessage) and msg.name == "get_modelinfo":
+            # Check if the message is an error message
+            if (
+                msg.status == "error"
+                and "ValueError('Unable to extract parameters from the model.')" in msg.content
+            ):
+                test_condition = True
+                break
+    assert test_condition

aiagents4pharma/talk2biomodels/tests/test_integration.py

@@ -0,0 +1,116 @@
+"""
+Test cases for Talk2Biomodels.
+"""
+
+import pandas as pd
+from langchain_core.messages import HumanMessage, ToolMessage
+from langchain_openai import ChatOpenAI
+
+from ..agents.t2b_agent import get_app
+
+LLM_MODEL = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
+
+def test_integration():
+    """
+    Test the integration of the tools.
+    """
+    unique_id = 1234567
+    app = get_app(unique_id, llm_model=LLM_MODEL)
+    config = {"configurable": {"thread_id": unique_id}}
+    # ##########################################
+    # ## Test simulate_model tool
+    # ##########################################
+    prompt = """Simulate the model BIOMD0000000537 for 100 hours and time 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": LLM_MODEL})
+    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 the custom_plotter tool when the
+    # simulation results are available but
+    # the species is not available
+    ##########################################
+    prompt = """Call the custom_plotter tool to make a plot
+        showing only species 'Infected cases'. Let me
+        know if these species were not found. Do not
+        invoke any other tool."""
+    # Update state
+    app.update_state(config, {"llm_model": LLM_MODEL})
+    # 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.
+    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
+
+    ##########################################
+    # Test custom_plotter tool when the
+    # simulation results are available
+    ##########################################
+    prompt = "Plot only CRP related species."
+
+    # Update state
+    app.update_state(config, {"llm_model": LLM_MODEL})
+    # 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["dic_data"]
+                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))

aiagents4pharma/talk2biomodels/tests/test_load_biomodel.py

@@ -0,0 +1,35 @@
+"""
+Test cases for Talk2Biomodels.
+"""
+
+import pytest
+
+from ..tools.load_biomodel import ModelData
+
+
+def test_model_data_valid_biomodel_id():
+    """
+    Test the ModelData class with valid
+    biomodel
+    """
+    # Test with string biomodel_id starting with 'BIOMD'
+    model_data = ModelData(biomodel_id="BIOMD0000000537")
+    assert model_data.biomodel_id == "BIOMD0000000537"
+
+    # Test with string biomodel_id starting with 'MODEL'
+    model_data = ModelData(biomodel_id="MODEL0000000537")
+    assert model_data.biomodel_id == "MODEL0000000537"
+
+
+def test_model_data_invalid_biomodel_id():
+    """
+    Test the ModelData class with invalid
+    biomodel
+    """
+    # Test with invalid string biomodel_id
+    with pytest.raises(ValueError):
+        ModelData(biomodel_id="12345")
+
+    # Test with float biomodel_id
+    with pytest.raises(ValueError):
+        ModelData(biomodel_id=123.45)

aiagents4pharma/talk2biomodels/tests/test_param_scan.py

@@ -0,0 +1,71 @@
+"""
+Test cases for Talk2Biomodels parameter scan tool.
+"""
+
+import pandas as pd
+from langchain_core.messages import HumanMessage, ToolMessage
+from langchain_openai import ChatOpenAI
+
+from ..agents.t2b_agent import get_app
+
+LLM_MODEL = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
+
+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, llm_model=LLM_MODEL)
+    config = {"configurable": {"thread_id": unique_id}}
+    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"))