aiagents4pharma 1.43.0__py3-none-any.whl → 1.45.0__py3-none-any.whl

aiagents4pharma/talk2biomodels/tools/load_arguments.py

@@ -8,45 +8,52 @@ before the experiment is run.
 
 import logging
 from dataclasses import dataclass
-from typing import Union, List, Optional, Annotated
-from pydantic import Field
+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: Union[int, float] = Field(
-        description="Duration of the simulation",
-        default=100)
-    interval: Union[int, float] = Field(
+
+    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)
+        default=100,
+    )
+
 
 @dataclass
 class SpeciesInitialData:
     """
     Dataclass for storing the species initial data.
     """
-    species_name: List[str] = Field(
+
+    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[Union[int, float]] = Field(
+        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=[])
+        default=[],
+    )
+
 
 @dataclass
 class TimeSpeciesNameConcentration:
@@ -54,10 +61,13 @@ class TimeSpeciesNameConcentration:
     Dataclass for storing the time,
     species name, and concentration data.
     """
-    time: Union[int, float] = Field(description="time point where the event occurs")
+
+    time: int | float = Field(description="time point where the event occurs")
     species_name: str = Field(description="species name")
-    species_concentration: Union[int, float] = Field(
-        description="species concentration at the time point")
+    species_concentration: int | float = Field(
+        description="species concentration at the time point"
+    )
+
 
 @dataclass
 class ReocurringData:
@@ -66,34 +76,42 @@ class ReocurringData:
     the concentration of the species resets to a certain
     value after a certain time interval.
     """
-    data: List[TimeSpeciesNameConcentration] = Field(
+
+    data: list[TimeSpeciesNameConcentration] = Field(
         description="List of time, name, and concentration data"
-                    " of species or parameters that reoccur",
-                    default=[])
+        " 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: Optional[TimeData] = Field(
-        description="time data",
-        default=None)
-    species_to_be_analyzed_before_experiment: Optional[SpeciesInitialData] = Field(
+
+    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: Optional[ReocurringData] = Field(
+        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)
+        default=None,
+    )
+
 
 def add_rec_events(model_object, reocurring_data):
     """
@@ -108,7 +126,9 @@ def add_rec_events(model_object, reocurring_data):
     """
     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)
+        basico.add_event(
+            f"{sn}_{tp}",
+            f"Time > {tp}",
+            [[sn, str(sc)]],
+            model=model_object.copasi_model,
+        )

aiagents4pharma/talk2biomodels/tools/load_biomodel.py

@@ -4,10 +4,13 @@
 Function for loading the BioModel.
 """
 
-from typing import Annotated, Any, Union
+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'.
@@ -18,14 +21,17 @@ def ensure_biomodel_id(value: Any) -> Any:
         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[Union[int, str], BeforeValidator(ensure_biomodel_id)] = None
+
+    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.

aiagents4pharma/talk2biomodels/tools/parameter_scan.py

@@ -6,69 +6,79 @@ Tool for parameter scan.
 
 import logging
 from dataclasses import dataclass
-from typing import Type, Union, List, Annotated, Optional
-import pandas as pd
+from typing import Annotated
+
 import basico
-from pydantic import BaseModel, Field
-from langgraph.types import Command
-from langgraph.prebuilt import InjectedState
-from langchain_core.tools import BaseTool
+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 .load_arguments import TimeData, SpeciesInitialData
 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_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[Union[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)
+        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: Optional[SpeciesInitialData] = 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)
+    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.")
+        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):
     """
@@ -86,24 +96,26 @@ def make_list_dic_scanned_data(dic_param_scan, arg_data, sys_bio_model, tool_cal
     """
     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)
+        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()
-        })
+        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:
+
+def run_parameter_scan(model_object, arg_data, dic_species_data, duration, interval) -> dict:
     """
     Run parameter scan on the model.
 
@@ -128,16 +140,18 @@ def run_parameter_scan(model_object,
 
     # 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()
+    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)
+            arg_data.parameter_scan_data.species_parameter_name,
+        )
         raise ValueError(
             "Invalid species or parameter name: "
-            f"{arg_data.parameter_scan_data.species_parameter_name}.")
+            f"{arg_data.parameter_scan_data.species_parameter_name}."
+        )
 
     # Dictionary to store the parameter scan results
     dic_param_scan_results = {}
@@ -165,12 +179,13 @@ def run_parameter_scan(model_object,
         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})
+                {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']
+            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]
@@ -182,37 +197,42 @@ def run_parameter_scan(model_object,
     # 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
+
+    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)
+        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
+    args_schema: type[BaseModel] = ParameterScanInput
 
-    def _run(self,
+    def _run(
+        self,
         tool_call_id: Annotated[str, InjectedToolCallId],
         state: Annotated[dict, InjectedState],
         sys_bio_model: ModelData = None,
-        arg_data: ArgumentData = None
+        arg_data: ArgumentData = None,
     ) -> Command:
         """
         Run the tool.
@@ -226,11 +246,9 @@ class ParameterScanTool(BaseTool):
         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)
+        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
@@ -243,9 +261,10 @@ class ParameterScanTool(BaseTool):
                 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
-                        )
+                        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:
@@ -257,37 +276,35 @@ class ParameterScanTool(BaseTool):
                 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)
+        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)
+        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():
+        }.items():
             if value:
                 dic_updated_state_for_model[key] = value
         # Return the updated state
         return Command(
-                update=dic_updated_state_for_model|{
+            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)
-                        )
-                    ],
-                }
-            )
+                        artifact=get_model_units(model_object),
+                    )
+                ],
+            }
+        )

aiagents4pharma/talk2biomodels/tools/query_article.py

@@ -5,58 +5,59 @@ Tool for asking questions to the article.
 """
 
 import logging
-from typing import Type, Annotated
-from pydantic import BaseModel, Field
+from typing import Annotated
+
+from langchain_community.document_loaders import PyPDFLoader
 from langchain_core.tools import BaseTool
 from langchain_core.vectorstores import InMemoryVectorStore
-from langchain_community.document_loaders import PyPDFLoader
 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
+    args_schema: type[BaseModel] = QueryArticleInput
 
-    def _run(self,
-             question: str,
-             state: Annotated[dict, InjectedState]):
+    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, "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'])
+        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']
+        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)
+        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