aiagents4pharma 0.0.0__py3-none-any.whl

aiagents4pharma/talk2biomodels/tools/search_models.py

@@ -0,0 +1,96 @@
+#!/usr/bin/env python3
+
+"""
+Tool for searching models based on search query.
+"""
+
+import logging
+from typing import Annotated
+
+import pandas as pd
+from basico import biomodels
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import BaseTool
+from langchain_core.tools.base import InjectedToolCallId
+from langgraph.types import Command
+from pydantic import BaseModel, Field
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class SearchModelsInput(BaseModel):
+    """
+    Input schema for the search models tool.
+    """
+
+    query: str = Field(description="Search models query", default=None)
+    num_query: int = Field(description="Top number of models to search", default=10, le=100)
+    tool_call_id: Annotated[str, InjectedToolCallId]
+
+
+# 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 SearchModelsTool(BaseTool):
+    """
+    Tool for returning the search results based on the search query.
+    """
+
+    name: str = "search_models"
+    description: str = "Search for only manually curated models in "
+    "the BioMmodels database based on keywords."
+    args_schema: type[BaseModel] = SearchModelsInput
+    return_direct: bool = False
+
+    def _run(
+        self,
+        tool_call_id: Annotated[str, InjectedToolCallId],
+        query: str = None,
+        num_query: int = 10,
+    ) -> dict:
+        """
+        Run the tool.
+
+        Args:
+            query (str): The search query.
+            num_query (int): The number of models to search.
+            tool_call_id (str): The tool call ID.
+
+        Returns:
+            dict: The answer to the question in the form of a dictionary.
+        """
+        logger.log(
+            logging.INFO,
+            "Searching models with the query and number %s, %s",
+            query,
+            num_query,
+        )
+        # Search for models based on the query
+        search_results = biomodels.search_for_model(query, num_results=num_query)
+        # Convert the search results to a pandas DataFrame
+        df = pd.DataFrame(search_results)
+        # Prepare a message to return
+        first_n = min(3, len(search_results))
+        content = f"Found {len(search_results)} manually curated models"
+        content += f" for the query: {query}."
+        # Pass the first 3 models to the LLM
+        # to avoid hallucinations
+        content += f" Here is the summary of the first {first_n} models:"
+        for i in range(first_n):
+            content += (
+                f"\nModel {i + 1}: {search_results[i]['name']} (ID: {search_results[i]['id']})"
+            )
+        # Return the updated state of the tool
+        return Command(
+            update={
+                # update the message history
+                "messages": [
+                    ToolMessage(
+                        content=content,
+                        tool_call_id=tool_call_id,
+                        artifact={"dic_data": df.to_dict(orient="records")},
+                    )
+                ],
+            }
+        )

aiagents4pharma/talk2biomodels/tools/simulate_model.py

