aiagents4pharma 0.0.0__py3-none-any.whl

aiagents4pharma/talk2biomodels/tests/test_query_article.py

@@ -0,0 +1,184 @@
+"""
+Test cases for Talk2Biomodels query_article tool.
+"""
+
+from unittest.mock import MagicMock, patch
+
+from langchain_core.messages import HumanMessage, ToolMessage
+from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings
+from langchain_openai import ChatOpenAI
+from pydantic import BaseModel, Field
+
+from ..agents.t2b_agent import get_app
+from ..tools.query_article import QueryArticle
+
+LLM_MODEL = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
+
+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, llm_model=LLM_MODEL)
+    config = {"configurable": {"thread_id": unique_id}}
+    # Update state by providing the pdf file name
+    # and the text embedding model
+    app.update_state(
+        config,
+        {
+            "pdf_file_name": "aiagents4pharma/talk2biomodels/tests/article_on_model_537.pdf",
+            "text_embedding_model": NVIDIAEmbeddings(model="nvidia/llama-3.2-nv-embedqa-1b-v2"),
+        },
+    )
+    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 = LLM_MODEL
+    # 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 article title contains key terms or reports access failure
+    keywords = ["Multiscale", "IL-6", "Immune", "Crohn"]
+    msg_lower = assistant_msg.lower()
+
+    # Count keyword matches and check for access failure
+    title_matches = sum(1 for kw in keywords if kw.lower() in article.title.lower())
+    msg_matches = sum(1 for kw in keywords if kw.lower() in msg_lower)
+    access_failed = any(
+        ind in msg_lower
+        for ind in [
+            "unable to access",
+            "cannot access",
+            "assistance with",
+            "request for assistance",
+        ]
+    )
+
+    # Test passes if keywords found OR system reports access failure
+    expected = "A Multiscale Model of IL-6–Mediated Immune Regulation in Crohn's Disease"
+    assert title_matches >= 2 or msg_matches >= 2 or access_failed, (
+        f"Expected key terms from '{expected}' or access failure, "
+        f"got title: '{article.title}' and message: '{assistant_msg}'"
+    )
+
+
+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, llm_model=LLM_MODEL)
+    config = {"configurable": {"thread_id": unique_id}}
+    prompt = "What is the title of the uploaded article?"
+    # Update state by providing the text embedding model
+    app.update_state(
+        config,
+        {"text_embedding_model": NVIDIAEmbeddings(model="nvidia/llama-3.2-nv-embedqa-1b-v2")},
+    )
+    # 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
+
+
+@patch("aiagents4pharma.talk2biomodels.tools.query_article.PyPDFLoader")
+@patch("aiagents4pharma.talk2biomodels.tools.query_article.InMemoryVectorStore")
+def test_query_article_similarity_search_and_return(mock_vector_store, mock_pdf_loader):
+    """
+    Test that lines 62-64 are covered: similarity search and return join operation.
+    """
+    # Mock PDF loader
+    mock_page = MagicMock()
+    mock_page.page_content = "Sample article content about research methodology"
+    mock_loader_instance = MagicMock()
+    mock_loader_instance.lazy_load.return_value = [mock_page]
+    mock_pdf_loader.return_value = mock_loader_instance
+
+    # Mock vector store and similarity search
+    mock_doc1 = MagicMock()
+    mock_doc1.page_content = "First relevant document content"
+    mock_doc2 = MagicMock()
+    mock_doc2.page_content = "Second relevant document content"
+    mock_vector_store_instance = MagicMock()
+    mock_vector_store_instance.similarity_search.return_value = [mock_doc1, mock_doc2]
+    mock_vector_store.from_documents.return_value = mock_vector_store_instance
+
+    # Create tool and run
+    tool = QueryArticle()
+    state = {
+        "pdf_file_name": "test_file.pdf",
+        "text_embedding_model": MagicMock(),
+    }
+
+    tool_input = {"question": "What is the methodology?", "state": state}
+    result = tool.invoke(tool_input)
+
+    # Verify similarity_search was called (line 62)
+    mock_vector_store_instance.similarity_search.assert_called_once_with("What is the methodology?")
+
+    # Verify return join operation (line 64)
+    expected_result = "First relevant document content\nSecond relevant document content"
+    assert result == expected_result
+    assert isinstance(result, str)
+    assert "\n" in result
+
+
+@patch("aiagents4pharma.talk2biomodels.tools.query_article.PyPDFLoader")
+@patch("aiagents4pharma.talk2biomodels.tools.query_article.InMemoryVectorStore")
+def test_query_article_empty_search_results(mock_vector_store, mock_pdf_loader):
+    """
+    Test edge case where similarity search returns empty results.
+    """
+    # Mock PDF loader
+    mock_page = MagicMock()
+    mock_page.page_content = "Sample content"
+    mock_loader_instance = MagicMock()
+    mock_loader_instance.lazy_load.return_value = [mock_page]
+    mock_pdf_loader.return_value = mock_loader_instance
+
+    # Mock vector store with empty search results
+    mock_vector_store_instance = MagicMock()
+    mock_vector_store_instance.similarity_search.return_value = []
+    mock_vector_store.from_documents.return_value = mock_vector_store_instance
+
+    tool = QueryArticle()
+    state = {
+        "pdf_file_name": "test_file.pdf",
+        "text_embedding_model": MagicMock(),
+    }
+
+    tool_input = {"question": "Nonexistent topic", "state": state}
+    result = tool.invoke(tool_input)
+
+    # Should return empty string when no documents found
+    assert result == ""
+    assert isinstance(result, str)

