aiagents4pharma 1.44.0__py3-none-any.whl → 1.45.1__py3-none-any.whl

aiagents4pharma/talk2biomodels/tests/__init__.py

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

aiagents4pharma/talk2biomodels/tests/test_api.py

@@ -1,27 +1,29 @@
-'''
+"""
 Test cases for Talk2Biomodels.
-'''
+"""
 
-from ..api.uniprot import search_uniprot_labels
+from ..api.kegg import fetch_from_api, fetch_kegg_names
 from ..api.ols import fetch_from_ols
-from ..api.kegg import fetch_kegg_names, fetch_from_api
+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
+    """
+    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)}"
@@ -29,10 +31,11 @@ def test_fetch_from_ols():
     assert label_1 == "plasma membrane"
     assert label_2.startswith("Error: 404")
 
+
 def test_fetch_kegg_names():
-    '''
+    """
     Test the fetch_kegg_names function.
-    '''
+    """
     ids = ["C00001", "C00002"]
     results = fetch_kegg_names(ids)
     assert results["C00001"] == "H2O"
@@ -42,10 +45,11 @@ def test_fetch_kegg_names():
     results = fetch_kegg_names([])
     assert not results
 
+
 def test_fetch_from_api():
-    '''
+    """
     Test the fetch_from_api function.
-    '''
+    """
     base_url = "https://rest.kegg.jp/get/"
     query = "C00001"
     entry_data = fetch_from_api(base_url, query)

aiagents4pharma/talk2biomodels/tests/test_ask_question.py

@@ -1,17 +1,19 @@
-'''
+"""
 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))
+    app = get_app(unique_id, llm_model=ChatOpenAI(model="gpt-4o-mini", temperature=0))
     config = {"configurable": {"thread_id": unique_id}}
 
     ##########################################
@@ -26,10 +28,7 @@ def test_ask_question_tool():
     prompt += "in serum at 1000 hours? The simulation name "
     prompt += "is `simulation_name`."
     # Invoke the tool
-    app.invoke(
-            {"messages": [HumanMessage(content=prompt)]},
-            config=config
-        )
+    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)

aiagents4pharma/talk2biomodels/tests/test_basico_model.py

@@ -1,12 +1,14 @@
-'''
+"""
 A test BasicoModel class for pytest unit testing.
-'''
+"""
 
+import basico
 import pandas as pd
 import pytest
-import basico
+
 from ..models.basico_model import BasicoModel
 
+
 @pytest.fixture(name="model")
 def model_fixture():
     """
@@ -14,16 +16,17 @@ def model_fixture():
     """
     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})
+    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
+    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
+    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
@@ -34,7 +37,8 @@ def test_with_biomodel_id(model):
     # 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})
+        model.update_parameters(parameters={"Pyruv": 0.5})
+
 
 def test_with_sbml_file():
     """
@@ -44,13 +48,15 @@ def test_with_sbml_file():
     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.

aiagents4pharma/talk2biomodels/tests/test_get_annotation.py

@@ -1,79 +1,74 @@
-'''
+"""
 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)
+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}
-        )
+    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
-                    )
+    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
-            )
+    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"
+    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
-    )
+    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
@@ -84,34 +79,32 @@ def test_invalid_species_provided(make_graph):
         # 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 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
+                # 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
-    )
+    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])
+    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.
@@ -124,10 +117,11 @@ def test_invalid_and_valid_species_provided(make_graph):
                 tool_status_success = True
                 break
     assert tool_status_success
-    assert set(extracted_species) == set(["NADH", "NAD"])
+    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.
@@ -136,16 +130,14 @@ def test_all_species_annotations(make_graph):
 
     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
-                        )
+        app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
         current_state = app.get_state(config)
 
         reversed_messages = current_state.values["messages"][::-1]
@@ -160,32 +152,39 @@ def test_all_species_annotations(make_graph):
                 # 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]
+                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 == '-')):
+                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"):
+                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"):
+                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
+        assert test_condition  # Expected output is validated

aiagents4pharma/talk2biomodels/tests/test_getmodelinfo.py

@@ -1,49 +1,48 @@
-'''
+"""
 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)
+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"]})
+    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
-                    )
+    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
-            )
+    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
@@ -53,28 +52,28 @@ def test_model_with_no_species():
         # 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):
+            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
-            )
+    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
@@ -84,8 +83,10 @@ def test_model_with_no_parameters():
         # 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):
+            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

@@ -1,34 +1,33 @@
-'''
+"""
 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)
+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
+    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.'''
+    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
-                    )
+    response = app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
     assistant_msg = response["messages"][-1].content
-    print (assistant_msg)
+    print(assistant_msg)
     # Check if the assistant message is a string
     assert isinstance(assistant_msg, str)
     ##########################################
@@ -40,14 +39,11 @@ def test_integration():
     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
-                    )
+    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
+    assert "211" in assistant_msg
 
     ##########################################
     # Test the custom_plotter tool when the
@@ -59,13 +55,9 @@ def test_integration():
         know if these species were not found. Do not
         invoke any other tool."""
     # Update state
-    app.update_state(config, {"llm_model": LLM_MODEL}
-                    )
+    app.update_state(config, {"llm_model": LLM_MODEL})
     # Test the tool get_modelinfo
-    response = app.invoke(
-                        {"messages": [HumanMessage(content=prompt)]},
-                        config=config
-                    )
+    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
@@ -92,13 +84,9 @@ def test_integration():
     prompt = "Plot only CRP related species."
 
     # Update state
-    app.update_state(config, {"llm_model": LLM_MODEL}
-                    )
+    app.update_state(config, {"llm_model": LLM_MODEL})
     # Test the tool get_modelinfo
-    response = app.invoke(
-                        {"messages": [HumanMessage(content=prompt)]},
-                        config=config
-                    )
+    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
@@ -106,9 +94,9 @@ def test_integration():
     reversed_messages = current_state.values["messages"][::-1]
     # Loop through the reversed messages
     # until a ToolMessage is found.
-    expected_header = ['Time', 'CRP{serum}', 'CRPExtracellular']
-    expected_header += ['CRP Suppression (%)', 'CRP (% of baseline)']
-    expected_header += ['CRP{liver}']
+    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):
@@ -116,7 +104,7 @@ def test_integration():
             # These may contain additional visuals that
             # need to be displayed to the user.
             if msg.name == "custom_plotter":
-                predicted_artifact = msg.artifact['dic_data']
+                predicted_artifact = msg.artifact["dic_data"]
                 break
     # Convert the artifact into a pandas dataframe
     # for easy comparison