aiagents4pharma 0.0.0__py3-none-any.whl

aiagents4pharma/talk2biomodels/tools/load_arguments.py

@@ -0,0 +1,134 @@
+#!/usr/bin/env python3
+
+"""
+A utility module for defining the dataclasses
+for the arguments to set up initial settings
+before the experiment is run.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Annotated
+
+import basico
+from pydantic import Field
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TimeData:
+    """
+    Dataclass for storing the time data.
+    """
+
+    duration: int | float = Field(description="Duration of the simulation", default=100)
+    interval: int | float = Field(
+        description="The interval is the time step or"
+        " the step size of the simulation. It is unrelated"
+        " to the step size of species concentration and parameter values.",
+        default=100,
+    )
+
+
+@dataclass
+class SpeciesInitialData:
+    """
+    Dataclass for storing the species initial data.
+    """
+
+    species_name: list[str] = Field(
+        description="List of species whose initial concentration is to be set."
+        " This does not include species that reoccur or the species whose"
+        " concentration is to be determined/observed at the end of the experiment."
+        " Do not hallucinate the species name.",
+        default=[],
+    )
+    species_concentration: list[int | float] = Field(
+        description="List of initial concentrations of species."
+        " This does not include species that reoccur or the species whose"
+        " concentration is to be determined/observed at the end of the experiment."
+        " Do not hallucinate the species concentration.",
+        default=[],
+    )
+
+
+@dataclass
+class TimeSpeciesNameConcentration:
+    """
+    Dataclass for storing the time,
+    species name, and concentration data.
+    """
+
+    time: int | float = Field(description="time point where the event occurs")
+    species_name: str = Field(description="species name")
+    species_concentration: int | float = Field(
+        description="species concentration at the time point"
+    )
+
+
+@dataclass
+class ReocurringData:
+    """
+    Dataclass for species that reoccur. In other words,
+    the concentration of the species resets to a certain
+    value after a certain time interval.
+    """
+
+    data: list[TimeSpeciesNameConcentration] = Field(
+        description="List of time, name, and concentration data"
+        " of species or parameters that reoccur",
+        default=[],
+    )
+
+
+@dataclass
+class ArgumentData:
+    """
+    Dataclass for storing the argument data.
+    """
+
+    experiment_name: Annotated[
+        str,
+        "An AI assigned _ separated name of"
+        " the experiment based on human query"
+        " and the context of the experiment."
+        " This must be set before the experiment is run.",
+    ]
+    time_data: TimeData | None = Field(description="time data", default=None)
+    species_to_be_analyzed_before_experiment: SpeciesInitialData | None = Field(
+        description="Data of species whose initial concentration"
+        " is to be set before the experiment. This does not include"
+        " species that reoccur or the species whose concentration"
+        " is to be determined at the end of the experiment.",
+        default=None,
+    )
+    reocurring_data: ReocurringData | None = Field(
+        description="List of concentration and time data of species that"
+        " reoccur. For example, a species whose concentration resets"
+        " to a certain value after a certain time interval.",
+        default=None,
+    )
+
+
+def add_rec_events(model_object, reocurring_data):
+    """
+    Add reocurring events to the model.
+
+    Args:
+        model_object: The model object.
+        reocurring_data: The reocurring data.
+
+    Returns:
+        None
+    """
+    for row in reocurring_data.data:
+        tp, sn, sc = row.time, row.species_name, row.species_concentration
+        basico.add_event(
+            f"{sn}_{tp}",
+            f"Time > {tp}",
+            [[sn, str(sc)]],
+            model=model_object.copasi_model,
+        )

aiagents4pharma/talk2biomodels/tools/load_biomodel.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+
+"""
+Function for loading the BioModel.
+"""
+
+from typing import Annotated, Any
+
+from pydantic import BaseModel, BeforeValidator
+
+from ..models.basico_model import BasicoModel
+
+
+def ensure_biomodel_id(value: Any) -> Any:
+    """
+    Ensure that the biomodel_id is an integer or a string starting with 'BIOMD' or 'MODEL'.
+    """
+    if isinstance(value, int):
+        return value
+    if isinstance(value, str) and (value.startswith("BIOMD") or value.startswith("MODEL")):
+        return value
+    raise ValueError("biomodel_id must be an integer or a string starting with 'BIOMD' or 'MODEL'.")
+
+
+class ModelData(BaseModel):
+    """
+    Base model for the model data.
+    """
+
+    biomodel_id: Annotated[int | str, BeforeValidator(ensure_biomodel_id)] = None
+    # sbml_file_path: Optional[str] = None
+    use_uploaded_sbml_file: bool = False
+
+
+def load_biomodel(sys_bio_model, sbml_file_path=None):
+    """
+    Load the BioModel.
+    """
+    model_object = None
+    if sys_bio_model.biomodel_id:
+        model_object = BasicoModel(biomodel_id=sys_bio_model.biomodel_id)
+    elif sbml_file_path:
+        model_object = BasicoModel(sbml_file_path=sbml_file_path)
+    return model_object

aiagents4pharma/talk2biomodels/tools/parameter_scan.py

@@ -0,0 +1,310 @@
+#!/usr/bin/env python3
+
+"""
+Tool for parameter scan.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Annotated
+
+import basico
+import pandas as pd
+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 SpeciesInitialData, TimeData
+from .load_biomodel import ModelData, load_biomodel
+from .utils import get_model_units
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ParameterScanData(BaseModel):
+    """
+    Dataclass for storing the parameter scan data.
+    """
+
+    species_names: list[str] = Field(
+        description="species to be observed after each scan."
+        " These are the species whose concentration"
+        " will be observed after the parameter scan."
+        " Do not make up this data.",
+        default=[],
+    )
+    species_parameter_name: str = Field(
+        description="Species or parameter name to be scanned."
+        " This is the species or parameter whose value will be scanned"
+        " over a range of values. This does not include the species"
+        " that are to be observed after the scan."
+        "Do not make up this data.",
+        default=None,
+    )
+    species_parameter_values: list[int | float] = Field(
+        description="Species or parameter values to be scanned."
+        " These are the values of the species or parameters that will be"
+        " scanned over a range of values. This does not include the "
+        "species that are to be observed after the scan."
+        "Do not make up this data.",
+        default=None,
+    )
+
+
+@dataclass
+class ArgumentData:
+    """
+    Dataclass for storing the argument data.
+    """
+
+    time_data: TimeData = Field(description="time data", default=None)
+    species_to_be_analyzed_before_experiment: SpeciesInitialData | None = Field(
+        description=" This is the initial condition of the model."
+        " This does not include species that reoccur or the species"
+        " whose concentration is to be determined/observed at the end"
+        " of the experiment. This also does not include the species"
+        " or the parameter that is to be scanned. Do not make up this data.",
+        default=None,
+    )
+    parameter_scan_data: ParameterScanData = Field(description="parameter scan data", default=None)
+    experiment_name: str = Field(
+        description="An AI assigned `_` separated unique name of"
+        " the parameter scan experiment based on human query."
+        " This must be unique for each experiment."
+    )
+
+
+def make_list_dic_scanned_data(dic_param_scan, arg_data, sys_bio_model, tool_call_id):
+    """
+    Prepare the list dictionary of scanned data
+    that will be passed to the state of the graph.
+
+    Args:
+        dic_param_scan: Dictionary of parameter scan results.
+        arg_data: The argument data.
+        sys_bio_model: The model data.
+        tool_call_id: The tool call ID.
+
+    Returns:
+        list: List of dictionary of scanned data.
+    """
+    list_dic_scanned_data = []
+    for species_name, df_param_scan in dic_param_scan.items():
+        logger.log(
+            logging.INFO,
+            "Parameter scan results for %s with shape %s",
+            species_name,
+            df_param_scan.shape,
+        )
+        # Prepare the list dictionary of scanned data
+        # that will be passed to the state of the graph
+        list_dic_scanned_data.append(
+            {
+                "name": arg_data.experiment_name + ":" + species_name,
+                "source": (sys_bio_model.biomodel_id if sys_bio_model.biomodel_id else "upload"),
+                "tool_call_id": tool_call_id,
+                "data": df_param_scan.to_dict(),
+            }
+        )
+    return list_dic_scanned_data
+
+
+def run_parameter_scan(model_object, arg_data, dic_species_data, duration, interval) -> dict:
+    """
+    Run parameter scan on the model.
+
+    Args:
+        model_object: The model object.
+        arg_data: The argument data.
+        dic_species_data: Dictionary of species data.
+        duration: Duration of the simulation.
+        interval: Interval between time points in the simulation.
+
+    Returns:
+        dict: Dictionary of parameter scan results. Each key is a species name
+        and each value is a DataFrame containing the results of the parameter scan.
+    """
+    # Extract all parameter names from the model
+    df_all_parameters = basico.model_info.get_parameters(model=model_object.copasi_model)
+    all_parameters = []
+    if df_all_parameters is not None:
+        # For example model 10 in the BioModels database
+        # has no parameters
+        all_parameters = df_all_parameters.index.tolist()
+
+    # Extract all species name from the model
+    df_all_species = basico.model_info.get_species(model=model_object.copasi_model)
+    all_species = df_all_species["display_name"].tolist()
+
+    # Verify if the given species or parameter names to be scanned are valid
+    if arg_data.parameter_scan_data.species_parameter_name not in all_parameters + all_species:
+        logger.error(
+            "Invalid species or parameter name: %s",
+            arg_data.parameter_scan_data.species_parameter_name,
+        )
+        raise ValueError(
+            "Invalid species or parameter name: "
+            f"{arg_data.parameter_scan_data.species_parameter_name}."
+        )
+
+    # Dictionary to store the parameter scan results
+    dic_param_scan_results = {}
+
+    # Loop through the species names that are to be observed
+    for species_name in arg_data.parameter_scan_data.species_names:
+        # Verify if the given species name to be observed is valid
+        if species_name not in all_species:
+            logger.error("Invalid species name: %s", species_name)
+            raise ValueError(f"Invalid species name: {species_name}.")
+
+        # Copy the model object to avoid modifying the original model
+        model_object_copy = model_object.model_copy()
+
+        # Update the fixed model species and parameters
+        # These are the initial conditions of the model
+        # set by the user
+        model_object_copy.update_parameters(dic_species_data)
+
+        # Initialize empty DataFrame to store results
+        # of the parameter scan
+        df_param_scan = pd.DataFrame()
+
+        # Loop through the parameter that are to be scanned
+        for param_value in arg_data.parameter_scan_data.species_parameter_values:
+            # Update the parameter value in the model
+            model_object_copy.update_parameters(
+                {arg_data.parameter_scan_data.species_parameter_name: param_value}
+            )
+            # Simulate the model
+            model_object_copy.simulate(duration=duration, interval=interval)
+            # If the column name 'Time' is not present in the results DataFrame
+            if "Time" not in df_param_scan.columns:
+                df_param_scan["Time"] = model_object_copy.simulation_results["Time"]
+            # Add the simulation results to the results DataFrame
+            col_name = f"{arg_data.parameter_scan_data.species_parameter_name}_{param_value}"
+            df_param_scan[col_name] = model_object_copy.simulation_results[species_name]
+
+        logger.log(logging.INFO, "Parameter scan results with shape %s", df_param_scan.shape)
+
+        # Add the results of the parameter scan to the dictionary
+        dic_param_scan_results[species_name] = df_param_scan
+    # return df_param_scan
+    return dic_param_scan_results
+
+
+class ParameterScanInput(BaseModel):
+    """
+    Input schema for the ParameterScan 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 parameter scan name and
+                                   data""",
+        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 ParameterScanTool(BaseTool):
+    """
+    Tool for parameter scan.
+    """
+
+    name: str = "parameter_scan"
+    description: str = """A tool to perform scanning of a given
+    parameter over a range of values and observe the effect on
+    the concentration of a given species"""
+    args_schema: type[BaseModel] = ParameterScanInput
+
+    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 parameter_scan 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_data = {}
+        if arg_data:
+            # Prepare the dictionary of species data
+            if arg_data.species_to_be_analyzed_before_experiment is not None:
+                dic_species_data = 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
+
+        # Run the parameter scan
+        dic_param_scan = run_parameter_scan(
+            model_object, arg_data, dic_species_data, duration, interval
+        )
+
+        logger.log(logging.INFO, "Parameter scan results ready")
+        # Prepare the list dictionary of scanned data
+        list_dic_scanned_data = make_list_dic_scanned_data(
+            dic_param_scan, arg_data, sys_bio_model, tool_call_id
+        )
+        # 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_scanned_data": list_dic_scanned_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"Parameter scan results of {arg_data.experiment_name}",
+                        tool_call_id=tool_call_id,
+                        artifact=get_model_units(model_object),
+                    )
+                ],
+            }
+        )