aiagents4pharma/talk2biomodels/tests/test_save_model.py

@@ -0,0 +1,47 @@
+"""
+Test cases for Talk2Biomodels.
+"""
+
+import tempfile
+
+from langchain_core.messages import HumanMessage
+from langchain_openai import ChatOpenAI
+
+from ..agents.t2b_agent import get_app
+
+LLM_MODEL = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
+
+def test_save_model_tool():
+    """
+    Test the save_model tool.
+    """
+    unique_id = 123
+    app = get_app(unique_id, llm_model=LLM_MODEL)
+    config = {"configurable": {"thread_id": unique_id}}
+    # Simulate a model
+    prompt = "Simulate model 64"
+    # Invoke the agent
+    app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
+    current_state = app.get_state(config)
+    assert current_state.values["model_as_string"][-1] is not None
+    # Save a model without simulating
+    prompt = "Save the model"
+    # Invoke the agent
+    app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
+    current_state = app.get_state(config)
+    assert current_state.values["model_as_string"][-1] is not None
+    # Create a temporary directory to save the model
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Save a model to the temporary directory
+        prompt = f"Simulate model 64 and save it model at {temp_dir}"
+        # Invoke the agent
+        app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
+        current_state = app.get_state(config)
+        assert current_state.values["model_as_string"][-1] is not None
+    # Simulate and save a model in non-existing path
+    prompt = "Simulate model 64 and then save the model at /xyz/"
+    # Invoke the agent
+    app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
+    current_state = app.get_state(config)
+    assert current_state.values["model_as_string"][-1] is not None

aiagents4pharma/talk2biomodels/tests/test_search_models.py