@@ -0,0 +1,137 @@
+#!/usr/bin/env python3
+
+"""
+Tool for simulating a model.
+"""
+
+import logging
+from typing import Annotated
+
+import basico
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import BaseTool
+from langchain_core.tools.base import InjectedToolCallId
+from langgraph.prebuilt import InjectedState
+from langgraph.types import Command
+from pydantic import BaseModel, Field
+
+from .load_arguments import ArgumentData, add_rec_events
+from .load_biomodel import ModelData, load_biomodel
+from .utils import get_model_units
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class SimulateModelInput(BaseModel):
+    """
+    Input schema for the SimulateModel tool.
+    """
+
+    sys_bio_model: ModelData = Field(description="model data", default=None)
+    arg_data: ArgumentData = Field(
+        description="""time, species, and reocurring data
+                                   as well as the simulation name""",
+        default=None,
+    )
+    tool_call_id: Annotated[str, InjectedToolCallId]
+    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 SimulateModelTool(BaseTool):
+    """
+    Tool for simulating a model.
+    """
+
+    name: str = "simulate_model"
+    description: str = "A tool to simulate a biomodel"
+    args_schema: type[BaseModel] = SimulateModelInput
+
+    def _run(
+        self,
+        tool_call_id: Annotated[str, InjectedToolCallId],
+        state: Annotated[dict, InjectedState],
+        sys_bio_model: ModelData = None,
+        arg_data: ArgumentData = None,
+    ) -> Command:
+        """
+        Run the tool.
+
+        Args:
+            tool_call_id (str): The tool call ID. This is injected by the system.
+            state (dict): The state of the tool.
+            sys_bio_model (ModelData): The model data.
+            arg_data (ArgumentData): The argument data.
+
+        Returns:
+            str: The result of the simulation.
+        """
+        logger.log(logging.INFO, "Calling simulate_model tool %s, %s", sys_bio_model, arg_data)
+        sbml_file_path = state["sbml_file_path"][-1] if len(state["sbml_file_path"]) > 0 else None
+        model_object = load_biomodel(sys_bio_model, sbml_file_path=sbml_file_path)
+        # Prepare the dictionary of species data
+        # that will be passed to the simulate method
+        # of the BasicoModel class
+        duration = 100.0
+        interval = 10
+        dic_species_to_be_analyzed_before_experiment = {}
+        if arg_data:
+            # Prepare the dictionary of species data
+            if arg_data.species_to_be_analyzed_before_experiment is not None:
+                dic_species_to_be_analyzed_before_experiment = dict(
+                    zip(
+                        arg_data.species_to_be_analyzed_before_experiment.species_name,
+                        arg_data.species_to_be_analyzed_before_experiment.species_concentration,
+                        strict=False,
+                    )
+                )
+            # Add reocurring events (if any) to the model
+            if arg_data.reocurring_data is not None:
+                add_rec_events(model_object, arg_data.reocurring_data)
+            # Set the duration and interval
+            if arg_data.time_data is not None:
+                duration = arg_data.time_data.duration
+                interval = arg_data.time_data.interval
+        # Update the model parameters
+        model_object.update_parameters(dic_species_to_be_analyzed_before_experiment)
+        logger.log(
+            logging.INFO,
+            "Following species/parameters updated in the model %s",
+            dic_species_to_be_analyzed_before_experiment,
+        )
+        # Simulate the model
+        df = model_object.simulate(duration=duration, interval=interval)
+        logger.log(logging.INFO, "Simulation results ready with shape %s", df.shape)
+        dic_simulated_data = {
+            "name": arg_data.experiment_name,
+            "source": (sys_bio_model.biomodel_id if sys_bio_model.biomodel_id else "upload"),
+            "tool_call_id": tool_call_id,
+            "data": df.to_dict(),
+        }
+        # Prepare the dictionary of updated state
+        dic_updated_state_for_model = {}
+        for key, value in {
+            "model_id": [sys_bio_model.biomodel_id],
+            "sbml_file_path": [sbml_file_path],
+            "dic_simulated_data": [dic_simulated_data],
+            "model_as_string": [basico.model_io.save_model_to_string()],
+        }.items():
+            if value:
+                dic_updated_state_for_model[key] = value
+        # Return the updated state of the tool
+        return Command(
+            update=dic_updated_state_for_model
+            | {
+                # update the message history
+                "messages": [
+                    ToolMessage(
+                        content=f"Simulation results of {arg_data.experiment_name}",
+                        tool_call_id=tool_call_id,
+                        artifact=get_model_units(model_object),
+                    )
+                ],
+            }
+        )

aiagents4pharma/talk2biomodels/tools/steady_state.py