aiagents4pharma/talk2biomodels/tools/query_article.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+
+"""
+Tool for asking questions to the article.
+"""
+
+import logging
+from typing import Annotated
+
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_core.tools import BaseTool
+from langchain_core.vectorstores import InMemoryVectorStore
+from langgraph.prebuilt import InjectedState
+from pydantic import BaseModel, Field
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class QueryArticleInput(BaseModel):
+    """
+    Input schema for the query_articles tool.
+    """
+
+    question: Annotated[str, Field(description="User question to search articles.")]
+    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 QueryArticle(BaseTool):
+    """
+    Tool to ask questions to the article.
+    """
+
+    name: str = "query_article"
+    description: str = "Ask questions to the article."
+    args_schema: type[BaseModel] = QueryArticleInput
+
+    def _run(self, question: str, state: Annotated[dict, InjectedState]):
+        """
+        Run the tool.
+
+        Args:
+            query (str): The search query.
+        """
+        logger.log(logging.INFO, "loading the article from %s", state["pdf_file_name"])
+        logger.log(logging.INFO, "searching the article with the question: %s", question)
+        # Load the article
+        loader = PyPDFLoader(state["pdf_file_name"])
+        # Load the pages of the article
+        pages = []
+        for page in loader.lazy_load():
+            pages.append(page)
+        # Set up text embedding model
+        text_embedding_model = state["text_embedding_model"]
+        logging.info("Loaded text embedding model %s", text_embedding_model)
+        # Create a vector store from the pages
+        vector_store = InMemoryVectorStore.from_documents(pages, text_embedding_model)
+        # Search the article with the question
+        docs = vector_store.similarity_search(question)
+        # Return the content of the pages
+        return "\n".join([doc.page_content for doc in docs])