@@ -0,0 +1,35 @@
+"""
+Test cases for Talk2Biomodels search models 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_search_models_tool():
+    """
+    Test the search_models tool.
+    """
+    unique_id = 12345
+    app = get_app(unique_id, llm_model=LLM_MODEL)
+    config = {"configurable": {"thread_id": unique_id}}
+    prompt = "Search for models on Crohn's disease."
+    app.update_state(config, {"llm_model": LLM_MODEL})
+    # Test the tool get_modelinfo
+    response = app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
+    # Extract the assistant artifact which contains
+    # all the search results
+    found_model_537 = False
+    for msg in response["messages"]:
+        if isinstance(msg, ToolMessage) and msg.name == "search_models":
+            msg_artifact = msg.artifact
+            for model in msg_artifact["dic_data"]:
+                if model["id"] == "BIOMD0000000537":
+                    found_model_537 = True
+            break
+    # Check if the model BIOMD0000000537 is found
+    assert found_model_537

aiagents4pharma/talk2biomodels/tests/test_simulate_model.py

@@ -0,0 +1,44 @@
+"""
+Test cases for Talk2Biomodels.
+"""
+
+from langchain_core.messages import HumanMessage
+from langchain_openai import ChatOpenAI
+
+from ..agents.t2b_agent import get_app
+
+LLM_MODEL = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
+
+def test_simulate_model_tool():
+    """
+    Test the simulate_model tool when simulating
+    multiple models.
+    """
+    unique_id = 123
+    app = get_app(unique_id, llm_model=LLM_MODEL)
+    config = {"configurable": {"thread_id": unique_id}}
+    # 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"]
+    # Check if the model_as_string is not None
+    assert current_state.values["model_as_string"][-1] is not None

aiagents4pharma/talk2biomodels/tests/test_steady_state.py

@@ -0,0 +1,86 @@
+"""
+Test cases for Talk2Biomodels steady state 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_steady_state_tool():
+    """
+    Test the steady_state tool.
+    """
+    unique_id = 123
+    app = get_app(unique_id, llm_model=LLM_MODEL)
+    config = {"configurable": {"thread_id": unique_id}}
+    app.update_state(config, {"llm_model": LLM_MODEL})
+    #########################################################
+    # 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 = """Bring model 64 to a steady state. 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

@@ -0,0 +1,67 @@
+"""
+This file contains the unit tests for the BioModel class.
+"""
+
+import pytest
+from pydantic import Field
+
+from ..models.sys_bio_model import SysBioModel
+
+
+class TestBioModel(SysBioModel):
+    """
+    A test BioModel class for unit testing.
+    """
+
+    biomodel_id: int | None = Field(None, description="BioModel ID of the model")
+    sbml_file_path: str | None = Field(None, description="Path to an SBML file")
+    name: str | None = Field(..., description="Name of the model")
+    description: str | None = Field("", description="Description of the model")
+    param1: float | None = Field(0.0, description="Parameter 1")
+    param2: float | None = Field(0.0, description="Parameter 2")
+
+    def get_model_metadata(self) -> dict[str, str | int]:
+        """
+        Get the metadata of the model.
+        """
+        return self.biomodel_id
+
+    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: int | float) -> list[float]:
+        """
+        Simulate the model.
+        """
+        return [self.param1 + self.param2 * t for t in range(int(duration))]
+
+
+def test_get_model_metadata():
+    """
+    Test the get_model_metadata method of the BioModel class.
+    """
+    model = TestBioModel(biomodel_id=123, name="Test Model", description="A test model")
+    metadata = model.get_model_metadata()
+    assert metadata == 123
+
+
+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):
+        TestBioModel(name="Test Model", description="A test model")
+
+
+def test_simulate():
+    """
+    Test the simulate method of the BioModel class.
+    """
+    model = TestBioModel(biomodel_id=123, name="Test Model", description="A test model")
+    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

@@ -0,0 +1,17 @@
+"""
+This file is used to import all the modules in the package.
+"""
+
+from . import (
+    ask_question,
+    custom_plotter,
+    get_annotation,
+    get_modelinfo,
+    load_biomodel,
+    parameter_scan,
+    query_article,
+    save_model,
+    search_models,
+    simulate_model,
+    steady_state,
+)

