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

aiagents4pharma/talk2biomodels/tools/ask_question.py

@@ -5,31 +5,38 @@ Tool for asking a question about the simulation results.
 """
 
 import logging
-from typing import Type, Annotated, Literal
-import hydra
+from typing import Annotated, Literal
+
 import basico
+import hydra
 import pandas as pd
-from pydantic import BaseModel, Field
 from langchain_core.tools.base import BaseTool
 from langchain_experimental.agents import create_pandas_dataframe_agent
 from langgraph.prebuilt import InjectedState
+from pydantic import BaseModel, Field
 
 # Initialize logger
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+
 class AskQuestionInput(BaseModel):
     """
     Input schema for the AskQuestion tool.
     """
+
     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.""")
+    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")
+        description="Context of the question"
+    )
     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.
@@ -37,17 +44,20 @@ class AskQuestionTool(BaseTool):
     """
     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 or steady state results."""
-    args_schema: Type[BaseModel] = AskQuestionInput
+    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:
+    def _run(
+        self,
+        question: str,
+        experiment_name: str,
+        question_context: Literal["simulation", "steady_state"],
+        state: Annotated[dict, InjectedState],
+    ) -> str:
         """
         Run the tool.
 
@@ -59,15 +69,16 @@ class AskQuestionTool(BaseTool):
         Returns:
             str: The answer to the question.
         """
-        logger.log(logging.INFO,
-                   "Calling ask_question tool %s, %s, %s",
-                   question,
-                   question_context,
-                   experiment_name)
+        logger.log(
+            logging.INFO,
+            "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=['tools/ask_question=default'])
+            cfg = hydra.compose(config_name="config", overrides=["tools/ask_question=default"])
             cfg = cfg.tools.ask_question
         # Get the context of the question
         # and based on the context, get the data
@@ -89,9 +100,7 @@ class AskQuestionTool(BaseTool):
         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]
-        )
+        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()
@@ -100,15 +109,16 @@ class AskQuestionTool(BaseTool):
         prompt_content += f"{basico.model_info.get_model_units()}\n\n"
         # Create a pandas dataframe agent
         df_agent = create_pandas_dataframe_agent(
-                        state['llm_model'],
-                        allow_dangerous_code=True,
-                        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)
+            state["llm_model"],
+            allow_dangerous_code=True,
+            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, stream_mode=None)
         # print (llm_result)

aiagents4pharma/talk2biomodels/tools/custom_plotter.py

@@ -5,13 +5,15 @@ Tool for plotting a custom y-axis of a simulation plot.
 """
 
 import logging
-from typing import Type, Annotated, List, Tuple, Union, Literal
-from pydantic import BaseModel, Field
+from typing import Annotated, Literal
+
 import hydra
 import pandas as pd
-from langchain_core.tools import BaseTool
 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
 
@@ -19,6 +21,7 @@ from .utils import get_model_units
 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.
@@ -31,6 +34,7 @@ def extract_relevant_species(question, species_names, state):
     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
@@ -41,44 +45,46 @@ def extract_relevant_species(question, species_names, state):
 
         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: Union[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."
-                )
+
+        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 = 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']
+    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}")])
+    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)
+    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.
@@ -89,21 +95,23 @@ 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
+    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
+                        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, Union[None, List[str]]]:
+    def _run(
+        self,
+        question: str,
+        sys_bio_model: ModelData,
+        simulation_name: str,
+        state: Annotated[dict, InjectedState],
+    ) -> tuple[str, None | list[str]]:
         """
         Run the tool.
 
@@ -118,7 +126,7 @@ class CustomPlotterTool(BaseTool):
         """
         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
+        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"]:
@@ -129,20 +137,20 @@ class CustomPlotterTool(BaseTool):
         # 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(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')
+        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)
+        print(results)
         if results.relevant_species is None:
-            raise ValueError("No species found in the simulation results \
-                             that matches the user prompt.")
+            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
@@ -151,7 +159,7 @@ class CustomPlotterTool(BaseTool):
                 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)
+        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)