aiagents4pharma 0.0.0__py3-none-any.whl

aiagents4pharma/talk2biomodels/tools/custom_plotter.py

@@ -0,0 +1,165 @@
+#!/usr/bin/env python3
+
+"""
+Tool for plotting a custom y-axis of a simulation plot.
+"""
+
+import logging
+from typing import Annotated, Literal
+
+import hydra
+import pandas as pd
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.tools import BaseTool
+from langgraph.prebuilt import InjectedState
+from pydantic import BaseModel, Field
+
+from .load_biomodel import ModelData, load_biomodel
+from .utils import get_model_units
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def extract_relevant_species(question, species_names, state):
+    """
+    Extract the relevant species from the user question.
+
+    Args:
+        question (str): The user question.
+        species_names (list): The species names available in the simulation results.
+        state (dict): The state of the graph.
+
+    Returns:
+        CustomHeader: The relevant species
+    """
+
+    # In the following code, we extract the species
+    # from the user question. We use Literal to restrict
+    # the species names to the ones available in the
+    # simulation results.
+    class CustomHeader(BaseModel):
+        """
+        A list of species based on user question.
+
+        This is a Pydantic model that restricts the species
+        names to the ones available in the simulation results.
+
+        If no species is relevant, set the attribute
+        `relevant_species` to None.
+        """
+
+        relevant_species: None | list[Literal[*species_names,]] = Field(
+            description="This is a list of species based on the user question."
+            "It is restricted to the species available in the simulation results."
+            "If no species is relevant, set this attribute to None."
+            "If the user asks for very specific species (for example, using the"
+            "keyword `only` in the question), set this attribute to correspond "
+            "to the species available in the simulation results, otherwise set it to None."
+        )
+
+    # Load hydra configuration
+    with hydra.initialize(version_base=None, config_path="../configs"):
+        cfg = hydra.compose(config_name="config", overrides=["tools/custom_plotter=default"])
+        cfg = cfg.tools.custom_plotter
+    # Get the system prompt
+    system_prompt = cfg.system_prompt_custom_header
+    # Create an instance of the LLM model
+    logging.log(logging.INFO, "LLM model: %s", state["llm_model"])
+    llm = state["llm_model"]
+    llm_with_structured_output = llm.with_structured_output(CustomHeader)
+    prompt = ChatPromptTemplate.from_messages([("system", system_prompt), ("human", "{input}")])
+    few_shot_structured_llm = prompt | llm_with_structured_output
+    return few_shot_structured_llm.invoke(question)
+
+
+class CustomPlotterInput(BaseModel):
+    """
+    Input schema for the custom plotter tool.
+    """
+
+    question: str = Field(description="Description of the plot")
+    sys_bio_model: ModelData = Field(description="model data", default=None)
+    simulation_name: str = Field(description="Name assigned to the simulation")
+    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.
+# 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 CustomPlotterTool(BaseTool):
+    """
+    Tool for custom plotting the y-axis of a plot.
+    """
+
+    name: str = "custom_plotter"
+    description: str = """A visualization tool designed to extract and display a subset
+                        of the larger simulation plot generated by the simulate_model tool.
+                        It allows users to specify particular species for the y-axis,
+                        providing a more targeted view of key species without the clutter
+                        of the full plot."""
+    args_schema: type[BaseModel] = CustomPlotterInput
+    response_format: str = "content_and_artifact"
+
+    def _run(
+        self,
+        question: str,
+        sys_bio_model: ModelData,
+        simulation_name: str,
+        state: Annotated[dict, InjectedState],
+    ) -> tuple[str, None | list[str]]:
+        """
+        Run the tool.
+
+        Args:
+            question (str): The question about the custom plot.
+            sys_bio_model (ModelData): The model data.
+            simulation_name (str): The name assigned to the simulation.
+            state (dict): The state of the graph.
+
+        Returns:
+            str: The answer to the question
+        """
+        logger.log(logging.INFO, "Calling custom_plotter tool %s, %s", question, sys_bio_model)
+        # Load the model
+        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)
+        dic_simulated_data = {}
+        for data in state["dic_simulated_data"]:
+            for key in data:
+                if key not in dic_simulated_data:
+                    dic_simulated_data[key] = []
+                dic_simulated_data[key] += [data[key]]
+        # Create a pandas dataframe from the dictionary
+        df = pd.DataFrame.from_dict(dic_simulated_data)
+        # Get the simulated data for the current tool call
+        df = pd.DataFrame(df[df["name"] == simulation_name]["data"].iloc[0])
+        # df = pd.DataFrame.from_dict(state['dic_simulated_data'])
+        species_names = df.columns.tolist()
+        # Exclude the time column
+        species_names.remove("Time")
+        logging.log(logging.INFO, "Species names: %s", species_names)
+        # Extract the relevant species from the user question
+        results = extract_relevant_species(question, species_names, state)
+        print(results)
+        if results.relevant_species is None:
+            raise ValueError(
+                "No species found in the simulation results \
+                             that matches the user prompt."
+            )
+        extracted_species = []
+        # Extract the species from the results
+        # that are available in the simulation results
+        for species in results.relevant_species:
+            if species in species_names:
+                extracted_species.append(species)
+        logging.info("Extracted species: %s", extracted_species)
+        # Include the time column
+        extracted_species.insert(0, "Time")
+        return f"Custom plot {simulation_name}", {
+            "dic_data": df[extracted_species].to_dict(orient="records")
+        } | get_model_units(model_object)

