aiagents4pharma 1.8.1__py3-none-any.whl → 1.8.3__py3-none-any.whl

aiagents4pharma/__init__.py

@@ -5,3 +5,4 @@ This file is used to import aiagents4pharma modules.
 from . import talk2biomodels
 from . import talk2cells
 from . import talk2knowledgegraphs
+from . import configs

aiagents4pharma/configs/__init__.py

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

aiagents4pharma/configs/config.yaml

@@ -0,0 +1,3 @@
+defaults:
+  - _self_
+  - talk2biomodels/agents/t2b_agent: default

aiagents4pharma/configs/talk2biomodels/__init__.py

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

aiagents4pharma/configs/talk2biomodels/agents/__init__.py

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

aiagents4pharma/configs/talk2biomodels/agents/t2b_agent/__init__.py

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

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

@@ -0,0 +1,6 @@
+_target_: talk2biomodels.agents.t2b_agent.get_app
+state_modifier: >
+  You are Talk2BioModels agent.
+  If the user asks for the uploaded model, 
+  then pass the use_uploaded_model argument
+    as True.

aiagents4pharma/talk2biomodels/agents/__init__.py

@@ -0,0 +1,5 @@
+'''
+This file is used to import all the modules in the package.
+'''
+# import everything from the module
+from . import t2b_agent

aiagents4pharma/talk2biomodels/agents/t2b_agent.py

@@ -0,0 +1,93 @@
+#/usr/bin/env python3
+
+'''
+This is the agent file for the Talk2BioModels agent.
+'''
+
+import logging
+from typing import Annotated
+import hydra
+from langchain_openai import ChatOpenAI
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import START, StateGraph
+from langgraph.prebuilt import create_react_agent, ToolNode, InjectedState
+from ..tools.search_models import SearchModelsTool
+from ..tools.get_modelinfo import GetModelInfoTool
+from ..tools.simulate_model import SimulateModelTool
+from ..tools.custom_plotter import CustomPlotterTool
+from ..tools.ask_question import AskQuestionTool
+from ..states.state_talk2biomodels import Talk2Biomodels
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def get_app(uniq_id, llm_model='gpt-4o-mini'):
+    '''
+    This function returns the langraph app.
+    '''
+    def agent_t2b_node(state: Annotated[dict, InjectedState]):
+        '''
+        This function calls the model.
+        '''
+        logger.log(logging.INFO, "Calling t2b_agent node with thread_id %s", uniq_id)
+        response = model.invoke(state, {"configurable": {"thread_id": uniq_id}})
+        return response
+
+    # Define the tools
+    simulate_model = SimulateModelTool()
+    custom_plotter = CustomPlotterTool()
+    ask_question = AskQuestionTool()
+    search_model = SearchModelsTool()
+    get_modelinfo = GetModelInfoTool()
+    tools = ToolNode([
+                    simulate_model,
+                    ask_question,
+                    custom_plotter,
+                    search_model,
+                    get_modelinfo
+                    ])
+
+    # Define the model
+    llm = ChatOpenAI(model=llm_model, temperature=0)
+    # Load hydra configuration
+    logger.log(logging.INFO, "Load Hydra configuration for Talk2BioModels agent.")
+    with hydra.initialize(version_base=None, config_path="../../configs"):
+        cfg = hydra.compose(config_name='config',
+                            overrides=['talk2biomodels/agents/t2b_agent=default'])
+        cfg = cfg.talk2biomodels.agents.t2b_agent
+    logger.log(logging.INFO, "state_modifier: %s", cfg.state_modifier)
+    # Create the agent
+    model = create_react_agent(
+                llm,
+                tools=tools,
+                state_schema=Talk2Biomodels,
+                state_modifier=cfg.state_modifier,
+                checkpointer=MemorySaver()
+            )
+
+    # Define a new graph
+    workflow = StateGraph(Talk2Biomodels)
+
+    # Define the two nodes we will cycle between
+    workflow.add_node("agent_t2b", agent_t2b_node)
+
+    # Set the entrypoint as the first node
+    # This means that this node is the first one called
+    workflow.add_edge(START, "agent_t2b")
+
+    # Initialize memory to persist state between graph runs
+    checkpointer = MemorySaver()
+
+    # Finally, we compile it!
+    # This compiles it into a LangChain Runnable,
+    # meaning you can use it as you would any other runnable.
+    # Note that we're (optionally) passing the memory
+    # when compiling the graph
+    app = workflow.compile(checkpointer=checkpointer)
+    logger.log(logging.INFO,
+               "Compiled the graph with thread_id %s and llm_model %s",
+               uniq_id,
+               llm_model)
+
+    return app

