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

aiagents4pharma/talk2biomodels/tools/search_models.py

@@ -4,46 +4,51 @@
 Tool for searching models based on search query.
 """
 
-from typing import Type, Annotated
 import logging
-from pydantic import BaseModel, Field
+from typing import Annotated
+
 import pandas as pd
 from basico import biomodels
-from langgraph.types import Command
-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.types import Command
+from pydantic import BaseModel, Field
 
 # Initialize logger
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+
 class SearchModelsInput(BaseModel):
     """
     Input schema for the search models tool.
     """
+
     query: str = Field(description="Search models query", default=None)
-    num_query: int = Field(description="Top number of models to search",
-                           default=10,
-                           le=100)
+    num_query: int = Field(description="Top number of models to search", default=10, le=100)
     tool_call_id: Annotated[str, InjectedToolCallId]
 
+
 # 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 SearchModelsTool(BaseTool):
     """
     Tool for returning the search results based on the search query.
     """
+
     name: str = "search_models"
     description: str = "Search for only manually curated models in "
     "the BioMmodels database based on keywords."
-    args_schema: Type[BaseModel] = SearchModelsInput
+    args_schema: type[BaseModel] = SearchModelsInput
     return_direct: bool = False
 
-    def _run(self,
-             tool_call_id: Annotated[str, InjectedToolCallId],
-             query: str = None,
-             num_query: int = 10) -> dict:
+    def _run(
+        self,
+        tool_call_id: Annotated[str, InjectedToolCallId],
+        query: str = None,
+        num_query: int = 10,
+    ) -> dict:
         """
         Run the tool.
 
@@ -55,8 +60,12 @@ class SearchModelsTool(BaseTool):
         Returns:
             dict: The answer to the question in the form of a dictionary.
         """
-        logger.log(logging.INFO, "Searching models with the query and number %s, %s",
-                   query, num_query)
+        logger.log(
+            logging.INFO,
+            "Searching models with the query and number %s, %s",
+            query,
+            num_query,
+        )
         # Search for models based on the query
         search_results = biomodels.search_for_model(query, num_results=num_query)
         # Convert the search results to a pandas DataFrame
@@ -69,17 +78,19 @@ class SearchModelsTool(BaseTool):
         # to avoid hallucinations
         content += f" Here is the summary of the first {first_n} models:"
         for i in range(first_n):
-            content += f"\nModel {i+1}: {search_results[i]['name']} (ID: {search_results[i]['id']})"
+            content += (
+                f"\nModel {i + 1}: {search_results[i]['name']} (ID: {search_results[i]['id']})"
+            )
         # Return the updated state of the tool
         return Command(
-                update={
-                    # update the message history
-                    "messages": [
-                        ToolMessage(
-                            content=content,
-                            tool_call_id=tool_call_id,
-                            artifact={'dic_data': df.to_dict(orient='records')}
-                            )
-                        ],
-                    }
-            )
+            update={
+                # update the message history
+                "messages": [
+                    ToolMessage(
+                        content=content,
+                        tool_call_id=tool_call_id,
+                        artifact={"dic_data": df.to_dict(orient="records")},
+                    )
+                ],
+            }
+        )

aiagents4pharma/talk2biomodels/tools/simulate_model.py

@@ -5,49 +5,56 @@ Tool for simulating a model.
 """
 
 import logging
-from typing import Type, Annotated
-from pydantic import BaseModel, Field
-from langgraph.types import Command
-from langgraph.prebuilt import InjectedState
-from langchain_core.tools import BaseTool
+from typing import Annotated
+
 from langchain_core.messages import ToolMessage
+from langchain_core.tools import BaseTool
 from langchain_core.tools.base import InjectedToolCallId
-from .load_biomodel import ModelData, load_biomodel
+from langgraph.prebuilt import InjectedState
+from langgraph.types import Command
+from pydantic import BaseModel, Field
+
 from .load_arguments import ArgumentData, add_rec_events
+from .load_biomodel import ModelData, load_biomodel
 from .utils import get_model_units
 
 # Initialize logger
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+
 class SimulateModelInput(BaseModel):
     """
     Input schema for the SimulateModel 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 simulation name""",
-                                   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 SimulateModelTool(BaseTool):
     """
     Tool for simulating a model.
     """
+
     name: str = "simulate_model"
     description: str = "A tool to simulate a biomodel"
-    args_schema: Type[BaseModel] = SimulateModelInput
+    args_schema: type[BaseModel] = SimulateModelInput
 