aiagents4pharma/talk2biomodels/tools/get_annotation.py

@@ -0,0 +1,342 @@
+#!/usr/bin/env python3
+
+"""
+This module contains the `GetAnnotationTool` for fetching species annotations
+based on the provided model and species names.
+"""
+
+import logging
+import math
+from dataclasses import dataclass
+from typing import Annotated, Literal
+
+import basico
+import hydra
+import pandas as pd
+from langchain_core.messages import ToolMessage
+from langchain_core.tools.base import BaseTool, InjectedToolCallId
+from langgraph.prebuilt import InjectedState
+from langgraph.types import Command
+from pydantic import BaseModel, Field
+
+from ..api.ols import search_ols_labels
+from ..api.uniprot import search_uniprot_labels
+
+# from langchain_openai import ChatOpenAI
+from .load_biomodel import ModelData, load_biomodel
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+ols_ontology_abbreviations = {"pato", "chebi", "sbo", "fma", "pr", "go"}
+
+
+def extract_relevant_species_names(model_object, arg_data, state):
+    """
+    Extract relevant species names based on the user question.
+    """
+    # Load hydra configuration
+    with hydra.initialize(version_base=None, config_path="../configs"):
+        cfg = hydra.compose(config_name="config", overrides=["tools/get_annotation=default"])
+        cfg = cfg.tools.get_annotation
+    logger.info(
+        "Loaded the following system prompt for the LLM to get a structured output: %s",
+        cfg.prompt,
+    )
+
+    # Extract all the species names from the model
+    df_species = basico.model_info.get_species(model=model_object.copasi_model)
+    if df_species is None:
+        raise ValueError("Unable to extract species from the model.")
+    # Get all the species names
+    all_species_names = df_species.index.tolist()
+
+    # Define a structured output for the LLM model
+    class CustomHeader(BaseModel):
+        """
+        A list of species based on user question.
+        """
+
+        relevant_species: None | list[Literal[*all_species_names,]] = Field(
+            description="""List of species based on user question.
+                If no relevant species are found, it must be None."""
+        )
+
+    # Create an instance of the LLM model
+    llm = state["llm_model"]
+    # Get the structured output from the LLM model
+    llm_with_structured_output = llm.with_structured_output(CustomHeader)
+    # Define the question for the LLM model using the prompt
+    question = cfg.prompt
+    question += f"Here is the user question: {arg_data.user_question}"
+    # Invoke the LLM model with the user question
+    results = llm_with_structured_output.invoke(question)
+    logging.info("Results from the LLM model: %s", results)
+    # Check if the returned species names are empty
+    if not results.relevant_species:
+        raise ValueError("Model does not contain the requested species.")
+    extracted_species = []
+    # Extract all the species names from the model
+    for species in results.relevant_species:
+        if species in all_species_names:
+            extracted_species.append(species)
+    logger.info("Extracted species: %s", extracted_species)
+    return extracted_species
+
+
+def prepare_content_msg(species_without_description: list[str]):
+    """
+    Prepare the content message.
+    """
+    content = "Successfully extracted annotations for the species."
+    if species_without_description:
+        content += f"""The descriptions for the following species
+                        were not found:
+                        {", ".join(species_without_description)}."""
+    return content
+
+
+@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.",
+    ]
+    user_question: Annotated[str, "Description of the user question"]
+
+
+class GetAnnotationInput(BaseModel):
+    """
+    Input schema for annotation tool.
+    """
+
+    arg_data: ArgumentData = Field(description="argument data")
+    sys_bio_model: ModelData = Field(description="model data")
+    tool_call_id: Annotated[str, InjectedToolCallId]
+    state: Annotated[dict, InjectedState]
+
+
+class GetAnnotationTool(BaseTool):
+    """
+    Tool for fetching species annotations based on the provided model and species names.
+    """
+
+    name: str = "get_annotation"
+    description: str = """A tool to extract annotations for a list of species names
+                        based on the provided model. Annotations include
+                        the species name, description, database, ID, link,
+                        and qualifier. The tool can handle multiple species
+                        in a single invoke."""
+    args_schema: type[BaseModel] = GetAnnotationInput
+    return_direct: bool = False
+
+    def _run(
+        self,
+        arg_data: ArgumentData,
+        tool_call_id: Annotated[str, InjectedToolCallId],
+        state: Annotated[dict, InjectedState],
+        sys_bio_model: ModelData = None,
+    ) -> str:
+        """
+        Run the tool.
+        """
+        logger.info(
+            "Running the GetAnnotationTool tool for species %s, %s",
+            arg_data.user_question,
+            arg_data.experiment_name,
+        )
+
+        # Prepare the model object
+        sbml_file_path = state["sbml_file_path"][-1] if state["sbml_file_path"] else None
+        model_object = load_biomodel(sys_bio_model, sbml_file_path=sbml_file_path)
+
+        # Extract relevant species names based on the user question
+        list_species_names = extract_relevant_species_names(model_object, arg_data, state)
+        print(list_species_names)
+
+        (annotations_df, species_without_description) = self._fetch_annotations(list_species_names)
+
+        # Process annotations
+        annotations_df = self._process_annotations(annotations_df)
+
+        # Prepare the simulated data
+        dic_annotations_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": annotations_df.to_dict(),
+        }
+
+        # Update the state with the annotations data
+        dic_updated_state_for_model = {}
+        for key, value in {
+            "model_id": [sys_bio_model.biomodel_id],
+            "sbml_file_path": [sbml_file_path],
+            "dic_annotations_data": [dic_annotations_data],
+        }.items():
+            if value:
+                dic_updated_state_for_model[key] = value
+
+        return Command(
+            update=dic_updated_state_for_model
+            | {
+                "messages": [
+                    ToolMessage(
+                        content=prepare_content_msg(species_without_description),
+                        artifact=True,
+                        tool_call_id=tool_call_id,
+                    )
+                ]
+            }
+        )
+
+    def _fetch_annotations(self, list_species_names: list[str]) -> tuple:
+        """
+        Fetch annotations for the given species names from the model.
+        In this method, we fetch the MIRIAM annotations for the species names.
+        If the annotation is not found, we add the species to the list of
+        species not found. If the annotation is found, we extract the descriptions
+        from the annotation and add them to the data list.
+
+        Args:
+            list_species_names (List[str]): List of species names to fetch annotations for.
+
+        Returns:
+            tuple: A tuple containing the annotations dataframe, species not found list,
+                   and description not found list.
+        """
+        description_not_found = []
+        data = []
+
+        # Loop through the species names
+        for species in list_species_names:
+            # Get the MIRIAM annotation for the species
+            annotation = basico.get_miriam_annotation(name=species)
+
+            # Extract the descriptions from the annotation
+            descriptions = annotation.get("descriptions", [])
+
+            if descriptions == []:
+                description_not_found.append(species)
+                continue
+
+            # Loop through the descriptions and add them to the data list
+            for desc in descriptions:
+                data.append(
+                    {
+                        "Species Name": species,
+                        "Link": desc["id"],
+                        "Qualifier": desc["qualifier"],
+                    }
+                )
+
+        # Create a dataframe from the data list
+        annotations_df = pd.DataFrame(data)
+
+        # Return the annotations dataframe and the species not found list
+        return annotations_df, description_not_found
+
+    def _process_annotations(self, annotations_df: pd.DataFrame) -> pd.DataFrame:
+        """
+        Process annotations dataframe to add additional information.
+        In this method, we add a new column for the ID, a new column for the database,
+        and a new column for the description. We then reorder the columns and process
+        the link to format it correctly.
+
+        Args:
+            annotations_df (pd.DataFrame): Annotations dataframe to process.
+
+        Returns:
+            pd.DataFrame: Processed annotations dataframe
+        """
+        logger.info("Processing annotations.")
+        # Add a new column for the ID
+        # Get the ID from the link key
+        annotations_df["Id"] = annotations_df["Link"].str.split("/").str[-1]
+
+        # Add a new column for the database
+        # Get the database from the link key
+        annotations_df["Database"] = annotations_df["Link"].str.split("/").str[-2]
+
+        # Fetch descriptions for the IDs based on the database type
+        # by qyerying the respective APIs
+        identifiers = annotations_df[["Id", "Database"]].to_dict(orient="records")
+        descriptions = self._fetch_descriptions(identifiers)
+
+        # Add a new column for the description
+        # Get the description from the descriptions dictionary
+        # based on the ID. If the description is not found, use '-'
+        annotations_df["Description"] = annotations_df["Id"].apply(
+            lambda x: descriptions.get(x, "-")
+        )
+        # annotations_df.index = annotations_df.index + 1
+
+        # Reorder the columns
+        annotations_df = annotations_df[
+            ["Species Name", "Description", "Database", "Id", "Link", "Qualifier"]
+        ]
+
+        # Process the link to format it correctly
+        annotations_df["Link"] = annotations_df["Link"].apply(self._process_link)
+
+        # Return the processed annotations dataframe
+        return annotations_df
+
+    def _process_link(self, link: str) -> str:
+        """
+        Process link to format it correctly.
+        """
+        for ols_ontology_abbreviation in ols_ontology_abbreviations:
+            if ols_ontology_abbreviation + "/" in link:
+                link = link.replace(f"{ols_ontology_abbreviation}/", "")
+        return link
+
+    def _fetch_descriptions(self, data: list[dict[str, str]]) -> dict[str, str]:
+        """
+        Fetch protein names or labels based on the database type.
+        """
+        logger.info("Fetching descriptions for the IDs.")
+        results = {}
+        grouped_data = {}
+
+        # In the following loop, we create a dictionary with database as the key
+        # and a list of identifiers as the value. If either the database or the
+        # identifier is NaN, we set it to None.
+        for entry in data:
+            identifier = entry.get("Id")
+            database = entry.get("Database")
+            # Check if database is NaN
+            if isinstance(database, float):
+                if math.isnan(database):
+                    database = None
+                    results[identifier or "unknown"] = "-"
+            else:
+                database = database.lower()
+                grouped_data.setdefault(database, []).append(identifier)
+
+        # In the following loop, we fetch the descriptions for the identifiers
+        # based on the database type.
+        # Constants
+
+        for database, identifiers in grouped_data.items():
+            if database == "uniprot":
+                results.update(search_uniprot_labels(identifiers))
+            elif database in ols_ontology_abbreviations:
+                annotations = search_ols_labels(
+                    [{"Id": id_, "Database": database} for id_ in identifiers]
+                )
+                for identifier in identifiers:
+                    results[identifier] = annotations.get(database, {}).get(identifier, "-")
+            else:
+                # For any other database types, do not fetch; mark as unknown
+                for identifier in identifiers:
+                    results[identifier] = "-"
+        return results

