aiagents4pharma 1.8.0__py3-none-any.whl → 1.15.0__py3-none-any.whl

aiagents4pharma/talk2biomodels/tests/test_sys_bio_model.py

@@ -0,0 +1,63 @@
+'''
+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")
+    param1: Optional[float] = Field(0.0, description="Parameter 1")
+    param2: Optional[float] = Field(0.0, description="Parameter 2")
+
+    def get_model_metadata(self) -> Dict[str, Union[str, int]]:
+        '''
+        Get the metadata of the model.
+        '''
+        return self.biomodel_id
+
+    def update_parameters(self, parameters):
+        '''
+        Update the model parameters.
+        '''
+        self.param1 = parameters.get('param1', 0.0)
+        self.param2 = parameters.get('param2', 0.0)
+
+    def simulate(self, duration: Union[int, float]) -> List[float]:
+        '''
+        Simulate the model.
+        '''
+        return [self.param1 + self.param2 * t for t in range(int(duration))]
+
+def test_get_model_metadata():
+    '''
+    Test the get_model_metadata method of the BioModel class.
+    '''
+    model = TestBioModel(biomodel_id=123, name="Test Model", description="A test model")
+    metadata = model.get_model_metadata()
+    assert metadata == 123
+
+def test_check_biomodel_id_or_sbml_file_path():
+    '''
+    Test the check_biomodel_id_or_sbml_file_path method of the BioModel class.
+    '''
+    with pytest.raises(ValueError):
+        TestBioModel(name="Test Model", description="A test model")
+
+def test_simulate():
+    '''
+    Test the simulate method of the BioModel class.
+    '''
+    model = TestBioModel(biomodel_id=123, name="Test Model", description="A test model")
+    model.update_parameters({'param1': 1.0, 'param2': 2.0})
+    results = model.simulate(duration=4.0)
+    assert results == [1.0, 3.0, 5.0, 7.0]

aiagents4pharma/talk2biomodels/tools/__init__.py

@@ -6,3 +6,8 @@ from . import simulate_model
 from . import ask_question
 from . import custom_plotter
 from . import get_modelinfo
+from . import parameter_scan
+from . import steady_state
+from . import load_biomodel
+from . import get_annotation
+from . import query_article

aiagents4pharma/talk2biomodels/tools/ask_question.py

@@ -5,11 +5,12 @@ Tool for asking a question about the simulation results.
 """
 
 import logging
-from typing import Type, Annotated
+from typing import Type, Annotated, Literal
+import hydra
+import basico
 import pandas as pd
 from pydantic import BaseModel, Field
 from langchain_core.tools.base import BaseTool
-from langchain.agents.agent_types import AgentType
 from langchain_experimental.agents import create_pandas_dataframe_agent
 from langchain_openai import ChatOpenAI
 from langgraph.prebuilt import InjectedState
@@ -22,7 +23,12 @@ class AskQuestionInput(BaseModel):
     """
     Input schema for the AskQuestion tool.
     """
-    question: str = Field(description="question about the simulation results")
+    question: str = Field(description="question about the simulation and steady state results")
+    experiment_name: str = Field(description="""Name assigned to the simulation
+                                            or steady state analysis when the tool 
+                                            simulate_model or steady_state is invoked.""")
+    question_context: Literal["simulation", "steady_state"] = Field(
+        description="Context of the question")
     state: Annotated[dict, InjectedState]
 
 # Note: It's important that every field has type hints.
@@ -30,43 +36,80 @@ class AskQuestionInput(BaseModel):
 # can lead to unexpected behavior.
 class AskQuestionTool(BaseTool):
     """
-    Tool for calculating the product of two numbers.
+    Tool for asking a question about the simulation or steady state results.
     """
     name: str = "ask_question"
-    description: str = "A tool to ask question about the simulation results."
+    description: str = """A tool to ask question about the
+                        simulation or steady state results."""
     args_schema: Type[BaseModel] = AskQuestionInput
     return_direct: bool = False
 
     def _run(self,
              question: str,
+             experiment_name: str,
+             question_context: Literal["simulation", "steady_state"],
              state: Annotated[dict, InjectedState]) -> str:
         """
         Run the tool.
 
         Args:
-            question (str): The question to ask about the simulation results.
+            question (str): The question to ask about the simulation or steady state results.
             state (dict): The state of the graph.
-            run_manager (Optional[CallbackManagerForToolRun]): The CallbackManagerForToolRun object.
+            experiment_name (str): The name assigned to the simulation or steady state analysis.
 
         Returns:
             str: The answer to the question.
         """
         logger.log(logging.INFO,
-                   "Calling ask_question tool %s", question)
-        # Check if the simulation results are available
-        if 'dic_simulated_data' not in state:
-            return "Please run the simulation first before \
-                asking a question about the simulation results."
-        df = pd.DataFrame.from_dict(state['dic_simulated_data'])
-        prompt_content = None
-        # if run_manager and 'prompt' in run_manager.metadata:
-        #     prompt_content = run_manager.metadata['prompt']
-        # Create a pandas dataframe agent with OpenAI
+                   "Calling ask_question tool %s, %s, %s",
+                   question,
+                   question_context,
+                   experiment_name)
+        # Load hydra configuration
+        with hydra.initialize(version_base=None, config_path="../../configs"):
+            cfg = hydra.compose(config_name='config',
+                                overrides=['talk2biomodels/tools/ask_question=default'])
+            cfg = cfg.talk2biomodels.tools.ask_question
+        # Get the context of the question
+        # and based on the context, get the data
+        # and prompt content to ask the question
+        if question_context == "steady_state":
+            dic_context = state["dic_steady_state_data"]
+            prompt_content = cfg.steady_state_prompt
+        else:
+            dic_context = state["dic_simulated_data"]
+            prompt_content = cfg.simulation_prompt
+        # Extract the
+        dic_data = {}
+        for data in dic_context:
+            for key in data:
+                if key not in dic_data:
+                    dic_data[key] = []
+                dic_data[key] += [data[key]]
+        # Create a pandas dataframe of the data
+        df_data = pd.DataFrame.from_dict(dic_data)
+        # Extract the data for the experiment
+        # matching the experiment name
+        df = pd.DataFrame(
+            df_data[df_data['name'] == experiment_name]['data'].iloc[0]
+        )
+        logger.log(logging.INFO, "Shape of the dataframe: %s", df.shape)
+        # # Extract the model units
+        # model_units = basico.model_info.get_model_units()
+        # Update the prompt content with the model units
+        prompt_content += "Following are the model units:\n"
+        prompt_content += f"{basico.model_info.get_model_units()}\n\n"
+        # Create a pandas dataframe agent
         df_agent = create_pandas_dataframe_agent(
                         ChatOpenAI(model=state['llm_model']),
                         allow_dangerous_code=True,
-                        agent_type=AgentType.OPENAI_FUNCTIONS,
+                        agent_type='tool-calling',
                         df=df,
+                        max_iterations=5,
+                        include_df_in_prompt=True,
+                        number_of_head_rows=df.shape[0],
+                        verbose=True,
                         prefix=prompt_content)
+        # Invoke the agent with the question
         llm_result = df_agent.invoke(question)
         return llm_result["output"]

aiagents4pharma/talk2biomodels/tools/custom_plotter.py

@@ -6,14 +6,11 @@ Tool for plotting a custom figure.
 
 import logging
 from typing import Type, List, TypedDict, Annotated, Tuple, Union, Literal
-from typing import Type, List, TypedDict, Annotated, Tuple, Union, Literal
 from pydantic import BaseModel, Field
 import pandas as pd
-import pandas as pd
 from langchain_openai import ChatOpenAI
 from langchain_core.tools import BaseTool
 from langgraph.prebuilt import InjectedState
-from langgraph.prebuilt import InjectedState
 
 # Initialize logger
 logging.basicConfig(level=logging.INFO)
@@ -24,7 +21,7 @@ class CustomPlotterInput(BaseModel):
     Input schema for the PlotImage tool.
     """
     question: str = Field(description="Description of the plot")
-    state: Annotated[dict, InjectedState]
+    simulation_name: str = Field(description="Name assigned to the simulation")
     state: Annotated[dict, InjectedState]
 
 # Note: It's important that every field has type hints.
@@ -41,10 +38,10 @@ class CustomPlotterTool(BaseTool):
     description: str = "A tool to make custom plots of the simulation results"
     args_schema: Type[BaseModel] = CustomPlotterInput
     response_format: str = "content_and_artifact"
-    response_format: str = "content_and_artifact"
 
     def _run(self,
              question: str,
+             simulation_name: str,
              state: Annotated[dict, InjectedState]
              ) -> Tuple[str, Union[None, List[str]]]:
         """
@@ -53,17 +50,24 @@ class CustomPlotterTool(BaseTool):
         Args:
             question (str): The question about the custom plot.
             state (dict): The state of the graph.
-            question (str): The question about the custom plot.
-            state (dict): The state of the graph.
 
         Returns:
             str: The answer to the question
         """
         logger.log(logging.INFO, "Calling custom_plotter tool %s", question)
-        # Check if the simulation results are available
-        # if 'dic_simulated_data' not in state:
-        #     return "Please run the simulation first before plotting the figure.", None
-        df = pd.DataFrame.from_dict(state['dic_simulated_data'])
+        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')
@@ -76,7 +80,8 @@ class CustomPlotterTool(BaseTool):
             A list of species based on user question.
             """
             relevant_species: Union[None, List[Literal[*species_names]]] = Field(
-                    description="List of species based on user question. If no relevant species are found, it will be None.")
+                    description="""List of species based on user question.
+                    If no relevant species are found, it will be None.""")
         # Create an instance of the LLM model
         llm = ChatOpenAI(model=state['llm_model'], temperature=0)
         llm_with_structured_output = llm.with_structured_output(CustomHeader)
@@ -90,5 +95,6 @@ class CustomPlotterTool(BaseTool):
         logger.info("Extracted species: %s", extracted_species)
         if len(extracted_species) == 0:
             return "No species found in the simulation results that matches the user prompt.", None
-        content = f"Plotted custom figure with species: {', '.join(extracted_species)}"
-        return content, extracted_species
+        # Include the time column
+        extracted_species.insert(0, 'Time')
+        return f"Custom plot {simulation_name}", df[extracted_species].to_dict(orient='records')

aiagents4pharma/talk2biomodels/tools/get_annotation.py

@@ -0,0 +1,304 @@
+#!/usr/bin/env python3
+
+"""
+This module contains the `GetAnnotationTool` for fetching species annotations 
+based on the provided model and species names.
+"""
+import math
+from typing import List, Annotated, Type
+import logging
+from dataclasses import dataclass
+from pydantic import BaseModel, Field
+import basico
+import pandas as pd
+from langgraph.types import Command
+from langgraph.prebuilt import InjectedState
+from langchain_core.tools.base import BaseTool
+from langchain_core.tools.base import InjectedToolCallId
+from langchain_core.messages import ToolMessage
+from .load_biomodel import ModelData, load_biomodel
+from ..api.uniprot import search_uniprot_labels
+from ..api.ols import search_ols_labels
+from ..api.kegg import fetch_kegg_annotations
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+ols_ontology_abbreviations = {'pato', 'chebi', 'sbo', 'fma', 'pr','go'}
+
+def prepare_content_msg(species_not_found: List[str],
+                        species_without_description: List[str]):
+    """
+    Prepare the content message.
+    """
+    content = 'Successfully extracted annotations for the species.'
+    if species_not_found:
+        content += f'''The following species do not exist, and
+                        hence their annotations were not extracted:
+                        {', '.join(species_not_found)}.'''
+    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."]
+    list_species_names: List[str] = Field(
+        default=None,
+        description='''List of species names to fetch annotations for.
+                      If not provided, annotations for all
+                      species in the model will be fetched.'''
+    )
+
+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.list_species_names,
+                    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 all the species names from the model
+        df_species = basico.model_info.get_species(model=model_object.copasi_model)
+
+        if df_species is None:
+            # for example this may happen with model 20
+            raise ValueError("Unable to extract species from the model.")
+        # Fetch annotations for the species names
+        arg_data.list_species_names = arg_data.list_species_names or df_species.index.tolist()
+
+        (annotations_df,
+         species_not_found,
+         species_without_description) = self._fetch_annotations(arg_data.list_species_names)
+
+        # Check if annotations are empty
+        # If empty, return a message
+        if annotations_df.empty:
+            logger.warning("The annotations dataframe is empty.")
+            return prepare_content_msg(species_not_found, species_without_description)
+
+        # 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_not_found,
+                                                    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.
+        """
+        species_not_found = []
+        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)
+            # If the annotation is not found, add the species to the list
+            if annotation is None:
+                species_not_found.append(species)
+                continue
+
+            # 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, species_not_found, 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}/", "")
+        if "kegg.compound" in link:
+            link = link.replace("kegg.compound/", "kegg.compound:")
+        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, "-")
+            elif database == 'kegg.compound':
+                data = [{"Id": identifier, "Database": "kegg.compound"}
+                        for identifier in identifiers]
+                annotations = fetch_kegg_annotations(data)
+                for identifier in identifiers:
+                    results[identifier] = annotations.get(database, {}).get(identifier, "-")
+            else:
+                for identifier in identifiers:
+                    results[identifier] = "-"
+        return results

aiagents4pharma/talk2biomodels/tools/get_modelinfo.py

@@ -25,12 +25,12 @@ class RequestedModelInfo:
     """
     Dataclass for storing the requested model information.
     """
-    species: bool = Field(description="Get species from the model.")
-    parameters: bool = Field(description="Get parameters from the model.")
-    compartments: bool = Field(description="Get compartments from the model.")
-    units: bool = Field(description="Get units from the model.")
-    description: bool = Field(description="Get description from the model.")
-    name: bool = Field(description="Get name from the model.")
+    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):
     """
@@ -47,8 +47,10 @@ class GetModelInfoTool(BaseTool):
     """
     This tool ise used extract model information.
     """
-    name: str = "get_parameters"
-    description: str = "A tool for extracting 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,
@@ -81,7 +83,7 @@ class GetModelInfoTool(BaseTool):
         # Extract species from the model
         if requested_model_info.species:
             df_species = basico.model_info.get_species(model=model_obj.copasi_model)
-            dic_results['Species'] = df_species.index.tolist()
+            dic_results['Species'] = df_species['display_name'].tolist()
             dic_results['Species'] = ','.join(dic_results['Species'])
 
         # Extract parameters from the model