@@ -0,0 +1,187 @@
+#!/usr/bin/env python3
+
+"""
+Tool for parameter scan.
+"""
+
+import logging
+from typing import Annotated
+
+import basico
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import BaseTool
+from langchain_core.tools.base import InjectedToolCallId
+from langgraph.prebuilt import InjectedState
+from langgraph.types import Command
+from pydantic import BaseModel, Field
+
+from .load_arguments import ArgumentData, add_rec_events
+from .load_biomodel import ModelData, load_biomodel
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def run_steady_state(model_object, dic_species_to_be_analyzed_before_experiment):
+    """
+    Run the steady state analysis.
+
+    Args:
+        model_object: The model object.
+        dic_species_to_be_analyzed_before_experiment: Dictionary of species data.
+
+    Returns:
+        DataFrame: The results of the steady state analysis.
+    """
+    # Update the fixed model species and parameters
+    # These are the initial conditions of the model
+    # set by the user
+    model_object.update_parameters(dic_species_to_be_analyzed_before_experiment)
+    logger.log(logging.INFO, "Running steady state analysis")
+    # Run the steady state analysis
+    output = basico.task_steadystate.run_steadystate(model=model_object.copasi_model)
+    if output == 0:
+        logger.error("Steady state analysis failed")
+        raise ValueError("A steady state was not found")
+    logger.log(logging.INFO, "Steady state analysis successful")
+    # Store the steady state results in a DataFrame
+    df_steady_state = basico.model_info.get_species(model=model_object.copasi_model).reset_index()
+    # print (df_steady_state)
+    # Rename the column name to species_name
+    df_steady_state.rename(columns={"name": "species_name"}, inplace=True)
+    # Rename the column concentration to steady_state_concentration
+    df_steady_state.rename(columns={"concentration": "steady_state_concentration"}, inplace=True)
+    # Rename the column transition_time to steady_state_transition_time
+    df_steady_state.rename(
+        columns={"transition_time": "steady_state_transition_time"}, inplace=True
+    )
+    # Drop some columns
+    df_steady_state.drop(
+        columns=[
+            "initial_particle_number",
+            "initial_expression",
+            "expression",
+            "particle_number",
+            "type",
+            "particle_number_rate",
+            "key",
+            "sbml_id",
+            "display_name",
+        ],
+        inplace=True,
+    )
+    logger.log(logging.INFO, "Steady state results with shape %s", df_steady_state.shape)
+    return df_steady_state
+
+
+class SteadyStateInput(BaseModel):
+    """
+    Input schema for the steady state tool.
+    """
+
+    sys_bio_model: ModelData = Field(description="model data", default=None)
+    arg_data: ArgumentData = Field(
+        description="time, species, and reocurring data"
+        " that must be set before the steady state analysis"
+        " as well as the experiment name",
+        default=None,
+    )
+    tool_call_id: Annotated[str, InjectedToolCallId]
+    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 SteadyStateTool(BaseTool):
+    """
+    Tool to bring a model to steady state.
+    """
+
+    name: str = "steady_state"
+    description: str = "A tool to bring a model to steady state."
+    args_schema: type[BaseModel] = SteadyStateInput
+
+    def _run(
+        self,
+        tool_call_id: Annotated[str, InjectedToolCallId],
+        state: Annotated[dict, InjectedState],
+        sys_bio_model: ModelData = None,
+        arg_data: ArgumentData = None,
+    ) -> Command:
+        """
+        Run the tool.
+
+        Args:
+            tool_call_id (str): The tool call ID. This is injected by the system.
+            state (dict): The state of the tool.
+            sys_bio_model (ModelData): The model data.
+            arg_data (ArgumentData): The argument data.
+
+        Returns:
+            Command: The updated state of the tool.
+        """
+        logger.log(
+            logging.INFO,
+            "Calling the steady_state tool %s, %s",
+            sys_bio_model,
+            arg_data,
+        )
+        # print (f'Calling steady_state tool {sys_bio_model}, {arg_data}, {tool_call_id}')
+        sbml_file_path = state["sbml_file_path"][-1] if len(state["sbml_file_path"]) > 0 else None
+        model_object = load_biomodel(sys_bio_model, sbml_file_path=sbml_file_path)
+        # Prepare the dictionary of species data
+        # that will be passed to the simulate method
+        # of the BasicoModel class
+        dic_species_to_be_analyzed_before_experiment = {}
+        if arg_data:
+            # Prepare the dictionary of species data
+            if arg_data.species_to_be_analyzed_before_experiment is not None:
+                dic_species_to_be_analyzed_before_experiment = dict(
+                    zip(
+                        arg_data.species_to_be_analyzed_before_experiment.species_name,
+                        arg_data.species_to_be_analyzed_before_experiment.species_concentration,
+                        strict=False,
+                    )
+                )
+            # Add reocurring events (if any) to the model
+            if arg_data.reocurring_data is not None:
+                add_rec_events(model_object, arg_data.reocurring_data)
+        # Run the parameter scan
+        df_steady_state = run_steady_state(
+            model_object, dic_species_to_be_analyzed_before_experiment
+        )
+        print(df_steady_state)
+        # Prepare the dictionary of scanned data
+        # that will be passed to the state of the graph
+        dic_steady_state_data = {
+            "name": arg_data.experiment_name,
+            "source": (sys_bio_model.biomodel_id if sys_bio_model.biomodel_id else "upload"),
+            "tool_call_id": tool_call_id,
+            "data": df_steady_state.to_dict(orient="records"),
+        }
+        # Prepare the dictionary of updated state for the model
+        dic_updated_state_for_model = {}
+        for key, value in {
+            "model_id": [sys_bio_model.biomodel_id],
+            "sbml_file_path": [sbml_file_path],
+            "dic_steady_state_data": [dic_steady_state_data],
+        }.items():
+            if value:
+                dic_updated_state_for_model[key] = value
+        # Return the updated state
+        return Command(
+            update=dic_updated_state_for_model
+            | {
+                # Update the message history
+                "messages": [
+                    ToolMessage(
+                        content=f"Steady state analysis of"
+                        f" {arg_data.experiment_name}"
+                        " was successful.",
+                        tool_call_id=tool_call_id,
+                        artifact={"dic_data": df_steady_state.to_dict(orient="records")},
+                    )
+                ],
+            }
+        )

