aiagents4pharma 1.44.0__py3-none-any.whl → 1.45.1__py3-none-any.whl

aiagents4pharma/talk2biomodels/tools/get_annotation.py

@@ -1,33 +1,37 @@
 #!/usr/bin/env python3
 
 """
-This module contains the `GetAnnotationTool` for fetching species annotations 
+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, Union, Literal
+
 import logging
+import math
 from dataclasses import dataclass
-import hydra
-from pydantic import BaseModel, Field
+from typing import Annotated, Literal
+
 import basico
+import hydra
 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 langchain_core.tools.base import BaseTool, InjectedToolCallId
+from langgraph.prebuilt import InjectedState
+from langgraph.types import Command
+from pydantic import BaseModel, Field
+
+from ..api.kegg import fetch_kegg_annotations
+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
-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'}
+ols_ontology_abbreviations = {"pato", "chebi", "sbo", "fma", "pr", "go"}
+
 
 def extract_relevant_species_names(model_object, arg_data, state):
     """
@@ -35,11 +39,12 @@ def extract_relevant_species_names(model_object, arg_data, state):
     """
     # Load hydra configuration
     with hydra.initialize(version_base=None, config_path="../configs"):
-        cfg = hydra.compose(config_name='config',
-                            overrides=['tools/get_annotation=default'])
+        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)
+    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)
@@ -53,17 +58,19 @@ def extract_relevant_species_names(model_object, arg_data, state):
         """
         A list of species based on user question.
         """
-        relevant_species: Union[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.""")
+
+        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']
+    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}'
+    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)
@@ -78,82 +85,95 @@ def extract_relevant_species_names(model_object, arg_data, state):
     logger.info("Extracted species: %s", extracted_species)
     return extracted_species
 
-def prepare_content_msg(species_without_description: List[str]):
+
+def prepare_content_msg(species_without_description: list[str]):
     """
     Prepare the content message.
     """
-    content = 'Successfully extracted annotations for the species.'
+    content = "Successfully extracted annotations for the species."
     if species_without_description:
-        content += f'''The descriptions for the following species
+        content += f"""The descriptions for the following species
                         were not found:
-                        {", ".join(species_without_description)}.'''
+                        {", ".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."]
+
+    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
+    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
+                        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:
+    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)
+        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
+        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)
+        print(list_species_names)
 
-        (annotations_df,
-         species_without_description) = self._fetch_annotations(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()
+            "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
@@ -161,24 +181,25 @@ class GetAnnotationTool(BaseTool):
         for key, value in {
             "model_id": [sys_bio_model.biomodel_id],
             "sbml_file_path": [sbml_file_path],
-            "dic_annotations_data": [dic_annotations_data]
+            "dic_annotations_data": [dic_annotations_data],
         }.items():
             if value:
                 dic_updated_state_for_model[key] = value
 
         return Command(
-            update=dic_updated_state_for_model | {
+            update=dic_updated_state_for_model
+            | {
                 "messages": [
                     ToolMessage(
                         content=prepare_content_msg(species_without_description),
                         artifact=True,
-                        tool_call_id=tool_call_id
+                        tool_call_id=tool_call_id,
                     )
                 ]
             }
         )
 
-    def _fetch_annotations(self, list_species_names: List[str]) -> tuple:
+    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.
@@ -210,11 +231,13 @@ class GetAnnotationTool(BaseTool):
 
             # 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"]
-                })
+                data.append(
+                    {
+                        "Species Name": species,
+                        "Link": desc["id"],
+                        "Qualifier": desc["qualifier"],
+                    }
+                )
 
         # Create a dataframe from the data list
         annotations_df = pd.DataFrame(data)
@@ -238,22 +261,23 @@ class GetAnnotationTool(BaseTool):
         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]
+        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]
+        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')
+        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["Description"] = annotations_df["Id"].apply(
+            lambda x: descriptions.get(x, "-")
+        )
         # annotations_df.index = annotations_df.index + 1
 
         # Reorder the columns
@@ -272,13 +296,13 @@ class GetAnnotationTool(BaseTool):
         Process link to format it correctly.
         """
         for ols_ontology_abbreviation in ols_ontology_abbreviations:
-            if ols_ontology_abbreviation +'/' in link:
+            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]:
+    def _fetch_descriptions(self, data: list[dict[str, str]]) -> dict[str, str]:
         """
         Fetch protein names or labels based on the database type.
         """
@@ -290,8 +314,8 @@ class GetAnnotationTool(BaseTool):
         # 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')
+            identifier = entry.get("Id")
+            database = entry.get("Database")
             # Check if database is NaN
             if isinstance(database, float):
                 if math.isnan(database):
@@ -306,18 +330,18 @@ class GetAnnotationTool(BaseTool):
         # Constants
 
         for database, identifiers in grouped_data.items():
-            if database == 'uniprot':
+            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
-                    ])
+                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]
+            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, "-")

aiagents4pharma/talk2biomodels/tools/get_modelinfo.py

@@ -5,26 +5,30 @@ Tool for get model information.
 """
 
 import logging
-from typing import Type, Optional, Annotated
 from dataclasses import dataclass
+from typing import Annotated
+
 import basico
-from pydantic import BaseModel, Field
-from langchain_core.tools import BaseTool
 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)
@@ -32,33 +36,38 @@ class RequestedModelInfo:
     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: Optional[ModelData] = None,
-             ) -> Command:
+    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.
 
@@ -71,14 +80,15 @@ class GetModelInfoTool(BaseTool):
         Returns:
             Command: The updated state of the tool.
         """
-        logger.log(logging.INFO,
-                   "Calling get_modelinfo tool %s, %s",
-                     sys_bio_model,
-                   requested_model_info)
+        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)
+        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:
@@ -87,15 +97,18 @@ class GetModelInfoTool(BaseTool):
                 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']]
+            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')
+            dic_results["Species"] = dic_results["Species"].to_dict(orient="records")
 
         # Extract parameters from the model
         if requested_model_info.parameters:
@@ -104,50 +117,43 @@ class GetModelInfoTool(BaseTool):
                 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']]
+            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')
+            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'])
+            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
+            dic_results["Description"] = model_obj.description
 
         # Extract description from the model
         if requested_model_info.name:
-            dic_results['Name'] = model_obj.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)
+            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():
+            "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
-                            )
-                        ],
-                    }
-            )
+            update=dic_updated_state_for_model
+            | {
+                # update the message history
+                "messages": [ToolMessage(content=dic_results, tool_call_id=tool_call_id)],
+            }
+        )