-    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.
@@ -61,13 +68,9 @@ class SimulateModelTool(BaseTool):
         Returns:
             str: The result of the simulation.
         """
-        logger.log(logging.INFO,
-                   "Calling simulate_model 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 simulate_model 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
@@ -78,8 +81,12 @@ class SimulateModelTool(BaseTool):
             # Prepare the dictionary of species data
             if arg_data.species_to_be_analyzed_before_experiment is not None:
                 dic_species_to_be_analyzed_before_experiment = dict(
-                    zip(arg_data.species_to_be_analyzed_before_experiment.species_name,
-                        arg_data.species_to_be_analyzed_before_experiment.species_concentration))
+                    zip(
+                        arg_data.species_to_be_analyzed_before_experiment.species_name,
+                        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:
                 add_rec_events(model_object, arg_data.reocurring_data)
@@ -89,17 +96,19 @@ class SimulateModelTool(BaseTool):
                 interval = arg_data.time_data.interval
         # Update the model parameters
         model_object.update_parameters(dic_species_to_be_analyzed_before_experiment)
-        logger.log(logging.INFO,
-                   "Following species/parameters updated in the model %s",
-                   dic_species_to_be_analyzed_before_experiment)
+        logger.log(
+            logging.INFO,
+            "Following species/parameters updated in the model %s",
+            dic_species_to_be_analyzed_before_experiment,
+        )
         # Simulate the model
         df = model_object.simulate(duration=duration, interval=interval)
         logger.log(logging.INFO, "Simulation results ready with shape %s", df.shape)
         dic_simulated_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': 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": df.to_dict(),
         }
         # Prepare the dictionary of updated state
         dic_updated_state_for_model = {}
@@ -107,19 +116,20 @@ class SimulateModelTool(BaseTool):
             "model_id": [sys_bio_model.biomodel_id],
             "sbml_file_path": [sbml_file_path],
             "dic_simulated_data": [dic_simulated_data],
-            }.items():
+        }.items():
             if value:
                 dic_updated_state_for_model[key] = value
         # Return the updated state of the tool
         return Command(
-                update=dic_updated_state_for_model|{
+            update=dic_updated_state_for_model
+            | {
                 # update the message history
                 "messages": [
                     ToolMessage(
                         content=f"Simulation 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/steady_state.py

@@ -5,23 +5,25 @@ Tool for parameter scan.
 """
 
 import logging
-from typing import Type, Annotated
+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
 from langchain_core.messages import ToolMessage
+from langchain_core.tools import BaseTool
 from langchain_core.tools.base import InjectedToolCallId
-from .load_biomodel import ModelData, load_biomodel
+from langgraph.prebuilt import InjectedState
+from langgraph.types import Command
+from pydantic import BaseModel, Field
+
 from .load_arguments import ArgumentData, add_rec_events
+from .load_biomodel import ModelData, load_biomodel
 
 # Initialize logger
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-def run_steady_state(model_object,
-                     dic_species_to_be_analyzed_before_experiment):
+
+def run_steady_state(model_object, dic_species_to_be_analyzed_before_experiment):
     """
     Run the steady state analysis.
 
@@ -47,58 +49,65 @@ def run_steady_state(model_object,
     df_steady_state = basico.model_info.get_species(model=model_object.copasi_model).reset_index()
     # print (df_steady_state)
     # Rename the column name to species_name
-    df_steady_state.rename(columns={'name': 'species_name'},
-                           inplace=True)
+    df_steady_state.rename(columns={"name": "species_name"}, inplace=True)
     # Rename the column concentration to steady_state_concentration
-    df_steady_state.rename(columns={'concentration': 'steady_state_concentration'},
-                           inplace=True)
+    df_steady_state.rename(columns={"concentration": "steady_state_concentration"}, inplace=True)
     # Rename the column transition_time to steady_state_transition_time
-    df_steady_state.rename(columns={'transition_time': 'steady_state_transition_time'},
-                           inplace=True)
+    df_steady_state.rename(
+        columns={"transition_time": "steady_state_transition_time"}, inplace=True
+    )
     # Drop some columns
-    df_steady_state.drop(columns=
-                         [
-                            'initial_particle_number',
-                            'initial_expression',
-                            'expression',
-                            'particle_number',
-                            'type',
-                            'particle_number_rate',
-                            'key',
-                            'sbml_id',
-                            'display_name'],
-                            inplace=True)
+    df_steady_state.drop(
+        columns=[
+            "initial_particle_number",
+            "initial_expression",
+            "expression",
+            "particle_number",
+            "type",
+            "particle_number_rate",
+            "key",
+            "sbml_id",
+            "display_name",
+        ],
+        inplace=True,
+    )
     logger.log(logging.INFO, "Steady state results with shape %s", df_steady_state.shape)
     return df_steady_state
 
+
 class SteadyStateInput(BaseModel):
     """
     Input schema for the steady state tool.
     """
-    sys_bio_model: ModelData = Field(description="model data",
-                                     default=None)
+
+    sys_bio_model: ModelData = Field(description="model data", default=None)
     arg_data: ArgumentData = Field(
         description="time, species, and reocurring data"
-                " that must be set before the steady state analysis"
-                " as well as the experiment name", default=None)
+        " that must be set before the steady state analysis"
+        " as well as the experiment name",
+        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 SteadyStateTool(BaseTool):
     """
     Tool to bring a model to steady state.
     """
+
     name: str = "steady_state"
     description: str = "A tool to bring a model to steady state."
-    args_schema: Type[BaseModel] = SteadyStateInput
+    args_schema: type[BaseModel] = SteadyStateInput
 
-    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.
@@ -112,12 +121,15 @@ class SteadyStateTool(BaseTool):
         Returns:
             Command: The updated state of the tool.
         """