aiagents4pharma/talk2biomodels/tools/utils.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python3
+
+"""
+Utility functions for T2B tools.
+"""
+
+import basico
+
+
+def get_model_units(model_object):
+    """
+    Get the units of the model.
+
+    Args:
+        model_object: The model object.
+
+    Returns:
+        dict: The units of the model.
+    """
+    model_units = basico.model_info.get_model_units(model=model_object.copasi_model)
+    model_units_y = model_units["quantity_unit"]
+    model_units_x = model_units["time_unit"]
+    return {"y_axis_label": model_units_y, "x_axis_label": model_units_x}

aiagents4pharma/talk2cells/README.md

@@ -0,0 +1 @@
+Please check out the README file in the root folder for more information.

aiagents4pharma/talk2cells/__init__.py

@@ -0,0 +1,5 @@
+"""
+This file is used to import the models and tools.
+"""
+
+from . import agents, states, tools

aiagents4pharma/talk2cells/agents/__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 scp_agent

aiagents4pharma/talk2cells/agents/scp_agent.py

@@ -0,0 +1,87 @@
+# /usr/bin/env python3
+
+"""
+This is the agent file for the Talk2Cells graph.
+"""
+
+import logging
+import os
+
+from langchain_openai import ChatOpenAI
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import START, StateGraph
+from langgraph.prebuilt import ToolNode, create_react_agent
+
+from ..states.state_talk2cells import Talk2Cells
+from ..tools.scp_agent.display_studies import display_studies
+from ..tools.scp_agent.search_studies import search_studies
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def get_app(uniq_id):
+    """
+    This function returns the langraph app.
+    """
+
+    def agent_scp_node(state: Talk2Cells):
+        """
+        This function calls the model.
+        """
+        logger.log(logging.INFO, "Creating SCP_Agent node with thread_id %s", uniq_id)
+        # Get the messages from the state
+        # messages = state['messages']
+        # Call the model
+        # inputs = {'messages': messages}
+        response = model.invoke(state, {"configurable": {"thread_id": uniq_id}})
+        # The response is a list of messages and may contain `tool calls`
+        # We return a list, because this will get added to the existing list
+        # return {"messages": [response]}
+        return response
+
+    # Define the tools
+    # tools = [search_studies, display_studies]
+    tools = ToolNode([search_studies, display_studies])
+
+    # Create the LLM
+    # And bind the tools to it
+    # model = ChatOpenAI(model="gpt-4o-mini", temperature=0).bind_tools(tools)
+
+    # Create an environment variable to store the LLM model
+    # Check if the environment variable AIAGENTS4PHARMA_LLM_MODEL is set
+    # If not, set it to 'gpt-4o-mini'
+    llm_model = os.getenv("AIAGENTS4PHARMA_LLM_MODEL", "gpt-4o-mini")
+    # print (f'LLM model: {llm_model}')
+    # llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+    llm = ChatOpenAI(model=llm_model, temperature=0)
+    model = create_react_agent(
+        llm,
+        tools=tools,
+        state_schema=Talk2Cells,
+        state_modifier=("You are Talk2Cells agent."),
+        checkpointer=MemorySaver(),
+    )
+
+    # Define a new graph
+    workflow = StateGraph(Talk2Cells)
+
+    # Define the two nodes we will cycle between
+    workflow.add_node("agent_scp", agent_scp_node)
+
+    # Set the entrypoint as `agent`
+    # This means that this node is the first one called
+    workflow.add_edge(START, "agent_scp")
+
+    # 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")
+
+    return app

aiagents4pharma/talk2cells/states/__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 state_talk2cells

aiagents4pharma/talk2cells/states/state_talk2cells.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+
+"""
+This is the state file for the Talk2Cells agent.
+"""
+
+from langgraph.prebuilt.chat_agent_executor import AgentState
+
+
+class Talk2Cells(AgentState):
+    """
+    The state for the Talk2Cells agent.
+    """
+
+    search_table: str

aiagents4pharma/talk2cells/tests/scp_agent/test_scp_agent.py

@@ -0,0 +1,22 @@
+"""
+Test cases for the search_studies
+"""
+
+# from ..tools.search_studies import search_studies
+from langchain_core.messages import HumanMessage
+
+from aiagents4pharma.talk2cells.agents.scp_agent import get_app
+
+
+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/__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 scp_agent

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, search_studies

aiagents4pharma/talk2cells/tools/scp_agent/display_studies.py

@@ -0,0 +1,27 @@
+#!/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"]