aiagents4pharma/talk2biomodels/tools/get_modelinfo.py

@@ -0,0 +1,159 @@
+#!/usr/bin/env python3
+
+"""
+Tool for get model information.
+"""
+
+import logging
+from dataclasses import dataclass
+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_biomodel import ModelData, load_biomodel
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class RequestedModelInfo:
+    """
+    Dataclass for storing the requested model information.
+    """
+
+    species: bool = Field(description="Get species from the model.", default=False)
+    parameters: bool = Field(description="Get parameters from the model.", default=False)
+    compartments: bool = Field(description="Get compartments from the model.", default=False)
+    units: bool = Field(description="Get units from the model.", default=False)
+    description: bool = Field(description="Get description from the model.", default=False)
+    name: bool = Field(description="Get name from the model.", default=False)
+
+
+class GetModelInfoInput(BaseModel):
+    """
+    Input schema for the GetModelInfo tool.
+    """
+
+    requested_model_info: RequestedModelInfo = Field(description="requested model information")
+    sys_bio_model: ModelData = Field(description="model data")
+    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 GetModelInfoTool(BaseTool):
+    """
+    This tool ise used extract model information.
+    """
+
+    name: str = "get_modelinfo"
+    description: str = """A tool for extracting name,
+                    description, species, parameters,
+                    compartments, and units from a model."""
+    args_schema: type[BaseModel] = GetModelInfoInput
+
+    def _run(
+        self,
+        requested_model_info: RequestedModelInfo,
+        tool_call_id: Annotated[str, InjectedToolCallId],
+        state: Annotated[dict, InjectedState],
+        sys_bio_model: ModelData | None = None,
+    ) -> Command:
+        """
+        Run the tool.
+
+        Args:
+            requested_model_info (RequestedModelInfo): The requested model information.
+            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.
+
+        Returns:
+            Command: The updated state of the tool.
+        """
+        logger.log(
+            logging.INFO,
+            "Calling get_modelinfo tool %s, %s",
+            sys_bio_model,
+            requested_model_info,
+        )
+        # print (state, 'state')
+        sbml_file_path = state["sbml_file_path"][-1] if len(state["sbml_file_path"]) > 0 else None
+        model_obj = load_biomodel(sys_bio_model, sbml_file_path=sbml_file_path)
+        dic_results = {}
+        # Extract species from the model
+        if requested_model_info.species:
+            df_species = basico.model_info.get_species(model=model_obj.copasi_model)
+            if df_species is None:
+                raise ValueError("Unable to extract species from the model.")
+            # Convert index into a column
+            df_species.reset_index(inplace=True)
+            dic_results["Species"] = df_species[
+                [
+                    "name",
+                    "compartment",
+                    "type",
+                    "unit",
+                    "initial_concentration",
+                    "display_name",
+                ]
+            ]
+            # Convert this into a dictionary
+            dic_results["Species"] = dic_results["Species"].to_dict(orient="records")
+
+        # Extract parameters from the model
+        if requested_model_info.parameters:
+            df_parameters = basico.model_info.get_parameters(model=model_obj.copasi_model)
+            if df_parameters is None:
+                raise ValueError("Unable to extract parameters from the model.")
+            # Convert index into a column
+            df_parameters.reset_index(inplace=True)
+            dic_results["Parameters"] = df_parameters[
+                ["name", "type", "unit", "initial_value", "display_name"]
+            ]
+            # Convert this into a dictionary
+            dic_results["Parameters"] = dic_results["Parameters"].to_dict(orient="records")
+
+        # Extract compartments from the model
+        if requested_model_info.compartments:
+            df_compartments = basico.model_info.get_compartments(model=model_obj.copasi_model)
+            dic_results["Compartments"] = df_compartments.index.tolist()
+            dic_results["Compartments"] = ",".join(dic_results["Compartments"])
+
+        # Extract description from the model
+        if requested_model_info.description:
+            dic_results["Description"] = model_obj.description
+
+        # Extract description from the model
+        if requested_model_info.name:
+            dic_results["Name"] = model_obj.name
+
+        # Extract time unit from the model
+        if requested_model_info.units:
+            dic_results["Units"] = basico.model_info.get_model_units(model=model_obj.copasi_model)
+
+        # 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],
+        }.items():
+            if value:
+                dic_updated_state_for_model[key] = value
+
+        return Command(
+            update=dic_updated_state_for_model
+            | {
+                # update the message history
+                "messages": [ToolMessage(content=dic_results, tool_call_id=tool_call_id)],
+            }
+        )