aiagents4pharma/talk2biomodels/states/__init__.py

@@ -0,0 +1,5 @@
+'''
+This file is used to import all the modules in the package.
+'''
+# import everything from the module
+from . import state_talk2biomodels

aiagents4pharma/talk2biomodels/states/state_talk2biomodels.py

@@ -0,0 +1,24 @@
+#!/usr/bin/env python3
+
+'''
+This is the state file for the Talk2BioModels agent.
+'''
+
+from typing import Annotated
+import operator
+from langgraph.prebuilt.chat_agent_executor import AgentState
+
+class Talk2Biomodels(AgentState):
+    """
+    The state for the Talk2BioModels agent.
+    """
+    model_id: Annotated[list, operator.add]
+    # sbml_file_path: str
+    # A StateGraph may receive a concurrent updates
+    # which is not supported by the StateGraph.
+    # Therefore, we need to use Annotated to specify
+    # the operator for the sbml_file_path field.
+    # https://langchain-ai.github.io/langgraph/troubleshooting/errors/INVALID_CONCURRENT_GRAPH_UPDATE/
+    sbml_file_path: Annotated[list, operator.add]
+    dic_simulated_data: dict
+    llm_model: str

aiagents4pharma/talk2biomodels/tests/__init__.py

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

aiagents4pharma/talk2biomodels/tests/test_basico_model.py