-        logger.log(logging.INFO, "Calling the steady_state tool %s, %s",
-                   sys_bio_model, arg_data)
+        logger.log(
+            logging.INFO,
+            "Calling the steady_state tool %s, %s",
+            sys_bio_model,
+            arg_data,
+        )
         # print (f'Calling steady_state tool {sys_bio_model}, {arg_data}, {tool_call_id}')
-        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)
+        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
@@ -126,44 +138,50 @@ class SteadyStateTool(BaseTool):
             # Prepare the dictionary of species data
             if arg_data.species_to_be_analyzed_before_experiment is not None:
                 dic_species_to_be_analyzed_before_experiment = dict(
-                    zip(arg_data.species_to_be_analyzed_before_experiment.species_name,
-                        arg_data.species_to_be_analyzed_before_experiment.species_concentration))
+                    zip(
+                        arg_data.species_to_be_analyzed_before_experiment.species_name,
+                        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:
                 add_rec_events(model_object, arg_data.reocurring_data)
         # Run the parameter scan
-        df_steady_state = run_steady_state(model_object,
-                                           dic_species_to_be_analyzed_before_experiment)
-        print (df_steady_state)
+        df_steady_state = run_steady_state(
+            model_object, dic_species_to_be_analyzed_before_experiment
+        )
+        print(df_steady_state)
         # Prepare the dictionary of scanned data
         # that will be passed to the state of the graph
         dic_steady_state_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': df_steady_state.to_dict(orient='records')
+            "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": df_steady_state.to_dict(orient="records"),
         }
         # 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_steady_state_data": [dic_steady_state_data]
-            }.items():
+            "dic_steady_state_data": [dic_steady_state_data],
+        }.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(
+                    ToolMessage(
                         content=f"Steady state analysis of"
-                                f" {arg_data.experiment_name}"
-                                " was successful.",
+                        f" {arg_data.experiment_name}"
+                        " was successful.",
                         tool_call_id=tool_call_id,
-                        artifact={'dic_data': df_steady_state.to_dict(orient='records')}
-                        )
-                    ],
-                }
-            )
+                        artifact={"dic_data": df_steady_state.to_dict(orient="records")},
+                    )
+                ],
+            }
+        )

aiagents4pharma/talk2biomodels/tools/utils.py

@@ -6,6 +6,7 @@ Utility functions for T2B tools.
 
 import basico
 
+
 def get_model_units(model_object):
     """
     Get the units of the model.
@@ -17,6 +18,6 @@ def get_model_units(model_object):
         dict: The units of the model.
     """
     model_units = basico.model_info.get_model_units(model=model_object.copasi_model)
-    model_units_y = model_units['quantity_unit']
-    model_units_x = model_units['time_unit']
-    return {'y_axis_label': model_units_y, 'x_axis_label': model_units_x}
+    model_units_y = model_units["quantity_unit"]
+    model_units_x = model_units["time_unit"]
+    return {"y_axis_label": model_units_y, "x_axis_label": model_units_x}

aiagents4pharma/talk2cells/README.md

@@ -0,0 +1 @@
+Please check out the README file in the root folder for more information.

aiagents4pharma/talk2cells/__init__.py

