aiagents4pharma 1.6.1__tar.gz → 1.7.0__tar.gz

{aiagents4pharma-1.6.1 → aiagents4pharma-1.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: aiagents4pharma
-Version: 1.6.1
+Version: 1.7.0
 Summary: AI Agents for drug discovery, drug development, and other pharmaceutical R&D
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
@@ -13,6 +13,7 @@ Requires-Dist: coverage==7.6.4
 Requires-Dist: einops==0.8.0
 Requires-Dist: gdown==5.2.0
 Requires-Dist: huggingface_hub==0.26.5
+Requires-Dist: hydra-core==1.3.2
 Requires-Dist: joblib==1.4.2
 Requires-Dist: langchain==0.3.7
 Requires-Dist: langchain-community==0.3.5
@@ -43,7 +44,10 @@ Requires-Dist: mkdocs-include-markdown-plugin==7.1.2
 Requires-Dist: mkdocstrings==0.27.0
 Requires-Dist: streamlit-feedback
 
-[![TESTS](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests.yml/badge.svg?branch=feat%2Finitial-setup)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests.yml)
+[![Talk2BioModels](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2biomodels.yml/badge.svg)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2biomodels.yml)
+[![Talk2Cells](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2cells.yml/badge.svg)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2cells.yml)
+[![Talk2KnowledgeGraphs](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2knowledgegraphs.yml/badge.svg)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2knowledgegraphs.yml)
+[![Talk2Competitors](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2competitors.yml/badge.svg)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2competitors.yml)
 
 <h1 align="center" style="border-bottom: none;">🤖 AIAgents4Pharma</h1>
 
@@ -54,6 +58,7 @@ Our toolkit currently consists of three intelligent agents, each designed to sim
 - **Talk2BioModels**: Engage directly with mathematical models in systems biology.
 - **Talk2Cells** *(Work in progress)*: Query and analyze sequencing data with ease.
 - **Talk2KnowledgeGraphs** *(Work in progress)*: Access and explore complex biological knowledge graphs for insightful data connections.
+- **Talk2Competitors** *(Coming soon)*: Get recommendations for articles related to your choice. Download, query, and write/retrieve them to your reference manager (currently supporting Zotero).
 
 ---
 
@@ -75,6 +80,8 @@ Our toolkit currently consists of three intelligent agents, each designed to sim
 
 **Talk2KnowledgeGraphs** is an agent designed to enable interaction with biological knowledge graphs (KGs). KGs integrate vast amounts of structured biological data into a format that highlights relationships between entities, such as proteins, genes, and diseases.
 
+### 4. Talk2KnowledgeGraphs *(Coming soon)*
+
 ## Getting Started
 
 ### Prerequisites
@@ -173,6 +180,7 @@ Check out our [CONTRIBUTING.md](CONTRIBUTING.md) for more information.
 - **User Interface**: Interactive web UI for all agents.
 - **Talk2Cells**: Integration of sequencing data analysis tools.
 - **Talk2KnowledgeGraphs**: Interface for biological knowledge graph interaction.
+- **Talk2Competitors**
 
 We’re excited to bring AIAgents4Pharma to the bioinformatics and pharmaceutical research community. Together, let’s make data-driven biological research more accessible and insightful. 
 

{aiagents4pharma-1.6.1 → aiagents4pharma-1.7.0}/README.md

@@ -1,4 +1,7 @@
-[![TESTS](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests.yml/badge.svg?branch=feat%2Finitial-setup)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests.yml)
+[![Talk2BioModels](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2biomodels.yml/badge.svg)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2biomodels.yml)
+[![Talk2Cells](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2cells.yml/badge.svg)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2cells.yml)
+[![Talk2KnowledgeGraphs](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2knowledgegraphs.yml/badge.svg)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2knowledgegraphs.yml)
+[![Talk2Competitors](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2competitors.yml/badge.svg)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2competitors.yml)
 
 <h1 align="center" style="border-bottom: none;">🤖 AIAgents4Pharma</h1>
 
@@ -9,6 +12,7 @@ Our toolkit currently consists of three intelligent agents, each designed to sim
 - **Talk2BioModels**: Engage directly with mathematical models in systems biology.
 - **Talk2Cells** *(Work in progress)*: Query and analyze sequencing data with ease.
 - **Talk2KnowledgeGraphs** *(Work in progress)*: Access and explore complex biological knowledge graphs for insightful data connections.
+- **Talk2Competitors** *(Coming soon)*: Get recommendations for articles related to your choice. Download, query, and write/retrieve them to your reference manager (currently supporting Zotero).
 
 ---
 
@@ -30,6 +34,8 @@ Our toolkit currently consists of three intelligent agents, each designed to sim
 
 **Talk2KnowledgeGraphs** is an agent designed to enable interaction with biological knowledge graphs (KGs). KGs integrate vast amounts of structured biological data into a format that highlights relationships between entities, such as proteins, genes, and diseases.
 
+### 4. Talk2KnowledgeGraphs *(Coming soon)*
+
 ## Getting Started
 
 ### Prerequisites
@@ -128,6 +134,7 @@ Check out our [CONTRIBUTING.md](CONTRIBUTING.md) for more information.
 - **User Interface**: Interactive web UI for all agents.
 - **Talk2Cells**: Integration of sequencing data analysis tools.
 - **Talk2KnowledgeGraphs**: Interface for biological knowledge graph interaction.
+- **Talk2Competitors**
 
 We’re excited to bring AIAgents4Pharma to the bioinformatics and pharmaceutical research community. Together, let’s make data-driven biological research more accessible and insightful. 
 

{aiagents4pharma-1.6.1 → aiagents4pharma-1.7.0}/aiagents4pharma/talk2biomodels/models/basico_model.py

@@ -5,20 +5,23 @@ BasicoModel class for loading and simulating SBML models
 using the basico package.
 """
 
+import logging
 from typing import Optional, Dict, Union
-from time import sleep
-from urllib.error import URLError
 from pydantic import Field, model_validator
 import pandas as pd
 import basico
 from .sys_bio_model import SysBioModel
 
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 class BasicoModel(SysBioModel):
     """
     Model that loads and simulates SBML models using the basico package.
-    Can load models from an SBML file or download them using a BioModels model_id.
+    Can load models from an SBML file or download them using a BioModels biomodel_id.
     """
-    model_id: Optional[int] = Field(None, description="BioModels model ID to download and load")
+    biomodel_id: Optional[int] = Field(None, description="BioModels model ID to download and load")
     sbml_file_path: Optional[str] = Field(None, description="Path to an SBML file to load")
     simulation_results: Optional[str] = None
     name: Optional[str] = Field("", description="Name of the model")
@@ -28,27 +31,21 @@ class BasicoModel(SysBioModel):
     copasi_model: Optional[object] = None  # Holds the loaded Copasi model
 
     @model_validator(mode="after")
-    def check_model_id_or_sbml_file_path(self):
+    def check_biomodel_id_or_sbml_file_path(self):
         """
-        Validate that either model_id or sbml_file_path is provided.
+        Validate that either biomodel_id or sbml_file_path is provided.
         """
-        if not self.model_id and not self.sbml_file_path:
-            raise ValueError("Either model_id or sbml_file_path must be provided.")
-        if self.model_id:
-            attempts = 0
-            max_retries = 5
-            while attempts < max_retries:
-                try:
-                    self.copasi_model = basico.load_biomodel(self.model_id)
-                    break
-                except URLError as e:
-                    attempts += 1
-                    sleep(10*attempts)
-                    if attempts >= max_retries:
-                        raise e
-            self.description = basico.biomodels.get_model_info(self.model_id)["description"]
+        if not self.biomodel_id and not self.sbml_file_path:
+            logger.error("Either biomodel_id or sbml_file_path must be provided.")
+            raise ValueError("Either biomodel_id or sbml_file_path must be provided.")
+        if self.biomodel_id:
+            self.copasi_model = basico.load_biomodel(self.biomodel_id)
+            self.description = basico.biomodels.get_model_info(self.biomodel_id)["description"]
+            self.name = basico.model_info.get_model_name(model=self.copasi_model)
         elif self.sbml_file_path:
             self.copasi_model = basico.load_model(self.sbml_file_path)
+            self.description = basico.model_info.get_notes(model=self.copasi_model)
+            self.name = basico.model_info.get_model_name(model=self.copasi_model)
         return self
 
     def simulate(self,
@@ -92,10 +89,16 @@ class BasicoModel(SysBioModel):
         df_result = basico.run_time_course(model=self.copasi_model,
                                         intervals=interval,
                                         duration=duration)
+        # Replace curly braces in column headers with square brackets
+        # Because curly braces in the world of LLMS are used for
+        # structured output
         df_result.columns = df_result.columns.str.replace('{', '[', regex=False).\
                     str.replace('}', ']', regex=False)
+        # Reset the index
         df_result.reset_index(inplace=True)
+        # Store the simulation results
         self.simulation_results = df_result
+        # Return copy of the simulation results
         return df_result.copy()
 
     def get_model_metadata(self) -> Dict[str, Union[str, int]]:

{aiagents4pharma-1.6.1 → aiagents4pharma-1.7.0}/aiagents4pharma/talk2biomodels/models/sys_bio_model.py

@@ -12,18 +12,18 @@ class SysBioModel(ABC, BaseModel):
     This class serves as a general structure for models, allowing
     different mathematical approaches to be implemented in subclasses.
     """
-    model_id: Optional[int] = Field(None, description="BioModel ID of the model")
+    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")
 
     @model_validator(mode="after")
-    def check_model_id_or_sbml_file_path(self):
+    def check_biomodel_id_or_sbml_file_path(self):
         """
-        Validate that either model_id or sbml_file_path is provided.
+        Validate that either biomodel_id or sbml_file_path is provided.
         """
-        if not self.model_id and not self.sbml_file_path:
-            raise ValueError("Either model_id or sbml_file_path must be provided.")
+        if not self.biomodel_id and not self.sbml_file_path:
+            raise ValueError("Either biomodel_id or sbml_file_path must be provided.")
         return self
 
     @abstractmethod

{aiagents4pharma-1.6.1 → aiagents4pharma-1.7.0}/aiagents4pharma/talk2biomodels/tools/__init__.py

@@ -1,10 +1,8 @@
 '''
 This file is used to import all the modules in the package.
 '''
-# import everything from the module
-from . import ask_question
+from . import search_models
 from . import simulate_model
+from . import ask_question
 from . import custom_plotter
-from . import fetch_parameters
-from . import model_description
-from . import search_models
+from . import get_modelinfo

aiagents4pharma-1.7.0/aiagents4pharma/talk2biomodels/tools/ask_question.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+
+"""
+Tool for asking a question about the simulation results.
+"""
+
+import logging
+from typing import Type, Annotated
+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
+
+# 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 results")
+    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 AskQuestionTool(BaseTool):
+    """
+    Tool for calculating the product of two numbers.
+    """
+    name: str = "ask_question"
+    description: str = "A tool to ask question about the simulation results."
+    args_schema: Type[BaseModel] = AskQuestionInput
+    return_direct: bool = False
+
+    def _run(self,
+             question: str,
+             state: Annotated[dict, InjectedState]) -> str:
+        """
+        Run the tool.
+
+        Args:
+            question (str): The question to ask about the simulation results.
+            state (dict): The state of the graph.
+            run_manager (Optional[CallbackManagerForToolRun]): The CallbackManagerForToolRun object.
+
+        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
+        df_agent = create_pandas_dataframe_agent(
+                        ChatOpenAI(model=state['llm_model']),
+                        allow_dangerous_code=True,
+                        agent_type=AgentType.OPENAI_FUNCTIONS,
+                        df=df,
+                        prefix=prompt_content)
+        llm_result = df_agent.invoke(question)
+        return llm_result["output"]

aiagents4pharma-1.7.0/aiagents4pharma/talk2biomodels/tools/custom_plotter.py

@@ -0,0 +1,84 @@
+#!/usr/bin/env python3
+
+"""
+Tool for plotting a custom figure.
+"""
+
+import logging
+from typing import Type, List, TypedDict, Annotated, Tuple, Union, Literal
+from pydantic import BaseModel, Field
+import pandas as pd
+from langchain_openai import ChatOpenAI
+from langchain_core.tools import BaseTool
+from langgraph.prebuilt import InjectedState
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class CustomPlotterInput(BaseModel):
+    """
+    Input schema for the PlotImage tool.
+    """
+    question: str = Field(description="Description of the plot")
+    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 CustomPlotterTool(BaseTool):
+    """
+    Tool for plotting a custom plot.
+    """
+    name: str = "custom_plotter"
+    description: str = "A tool to plot a custom figure."
+    args_schema: Type[BaseModel] = CustomPlotterInput
+    response_format: str = "content_and_artifact"
+
+    def _run(self,
+             question: str,
+             state: Annotated[dict, InjectedState]
+             ) -> Tuple[str, Union[None, List[str]]]:
+        """
+        Run the tool.
+
+        Args:
+            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'])
+        species_names = df.columns.tolist()
+        # Exclude the time column
+        species_names.remove('Time')
+        # 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
+        # simulation results.
+        class CustomHeader(TypedDict):
+            """
+            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.")
+        # Create an instance of the LLM model
+        llm = ChatOpenAI(model=state['llm_model'], temperature=0)
+        llm_with_structured_output = llm.with_structured_output(CustomHeader)
+        results = llm_with_structured_output.invoke(question)
+        extracted_species = []
+        # Extract the species from the results
+        # that are available in the simulation results
+        for species in results['relevant_species']:
+            if species in species_names:
+                extracted_species.append(species)
+        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

aiagents4pharma-1.7.0/aiagents4pharma/talk2biomodels/tools/get_modelinfo.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+
+"""
+Tool for get model information.
+"""
+
+import logging
+from typing import Type, Optional, Annotated
+from dataclasses import dataclass
+import basico
+from pydantic import BaseModel, Field
+from langchain_core.tools import BaseTool
+from langchain_core.messages import ToolMessage
+from langchain_core.tools.base import InjectedToolCallId
+from langgraph.prebuilt import InjectedState
+from langgraph.types import Command
+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.")
+    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.")
+
+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_parameters"
+    description: str = "A tool for extracting model information."
+    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:
+        """
+        Run the tool.
+
+        Args:
+            requested_model_info (RequestedModelInfo): The requested model information.
+            tool_call_id (str): The tool call ID. This is injected by the system.
+            state (dict): The state of the tool.
+            sys_bio_model (ModelData): The model data.
+
+        Returns:
+            Command: The updated state of the tool.
+        """
+        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)
+        dic_results = {}
+        # 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'] = ','.join(dic_results['Species'])
+
+        # Extract parameters from the model
+        if requested_model_info.parameters:
+            df_parameters = basico.model_info.get_parameters(model=model_obj.copasi_model)
+            dic_results['Parameters'] = df_parameters.index.tolist()
+            dic_results['Parameters'] = ','.join(dic_results['Parameters'])
+
+        # 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'])
+
+        # Extract description from the model
+        if requested_model_info.description:
+            dic_results['Description'] = model_obj.description
+
+        # Extract description from the model
+        if requested_model_info.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)
+
+        # 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():
+            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
+                            )
+                        ],
+                    }
+            )

aiagents4pharma-1.7.0/aiagents4pharma/talk2biomodels/tools/load_biomodel.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+
+"""
+Function for loading the BioModel.
+"""
+
+from dataclasses import dataclass
+from ..models.basico_model import BasicoModel
+
+@dataclass
+class ModelData:
+    """
+    Dataclass for storing the model data.
+    """
+    biomodel_id: int = 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.
+    """
+    model_object = None
+    if sys_bio_model.biomodel_id:
+        model_object = BasicoModel(biomodel_id=sys_bio_model.biomodel_id)
+    elif sbml_file_path:
+        model_object = BasicoModel(sbml_file_path=sbml_file_path)
+    return model_object
+    # return None

{aiagents4pharma-1.6.1 → aiagents4pharma-1.7.0}/aiagents4pharma/talk2biomodels/tools/search_models.py

@@ -4,21 +4,21 @@
 Tool for searching models based on search query.
 """
 
-from urllib.error import URLError
-from time import sleep
-from typing import Type
+from typing import Type, Annotated
 from pydantic import BaseModel, Field
 from basico import biomodels
 from langchain_core.tools import BaseTool
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_openai import ChatOpenAI
+from langgraph.prebuilt import InjectedState
 
 class SearchModelsInput(BaseModel):
     """
     Input schema for the search models tool.
     """
     query: str = Field(description="Search models query", default=None)
+    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.
@@ -31,7 +31,9 @@ class SearchModelsTool(BaseTool):
     args_schema: Type[BaseModel] = SearchModelsInput
     return_direct: bool = True
 
-    def _run(self, query: str) -> str:
+    def _run(self,
+             query: str,
+             state: Annotated[dict, InjectedState]) -> dict:
         """
         Run the tool.
 
@@ -39,20 +41,10 @@ class SearchModelsTool(BaseTool):
             query (str): The search query.
 
         Returns:
-            str: The answer to the question.
+            dict: The answer to the question in the form of a dictionary.
         """
-        attempts = 0
-        max_retries = 3
-        while attempts < max_retries:
-            try:
-                search_results = biomodels.search_for_model(query)
-                break
-            except URLError as e:
-                attempts += 1
-                sleep(10)
-                if attempts >= max_retries:
-                    raise e
-        llm = ChatOpenAI(model="gpt-4o-mini")
+        search_results = biomodels.search_for_model(query)
+        llm = ChatOpenAI(model=state['llm_model'])
         # Check if run_manager's metadata has the key 'prompt_content'
         prompt_content = f'''
                         Convert the input into a table.
@@ -80,15 +72,3 @@ class SearchModelsTool(BaseTool):
         parser = StrOutputParser()
         chain = prompt_template | llm | parser
         return chain.invoke({"input": search_results})
-
-    def get_metadata(self):
-        """
-        Get metadata for the tool.
-
-        Returns:
-            dict: The metadata for the tool.
-        """
-        return {
-            "name": self.name,
-            "description": self.description
-        }