@@ -0,0 +1,55 @@
+'''
+A test BasicoModel class for pytest unit testing.
+'''
+
+import pandas as pd
+import pytest
+import basico
+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
+    # check if the simulation results are a pandas DataFrame object
+    assert isinstance(model.simulate(parameters={'Pyruvate': 0.5, 'KmPFKF6P': 1.5},
+                                     duration=2,
+                                     interval=2),
+                    pd.DataFrame)
+    assert isinstance(model.simulate(parameters={None: None}, duration=2, interval=2),
+                    pd.DataFrame)
+    assert model.description == basico.biomodels.get_model_info(model.biomodel_id)["description"]
+
+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)
+    assert isinstance(model_object.simulate(parameters={'NADH': 0.5}, 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_langgraph.py

@@ -0,0 +1,189 @@
+'''
+Test cases
+'''
+
+from langchain_core.messages import HumanMessage, ToolMessage
+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": ["BIOMD0000000537.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_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
+
+def test_ask_question_tool():
+    '''
+    Test the ask_question tool without the simulation results.
+    '''
+    unique_id = 12345
+    app = get_app(unique_id, llm_model='gpt-4o-mini')
+    config = {"configurable": {"thread_id": unique_id}}
+
+    ##########################################
+    # Test ask_question tool when simulation
+    # results are not available
+    ##########################################
+    # Update state
+    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    prompt = "Call the ask_question tool to answer the "
+    prompt += "question: What is the concentration of CRP "
+    prompt += "in serum at 1000 hours?"
+
+    # 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_simulate_model_tool():
+    '''
+    Test the simulate_model tool.
+    '''
+    unique_id = 123
+    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 2016 hours and intervals"
+    prompt += " 2016 with an initial concentration of `DoseQ2W` "
+    prompt += "set to 300 and `Dose` set to 0. Reset the concentration"
+    prompt += " of `NAD` to 100 every 500 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 at 1000 hours? "
+    # prompt += "Show only the concentration, rounded to 1 decimal place."
+    # prompt += "For example, if the concentration is 0.123456, "
+    # prompt += "your response should be `0.1`. Do not return any other information."
+    # 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 "1.7" 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_artifact = ['CRP[serum]', 'CRPExtracellular']
+    expected_artifact += ['CRP Suppression (%)', 'CRP (% of baseline)']
+    expected_artifact += ['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
+    # Check if the two artifacts are equal
+    # assert expected_artifact in predicted_artifact
+    assert set(expected_artifact).issubset(set(predicted_artifact))
+    ##########################################
+    # Test custom_plotter tool when the
+    # simulation results are available but
+    # the species is not available
+    ##########################################
+    prompt = "Plot the species `TP53`."
+
+    # 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_sys_bio_model.py

@@ -0,0 +1,57 @@
+'''
+This file contains the unit tests for the BioModel class.
+'''
+
+from typing import List, Dict, Union, Optional
+from pydantic import Field
+import pytest
+from ..models.sys_bio_model import SysBioModel
+
+class TestBioModel(SysBioModel):
+    '''
+    A test BioModel class for unit testing.
+    '''
+
+    biomodel_id: Optional[int] = Field(None, description="BioModel ID of the model")
+    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")
+
+    def get_model_metadata(self) -> Dict[str, Union[str, int]]:
+        '''
+        Get the metadata of the model.
+        '''
+        return self.biomodel_id
+
+    def simulate(self,
+                 parameters: Dict[str, Union[float, int]],
+                 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))]
+
+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")
+    results = model.simulate(parameters={'param1': 1.0, 'param2': 2.0}, duration=4.0)
+    assert results == [1.0, 3.0, 5.0, 7.0]

aiagents4pharma/talk2cells/tests/scp_agent/test_scp_agent.py

@@ -0,0 +1,23 @@
+'''
+Test cases for the search_studies
+'''
+
+# from ..tools.search_studies import search_studies
+from aiagents4pharma.talk2cells.agents.scp_agent import get_app
+from langchain_core.messages import HumanMessage
+
+def test_agent_scp():
+    '''
+    Test the agent_scp.
+    '''
+    unique_id = 12345
+    app = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    prompt = "Search for studies on Crohns Disease."
+    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/talk2cells/tools/scp_agent/__init__.py

@@ -0,0 +1,6 @@
+'''
+This file is used to import all the modules in the package.
+'''
+# import everything from the module
+from . import display_studies
+from . import search_studies

aiagents4pharma/talk2cells/tools/scp_agent/display_studies.py

@@ -0,0 +1,25 @@
+#!/usr/bin/env python3
+
+'''
+This tool is used to display the table of studies.
+'''
+
+import logging
+from typing import Annotated
+from langchain_core.tools import tool
+from langgraph.prebuilt import InjectedState
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+@tool('display_studies')
+def display_studies(state: Annotated[dict, InjectedState]):
+    """
+    Display the table of studies.
+
+    Args:
+        state (dict): The state of the agent.
+    """
+    logger.log(logging.INFO, "Calling the tool display_studies")
+    return state["search_table"]

aiagents4pharma/talk2cells/tools/scp_agent/search_studies.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+
+'''
+A tool to fetch studies from the Single Cell Portal.
+'''
+
+import logging
+from typing import Annotated
+import requests
+from langchain_core.tools import tool
+from langchain_core.tools.base import InjectedToolCallId
+from langchain_core.messages import ToolMessage
+from langgraph.types import Command
+import pandas as pd
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+@tool('search_studies')
+def search_studies(search_term: str,
+                   tool_call_id: Annotated[str, InjectedToolCallId],
+                   limit: int = 5):
+    """
+    Fetch studies from single cell portal
+
+    Args:
+        search_term (str): The search term to use. Example: "COVID-19", "cancer", etc.
+        limit (int): The number of papers to return. Default is 5.
+
+    """
+    logger.log(logging.INFO, "Calling the tool search_studies")
+    scp_endpoint = 'https://singlecell.broadinstitute.org/single_cell/api/v1/search?type=study'
+    # params = {'terms': search_term, 'facets': 'MONDO_0005011'}
+    params = {'terms': search_term}
+    status_code = 0
+    while status_code != 200:
+        # Make a GET request to the single cell portal
+        search_response = requests.get(scp_endpoint,
+                                       params=params,
+                                       timeout=10,
+                                       verify=False)
+        status_code = search_response.status_code
+        logger.log(logging.INFO, "Status code %s received from SCP")
+
+    # Select the columns to display in the table
+    selected_columns = ["study_source", "name", "study_url", "gene_count", "cell_count"]
+
+    # Extract the data from the response
+    # with the selected columns
+    df = pd.DataFrame(search_response.json()['studies'])[selected_columns]
+
+    # Convert column 'Study Name' into clickable
+    # hyperlinks from the column 'Study URL'
+    scp_api_url = 'https://singlecell.broadinstitute.org'
+    df['name'] = df.apply(
+            lambda x: f"<a href=\"{scp_api_url}/{x['study_url']}\">{x['name']}</a>",
+            axis=1)
+
+    # Excldue the column 'Study URL' from the dataframe
+    df = df.drop(columns=['study_url'])
+
+    # Add a new column a the beginning of the dataframe with row numbers
+    df.insert(0, 'S/N', range(1, 1 + len(df)))
+
+    # Update the state key 'search_table' with the dataframe in markdown format
+    return Command(
+        update={
+            # update the state keys
+            "search_table": df.to_markdown(tablefmt="grid"),
+            # update the message history
+            "messages": [
+                ToolMessage(
+                    f"Successfully fetched {limit} studies on {search_term}.",
+                    tool_call_id=tool_call_id
+                )
+            ],
+        }
+    )