@@ -1,6 +1,5 @@
-'''
+"""
 This file is used to import the models and tools.
-'''
-from . import agents
-from . import states
-from . import tools
+"""
+
+from . import agents, states, tools

aiagents4pharma/talk2cells/agents/__init__.py

@@ -1,5 +1,6 @@
-'''
+"""
 This file is used to import all the modules in the package.
-'''
+"""
+
 # import everything from the module
 from . import scp_agent

aiagents4pharma/talk2cells/agents/scp_agent.py

@@ -1,31 +1,35 @@
-#/usr/bin/env python3
+# /usr/bin/env python3
 
-'''
+"""
 This is the agent file for the Talk2Cells graph.
-'''
+"""
 
 import logging
 import os
+
 from langchain_openai import ChatOpenAI
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import START, StateGraph
-from langgraph.prebuilt import create_react_agent, ToolNode
-from ..tools.scp_agent.search_studies import search_studies
-from ..tools.scp_agent.display_studies import display_studies
+from langgraph.prebuilt import ToolNode, create_react_agent
+
 from ..states.state_talk2cells import Talk2Cells
+from ..tools.scp_agent.display_studies import display_studies
+from ..tools.scp_agent.search_studies import search_studies
 
 # Initialize logger
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+
 def get_app(uniq_id):
-    '''
+    """
     This function returns the langraph app.
-    '''
+    """
+
     def agent_scp_node(state: Talk2Cells):
-        '''
+        """
         This function calls the model.
-        '''
+        """
         logger.log(logging.INFO, "Creating SCP_Agent node with thread_id %s", uniq_id)
         # Get the messages from the state
         # messages = state['messages']
@@ -48,19 +52,17 @@ def get_app(uniq_id):
     # Create an environment variable to store the LLM model
     # Check if the environment variable AIAGENTS4PHARMA_LLM_MODEL is set
     # If not, set it to 'gpt-4o-mini'
-    llm_model = os.getenv('AIAGENTS4PHARMA_LLM_MODEL', 'gpt-4o-mini')
+    llm_model = os.getenv("AIAGENTS4PHARMA_LLM_MODEL", "gpt-4o-mini")
     # print (f'LLM model: {llm_model}')
     # llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
     llm = ChatOpenAI(model=llm_model, temperature=0)
     model = create_react_agent(
-                            llm,
-                            tools=tools,
-                            state_schema=Talk2Cells,
-                            state_modifier=(
-                                            "You are Talk2Cells agent."
-                                            ),
-                            checkpointer=MemorySaver()
-                        )
+        llm,
+        tools=tools,
+        state_schema=Talk2Cells,
+        state_modifier=("You are Talk2Cells agent."),
+        checkpointer=MemorySaver(),
+    )
 
     # Define a new graph
     workflow = StateGraph(Talk2Cells)

aiagents4pharma/talk2cells/states/__init__.py

@@ -1,5 +1,6 @@
-'''
+"""
 This file is used to import all the modules in the package.
-'''
+"""
+
 # import everything from the module
 from . import state_talk2cells

aiagents4pharma/talk2cells/states/state_talk2cells.py

@@ -1,13 +1,15 @@
 #!/usr/bin/env python3
 
-'''
+"""
 This is the state file for the Talk2Cells agent.
-'''
+"""
 
 from langgraph.prebuilt.chat_agent_executor import AgentState
 
+
 class Talk2Cells(AgentState):
     """
     The state for the Talk2Cells agent.
     """
+
     search_table: str

aiagents4pharma/talk2cells/tests/scp_agent/test_scp_agent.py

@@ -1,23 +1,22 @@
-'''
+"""
 Test cases for the search_studies
-'''
+"""
 
 # from ..tools.search_studies import search_studies
-from aiagents4pharma.talk2cells.agents.scp_agent import get_app
 from langchain_core.messages import HumanMessage
 
+from aiagents4pharma.talk2cells.agents.scp_agent import get_app
+
+
 def test_agent_scp():
-    '''
+    """
     Test the agent_scp.
-    '''
+    """
     unique_id = 12345
     app = get_app(unique_id)
     config = {"configurable": {"thread_id": unique_id}}
     prompt = "Search for studies on Crohns Disease."
-    response = app.invoke(
-                        {"messages": [HumanMessage(content=prompt)]},
-                        config=config
-                    )
+    response = app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
     assistant_msg = response["messages"][-1].content
     # Check if the assistant message is a string
     assert isinstance(assistant_msg, str)