aiagents4pharma 1.9.0__py3-none-any.whl → 1.15.0__py3-none-any.whl

aiagents4pharma/talk2biomodels/tests/test_get_annotation.py

@@ -0,0 +1,171 @@
+'''
+Test cases for Talk2Biomodels get_annotation tool.
+'''
+
+import random
+import pytest
+from langchain_core.messages import HumanMessage, ToolMessage
+from ..agents.t2b_agent import get_app
+from ..tools.get_annotation import prepare_content_msg
+
+@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)
+    config = {"configurable": {"thread_id": unique_id}}
+    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_specific_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"
+
+    # Test with an invalid species name
+    app, config = make_graph
+    prompt = "Extract annotations of 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 'NADH' in msg.content:
+                #If artifact is None, it means no annotation was found
+                # (likely due to an invalid species).
+                #If artifact contains data, the tool successfully retrieved annotations.
+                test_condition = True
+                break
+    # assert test_condition
+    assert test_condition, "Expected rejection message for NADH but did not find it."
+
+    # 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"]
+    reversed_messages = current_state.values["messages"][::-1]
+    # Loop through the reversed messages until a
+    # ToolMessage is found.
+    artifact_was_none = 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":
+            # print (msg.artifact, msg.content)
+
+            if msg.artifact is True and 'IL7' in msg.content:
+                artifact_was_none = True
+                break
+    assert artifact_was_none
+
+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.
+
+    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,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_query_article.py

@@ -0,0 +1,76 @@
+'''
+Test cases for Talk2Biomodels query_article tool.
+'''
+
+from pydantic import BaseModel, Field
+from langchain_core.messages import HumanMessage, ToolMessage
+from langchain_openai import ChatOpenAI
+from ..agents.t2b_agent import get_app
+
+class Article(BaseModel):
+    '''
+    Article schema.
+    '''
+    title: str = Field(description="Title of the article.")
+
+def test_query_article_with_an_article():
+    '''
+    Test the query_article tool by providing an article.
+    '''
+    unique_id = 12345
+    app = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    # Update state by providing the pdf file name
+    app.update_state(config,
+      {"pdf_file_name": "aiagents4pharma/talk2biomodels/tests/article_on_model_537.pdf"})
+    prompt = "What is the title of the article?"
+    # Test the tool query_article
+    response = app.invoke(
+                        {"messages": [HumanMessage(content=prompt)]},
+                        config=config
+                    )
+    # Get the response from the tool
+    assistant_msg = response["messages"][-1].content
+    # Prepare a LLM that can be used as a judge
+    llm = ChatOpenAI(model='gpt-4o-mini', temperature=0)
+    # Make it return a structured output
+    structured_llm = llm.with_structured_output(Article)
+    # Prepare a prompt for the judge
+    prompt = "Given the text below, what is the title of the article?"
+    prompt += f"\n\n{assistant_msg}"
+    # Get the structured output
+    article = structured_llm.invoke(prompt)
+    # Check if the article title is correct
+    expected_title = "A Multiscale Model of IL-6–Mediated "
+    expected_title += "Immune Regulation in Crohn’s Disease"
+    # Check if the article title is correct
+    assert article.title == expected_title
+
+def test_query_article_without_an_article():
+    '''
+    Test the query_article tool without providing an article.
+    The status of the tool should be error.
+    '''
+    unique_id = 12345
+    app = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    prompt = "What is the title of the uploaded article?"
+    # Test the tool query_article
+    app.invoke(
+            {"messages": [HumanMessage(content=prompt)]},
+            config=config
+        )
+    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.
+    tool_status_is_error = False
+    for msg in reversed_messages:
+        if isinstance(msg, ToolMessage):
+            # Skip until it finds a ToolMessage
+            if msg.name == "query_article" and msg.status == "error":
+                tool_status_is_error = True
+                break
+    assert tool_status_is_error

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/tests/test_sys_bio_model.py

@@ -16,6 +16,8 @@ class TestBioModel(SysBioModel):
     sbml_file_path: Optional[str] = Field(None, description="Path to an SBML file")
     name: Optional[str] = Field(..., description="Name of the model")
     description: Optional[str] = Field("", description="Description of the model")
+    param1: Optional[float] = Field(0.0, description="Parameter 1")
+    param2: Optional[float] = Field(0.0, description="Parameter 2")
 
     def get_model_metadata(self) -> Dict[str, Union[str, int]]:
         '''
@@ -23,15 +25,18 @@ class TestBioModel(SysBioModel):
         '''
         return self.biomodel_id
 
-    def simulate(self,
-                 parameters: Dict[str, Union[float, int]],
-                 duration: Union[int, float]) -> List[float]:
+    def update_parameters(self, parameters):
+        '''
+        Update the model parameters.
+        '''
+        self.param1 = parameters.get('param1', 0.0)
+        self.param2 = parameters.get('param2', 0.0)
+
+    def simulate(self, duration: Union[int, float]) -> List[float]:
         '''
         Simulate the model.
         '''
-        param1 = parameters.get('param1', 0.0)
-        param2 = parameters.get('param2', 0.0)
-        return [param1 + param2 * t for t in range(int(duration))]
+        return [self.param1 + self.param2 * t for t in range(int(duration))]
 
 def test_get_model_metadata():
     '''
@@ -53,5 +58,6 @@ def test_simulate():
     Test the simulate method of the BioModel class.
     '''
     model = TestBioModel(biomodel_id=123, name="Test Model", description="A test model")
-    results = model.simulate(parameters={'param1': 1.0, 'param2': 2.0}, duration=4.0)
+    model.update_parameters({'param1': 1.0, 'param2': 2.0})
+    results = model.simulate(duration=4.0)
     assert results == [1.0, 3.0, 5.0, 7.0]

aiagents4pharma/talk2biomodels/tools/__init__.py

@@ -6,4 +6,8 @@ from . import simulate_model
 from . import ask_question
 from . import custom_plotter
 from . import get_modelinfo
+from . import parameter_scan
+from . import steady_state
 from . import load_biomodel
+from . import get_annotation
+from . import query_article