aiagents4pharma/talk2biomodels/tools/save_model.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+
+"""
+Tool for saving models.
+"""
+
+import logging
+import os
+from typing import Annotated
+
+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
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class SaveModelInput(BaseModel):
+    """
+    Input schema for the save model tool.
+    """
+
+    path_to_folder: str = Field(
+        description="Path to folder to save the model. Keep it to . if not provided.", default="."
+    )
+    output_filename: str = Field(
+        description="Filename to save the model as. Default is 'saved_model.xml'.",
+        default="saved_model.xml",
+    )
+    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 SaveModelTool(BaseTool):
+    """
+    Tool for saving a model.
+    """
+
+    name: str = "save_model"
+    description: str = "A tool to save the current biomodel to a \
+                        user specified path with the default filename\
+                         'saved_model.xml'"
+    args_schema: type[BaseModel] = SaveModelInput
+    return_direct: bool = False
+
+    def _run(
+        self,
+        tool_call_id: Annotated[str, InjectedToolCallId],
+        state: Annotated[dict, InjectedState],
+        path_to_folder: str = ".",
+        output_filename: str = "saved_model.xml",
+    ) -> Command:
+        """
+        Run the tool.
+
+        Args:
+            path (str): The path to save the model.
+            tool_call_id (str): The tool call ID.
+
+        Returns:
+
+        """
+        logger.log(
+            logging.INFO,
+            "Saving model to path: %s with filename: %s",
+            path_to_folder,
+            output_filename,
+        )
+        # Check if path does not exist
+        if not os.path.exists(path_to_folder):
+            content = f"Error: Path {path_to_folder} does not exist."
+            logger.error(content)
+        else:
+            logger.info("Saving now")
+            # Save the model to the specified path
+            with open(os.path.join(path_to_folder, output_filename), "w", encoding="utf-8") as f:
+                f.write(state["model_as_string"][-1])
+            content = f"Model saved successfully to {path_to_folder}/{output_filename}."
+            logger.info(content)
+        # Return the updated state of the tool
+        return Command(
+            update={
+                # update the message history
+                "messages": [
+                    ToolMessage(
+                        content=content,
+                        tool_call_id=tool_call_id,
+                    )
+                ],
+            }
+        )