aiagents4pharma/talk2biomodels/tools/ask_question.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python3
+
+"""
+Tool for asking a question about the simulation results.
+"""
+
+import logging
+from typing import Annotated, Literal
+
+import basico
+import hydra
+import pandas as pd
+from langchain_core.tools.base import BaseTool
+from langchain_experimental.agents import create_pandas_dataframe_agent
+from langgraph.prebuilt import InjectedState
+from pydantic import BaseModel, Field
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class AskQuestionInput(BaseModel):
+    """
+    Input schema for the AskQuestion tool.
+    """
+
+    question: str = Field(description="question about the simulation and steady state results")
+    experiment_name: str = Field(
+        description="""Name assigned to the simulation
+                                            or steady state analysis when the tool
+                                            simulate_model or steady_state is invoked."""
+    )
+    question_context: Literal["simulation", "steady_state"] = Field(
+        description="Context of the question"
+    )
+    state: Annotated[dict, InjectedState]
+
+
+# Note: It's important that every field has type hints.
+# BaseTool is a Pydantic class and not having type hints
+# can lead to unexpected behavior.
+class AskQuestionTool(BaseTool):
+    """
+    Tool for asking a question about the simulation or steady state results.
+    """
+
+    name: str = "ask_question"
+    description: str = """A tool to ask question about the
+                        simulation or steady state results."""
+    args_schema: type[BaseModel] = AskQuestionInput
+    return_direct: bool = False
+
+    def _run(
+        self,
+        question: str,
+        experiment_name: str,
+        question_context: Literal["simulation", "steady_state"],
+        state: Annotated[dict, InjectedState],
+    ) -> str:
+        """
+        Run the tool.
+
+        Args:
+            question (str): The question to ask about the simulation or steady state results.
+            state (dict): The state of the graph.
+            experiment_name (str): The name assigned to the simulation or steady state analysis.
+
+        Returns:
+            str: The answer to the question.
+        """
+        logger.log(
+            logging.INFO,
+            "Calling ask_question tool %s, %s, %s",
+            question,
+            question_context,
+            experiment_name,
+        )
+        # Load hydra configuration
+        with hydra.initialize(version_base=None, config_path="../configs"):
+            cfg = hydra.compose(config_name="config", overrides=["tools/ask_question=default"])
+            cfg = cfg.tools.ask_question
+        # Get the context of the question
+        # and based on the context, get the data
+        # and prompt content to ask the question
+        if question_context == "steady_state":
+            dic_context = state["dic_steady_state_data"]
+            prompt_content = cfg.steady_state_prompt
+        else:
+            dic_context = state["dic_simulated_data"]
+            prompt_content = cfg.simulation_prompt
+        # Extract the
+        dic_data = {}
+        for data in dic_context:
+            for key in data:
+                if key not in dic_data:
+                    dic_data[key] = []
+                dic_data[key] += [data[key]]
+        # Create a pandas dataframe of the data
+        df_data = pd.DataFrame.from_dict(dic_data)
+        # Extract the data for the experiment
+        # matching the experiment name
+        df = pd.DataFrame(df_data[df_data["name"] == experiment_name]["data"].iloc[0])
+        logger.log(logging.INFO, "Shape of the dataframe: %s", df.shape)
+        # # Extract the model units
+        # model_units = basico.model_info.get_model_units()
+        # Update the prompt content with the model units
+        prompt_content += "Following are the model units:\n"
+        prompt_content += f"{basico.model_info.get_model_units()}\n\n"
+        # Create a pandas dataframe agent
+        df_agent = create_pandas_dataframe_agent(
+            state["llm_model"],
+            allow_dangerous_code=True,
+            agent_type="tool-calling",
+            df=df,
+            max_iterations=5,
+            include_df_in_prompt=True,
+            number_of_head_rows=df.shape[0],
+            verbose=True,
+            prefix=prompt_content,
+        )
+        # Invoke the agent with the question
+        llm_result = df_agent.invoke(question, stream_mode=None)
+        # print (llm_result)
+        return llm_result["output"]