aiagents4pharma 1.24.1__tar.gz → 1.25.1__tar.gz

{aiagents4pharma-1.24.1 → aiagents4pharma-1.25.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: aiagents4pharma
-Version: 1.24.1
+Version: 1.25.1
 Summary: AI Agents for drug discovery, drug development, and other pharmaceutical R&D.
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
@@ -42,7 +42,6 @@ Requires-Dist: sentence_transformers==3.3.1
 Requires-Dist: tabulate==0.9.0
 Requires-Dist: torch==2.2.2
 Requires-Dist: torch_geometric==2.6.1
-Requires-Dist: tqdm==4.66.6
 Requires-Dist: transformers==4.48.0
 Requires-Dist: mkdocs==1.6.1
 Requires-Dist: mkdocs-jupyter==0.25.1
@@ -56,7 +55,6 @@ Requires-Dist: h5py==3.13.0
 Requires-Dist: igraph==0.11.8
 Requires-Dist: ipykernel==6.29.5
 Requires-Dist: ipython==8.32.0
-Requires-Dist: numpy==2.1.3
 Requires-Dist: nbformat==5.10.4
 Requires-Dist: scipy==1.15.2
 Requires-Dist: tqdm==4.67.1
@@ -83,6 +81,7 @@ Our toolkit currently consists of the following agents:
 - **Talk2KnowledgeGraphs** _(v1 in progress)_: Access and explore complex biological knowledge graphs for insightful data connections.
 - **Talk2Scholars** _(v1 in progress)_: Get recommendations for articles related to your choice. Download, query, and write/retrieve them to your reference manager (currently supporting Zotero).
 - **Talk2Cells** _(v1 in progress)_: Query and analyze sequencing data with ease.
+- **Talk2AIAgents4Pharma** _(v1 in progress)_: Converse with all the agents above (currently supports T2B and T2KG)
 
 ![AIAgents4Pharma](docs/assets/AIAgents4Pharma.png)
 
@@ -104,49 +103,47 @@ Check out the tutorials on each agent for detailed instrcutions.
 
 _Both `Talk2Biomodels` and `Talk2Scholars` are now available on Docker Hub._
 
-#### **Running Talk2Biomodels**
+1. **Pull the Docker images**
 
-1. **Pull the Docker image**
    ```bash
    docker pull virtualpatientengine/talk2biomodels
    ```
-2. **Run the container**
+
+   ```bash
+   docker pull virtualpatientengine/talk2scholars
+   ```
+
+2. **Run the containers**
+
    ```bash
    docker run -d \
+     --name talk2biomodels \
      -e OPENAI_API_KEY=<your_openai_api_key> \
      -e NVIDIA_API_KEY=<your_nvidia_api_key> \
      -p 8501:8501 \
      virtualpatientengine/talk2biomodels
    ```
-3. **Access the Web App**  
-    Open your browser and go to:
-   ```
-   http://localhost:8501
-   ```
-   _You can create a free account at NVIDIA and apply for their
-   free credits [here](https://build.nvidia.com/explore/discover)._
 
-#### **Running Talk2Scholars**
-
-1. **Pull the Docker image**
-   ```bash
-   docker pull virtualpatientengine/talk2scholars
-   ```
-2. **Run the container**
    ```bash
    docker run -d \
+     --name talk2scholars \
      -e OPENAI_API_KEY=<your_openai_api_key> \
      -e ZOTERO_API_KEY=<your_zotero_api_key> \
      -e ZOTERO_USER_ID=<your_zotero_user_id> \
      -p 8501:8501 \
      virtualpatientengine/talk2scholars
    ```
-3. **Access the Web App**  
-   Open your browser and go to:
+
+3. **Access the Web App**
+    Open your browser and go to:
+
    ```
    http://localhost:8501
    ```
 
+   _You can create a free account at NVIDIA and apply for their
+   free credits [here](https://build.nvidia.com/explore/discover)._
+
 #### **Notes**
 
 - Ensure you **replace `<your_openai_api_key>`, `<your_nvidia_api_key>`, `<your_zotero_api_key>`, and `<your_zotero_user_id>`** with your actual credentials.
@@ -167,6 +164,33 @@ _Both `Talk2Biomodels` and `Talk2Scholars` are now available on Docker Hub._
    ```bash
    pip install -r requirements.txt
    ```
+   The current version of T2KG requires additional Ollama library to be installed.
+
+   Ollama can be easily downloaded and installed from the following link: [https://ollama.com/download](https://ollama.com/download)
+
+   As an alternative, use the following commands to install the library using terminal and to pull necessary model:
+
+   - Ubuntu:
+      ```
+      curl -fsSL https://ollama.com/install.sh | sh
+      ollama pull nomic-embed-text
+      ```
+   - Windows:
+      ```
+      curl -L https://ollama.com/download/ollama-windows-amd64.zip -o ollama-windows-amd64.zip
+      tar -xzf .\ollama-windows-amd64.zip
+      start ollama serve
+      ollama pull nomic-embed-text
+      ```
+   - macOS:
+      ```
+      brew install ollama
+      ollama pull nomic-embed-text
+      ```
+   A list of pulled Ollama models can be checked using the following command:
+   ```
+   ollama list
+   ```
 3. **Initialize OPENAI_API_KEY and NVIDIA_API_KEY**
 
    ```bash
@@ -192,10 +216,10 @@ _Both `Talk2Biomodels` and `Talk2Scholars` are now available on Docker Hub._
    ```
 
    _Please note that this will create a new tracing project in your Langsmith
-   account with the name `T2X-xxxx`, where `X` can be `B` (Biomodels), `S` (Scholars),
-   `KG` (KnowledgeGraphs), or `C` (Cells). If you skip the previous step, it will
-   default to the name `default`. `xxxx` will be the 4-digit ID created for the
-   session._
+   account with the name `T2X-xxxx`, where `X` can be `AA4P` (Main Agent),
+   `B` (Biomodels), `S` (Scholars), `KG` (KnowledgeGraphs), or `C` (Cells).
+   If you skip the previous step, it will default to the name `default`.
+   `xxxx` will be the 4-digit ID created for the session._
 
 6. **Launch the app:**
    ```bash

{aiagents4pharma-1.24.1 → aiagents4pharma-1.25.1}/README.md

@@ -16,6 +16,7 @@ Our toolkit currently consists of the following agents:
 - **Talk2KnowledgeGraphs** _(v1 in progress)_: Access and explore complex biological knowledge graphs for insightful data connections.
 - **Talk2Scholars** _(v1 in progress)_: Get recommendations for articles related to your choice. Download, query, and write/retrieve them to your reference manager (currently supporting Zotero).
 - **Talk2Cells** _(v1 in progress)_: Query and analyze sequencing data with ease.
+- **Talk2AIAgents4Pharma** _(v1 in progress)_: Converse with all the agents above (currently supports T2B and T2KG)
 
 ![AIAgents4Pharma](docs/assets/AIAgents4Pharma.png)
 
@@ -37,49 +38,47 @@ Check out the tutorials on each agent for detailed instrcutions.
 
 _Both `Talk2Biomodels` and `Talk2Scholars` are now available on Docker Hub._
 
-#### **Running Talk2Biomodels**
+1. **Pull the Docker images**
 
-1. **Pull the Docker image**
    ```bash
    docker pull virtualpatientengine/talk2biomodels
    ```
-2. **Run the container**
+
+   ```bash
+   docker pull virtualpatientengine/talk2scholars
+   ```
+
+2. **Run the containers**
+
    ```bash
    docker run -d \
+     --name talk2biomodels \
      -e OPENAI_API_KEY=<your_openai_api_key> \
      -e NVIDIA_API_KEY=<your_nvidia_api_key> \
      -p 8501:8501 \
      virtualpatientengine/talk2biomodels
    ```
-3. **Access the Web App**  
-    Open your browser and go to:
-   ```
-   http://localhost:8501
-   ```
-   _You can create a free account at NVIDIA and apply for their
-   free credits [here](https://build.nvidia.com/explore/discover)._
 
-#### **Running Talk2Scholars**
-
-1. **Pull the Docker image**
-   ```bash
-   docker pull virtualpatientengine/talk2scholars
-   ```
-2. **Run the container**
    ```bash
    docker run -d \
+     --name talk2scholars \
      -e OPENAI_API_KEY=<your_openai_api_key> \
      -e ZOTERO_API_KEY=<your_zotero_api_key> \
      -e ZOTERO_USER_ID=<your_zotero_user_id> \
      -p 8501:8501 \
      virtualpatientengine/talk2scholars
    ```
-3. **Access the Web App**  
-   Open your browser and go to:
+
+3. **Access the Web App**
+    Open your browser and go to:
+
    ```
    http://localhost:8501
    ```
 
+   _You can create a free account at NVIDIA and apply for their
+   free credits [here](https://build.nvidia.com/explore/discover)._
+
 #### **Notes**
 
 - Ensure you **replace `<your_openai_api_key>`, `<your_nvidia_api_key>`, `<your_zotero_api_key>`, and `<your_zotero_user_id>`** with your actual credentials.
@@ -100,6 +99,33 @@ _Both `Talk2Biomodels` and `Talk2Scholars` are now available on Docker Hub._
    ```bash
    pip install -r requirements.txt
    ```
+   The current version of T2KG requires additional Ollama library to be installed.
+
+   Ollama can be easily downloaded and installed from the following link: [https://ollama.com/download](https://ollama.com/download)
+
+   As an alternative, use the following commands to install the library using terminal and to pull necessary model:
+
+   - Ubuntu:
+      ```
+      curl -fsSL https://ollama.com/install.sh | sh
+      ollama pull nomic-embed-text
+      ```
+   - Windows:
+      ```
+      curl -L https://ollama.com/download/ollama-windows-amd64.zip -o ollama-windows-amd64.zip
+      tar -xzf .\ollama-windows-amd64.zip
+      start ollama serve
+      ollama pull nomic-embed-text
+      ```
+   - macOS:
+      ```
+      brew install ollama
+      ollama pull nomic-embed-text
+      ```
+   A list of pulled Ollama models can be checked using the following command:
+   ```
+   ollama list
+   ```
 3. **Initialize OPENAI_API_KEY and NVIDIA_API_KEY**
 
    ```bash
@@ -125,10 +151,10 @@ _Both `Talk2Biomodels` and `Talk2Scholars` are now available on Docker Hub._
    ```
 
    _Please note that this will create a new tracing project in your Langsmith
-   account with the name `T2X-xxxx`, where `X` can be `B` (Biomodels), `S` (Scholars),
-   `KG` (KnowledgeGraphs), or `C` (Cells). If you skip the previous step, it will
-   default to the name `default`. `xxxx` will be the 4-digit ID created for the
-   session._
+   account with the name `T2X-xxxx`, where `X` can be `AA4P` (Main Agent),
+   `B` (Biomodels), `S` (Scholars), `KG` (KnowledgeGraphs), or `C` (Cells).
+   If you skip the previous step, it will default to the name `default`.
+   `xxxx` will be the 4-digit ID created for the session._
 
 6. **Launch the app:**
    ```bash

{aiagents4pharma-1.24.1 → aiagents4pharma-1.25.1}/aiagents4pharma/__init__.py

@@ -7,4 +7,5 @@ from . import (
     talk2cells,
     talk2scholars,
     talk2knowledgegraphs,
+    talk2aiagents4pharma
 )

aiagents4pharma-1.25.1/aiagents4pharma/talk2aiagents4pharma/__init__.py

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

aiagents4pharma-1.25.1/aiagents4pharma/talk2aiagents4pharma/agents/__init__.py

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

aiagents4pharma-1.25.1/aiagents4pharma/talk2aiagents4pharma/agents/main_agent.py

@@ -0,0 +1,68 @@
+#/usr/bin/env python3
+
+'''
+This is the main agent file for the AIAgents4Pharma.
+'''
+
+import logging
+import hydra
+from langgraph_supervisor import create_supervisor
+from langchain_openai import ChatOpenAI
+from langchain_core.language_models.chat_models import BaseChatModel
+from langgraph.checkpoint.memory import MemorySaver
+from ...talk2biomodels.agents.t2b_agent import get_app as get_app_t2b
+from ...talk2knowledgegraphs.agents.t2kg_agent import get_app as get_app_t2kg
+from ..states.state_talk2aiagents4pharma import Talk2AIAgents4Pharma
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def get_app(uniq_id, llm_model: BaseChatModel):
+    '''
+    This function returns the langraph app.
+    '''
+    if llm_model.model_name == 'gpt-4o-mini':
+        llm_model = ChatOpenAI(model='gpt-4o-mini',
+                               temperature=0,
+                               model_kwargs={"parallel_tool_calls": False})
+    # Load hydra configuration
+    logger.log(logging.INFO, "Launching AIAgents4Pharma_Agent with thread_id %s", uniq_id)
+    with hydra.initialize(version_base=None, config_path="../configs"):
+        cfg = hydra.compose(config_name='config',
+                            overrides=['agents/main_agent=default'])
+        cfg = cfg.agents.main_agent
+    logger.log(logging.INFO, "System_prompt of T2AA4P: %s", cfg.system_prompt)
+    with hydra.initialize(version_base=None, config_path="../../talk2biomodels/configs"):
+        cfg_t2b = hydra.compose(config_name='config',
+                            overrides=['agents/t2b_agent=default'])
+        cfg_t2b = cfg_t2b.agents.t2b_agent
+    with hydra.initialize(version_base=None, config_path="../../talk2knowledgegraphs/configs"):
+        cfg_t2kg = hydra.compose(config_name='config',
+                            overrides=['agents/t2kg_agent=default'])
+        cfg_t2kg = cfg_t2kg.agents.t2kg_agent
+    system_prompt = cfg.system_prompt
+    system_prompt += "\n\nHere is the system prompt of T2B agent\n"
+    system_prompt += cfg_t2b.state_modifier
+    system_prompt += "\n\nHere is the system prompt of T2KG agent\n"
+    system_prompt += cfg_t2kg.state_modifier
+    # Create supervisor workflow
+    workflow = create_supervisor(
+        [
+            get_app_t2b(uniq_id, llm_model),    # Talk2BioModels
+            get_app_t2kg(uniq_id, llm_model)    # Talk2KnowledgeGraphs
+        ],
+        model=llm_model,
+        state_schema=Talk2AIAgents4Pharma,
+        # Full history is needed to extract
+        # the tool artifacts
+        output_mode="full_history",
+        add_handoff_back_messages=False,
+        prompt=system_prompt
+    )
+
+    # Compile and run
+    app = workflow.compile(checkpointer=MemorySaver(),
+                           name="AIAgents4Pharma_Agent")
+
+    return app

aiagents4pharma-1.25.1/aiagents4pharma/talk2aiagents4pharma/configs/__init__.py

@@ -0,0 +1,5 @@
+'''
+Import all the modules in the package
+'''
+
+from . import agents

aiagents4pharma-1.25.1/aiagents4pharma/talk2aiagents4pharma/configs/agents/__init__.py

@@ -0,0 +1,5 @@
+'''
+Import all the modules in the package
+'''
+
+from . import main_agent

aiagents4pharma-1.25.1/aiagents4pharma/talk2aiagents4pharma/configs/agents/main_agent/default.yaml

@@ -0,0 +1,12 @@
+_target_: agents.main_agent.get_app
+system_prompt: >
+  You are Talk2AIAgents4Pharma agent.
+  You are managing a team of the following 2 agents:
+   
+  1. Talk2Biomodels (T2B) agent: This agent can operate
+  on mathematical models of biological systems. This
+  agent can also query an uploaded document/pdf/article. 
+   
+  2. Talk2KnowledgeGraphs (T2KG) agent: This agent can
+  reason over a knowledge graph of biological entities
+  and their relationships.

aiagents4pharma-1.25.1/aiagents4pharma/talk2aiagents4pharma/configs/config.yaml

@@ -0,0 +1,3 @@
+defaults:
+  - _self_
+  - agents/main_agent: default

aiagents4pharma-1.25.1/aiagents4pharma/talk2aiagents4pharma/states/__init__.py

@@ -0,0 +1,4 @@
+'''
+This file is used to import all the models in the package.
+'''
+from . import state_talk2aiagents4pharma

aiagents4pharma-1.25.1/aiagents4pharma/talk2aiagents4pharma/states/state_talk2aiagents4pharma.py

@@ -0,0 +1,16 @@
+"""
+This is the state file for the Talk2AIAgents4Pharma agent.
+"""
+
+from ...talk2biomodels.states.state_talk2biomodels import Talk2Biomodels
+from ...talk2knowledgegraphs.states.state_talk2knowledgegraphs import Talk2KnowledgeGraphs
+
+class Talk2AIAgents4Pharma(Talk2Biomodels,
+                           Talk2KnowledgeGraphs):
+    """
+    The state for the Talk2AIAgents4Pharma agent.
+
+    This class inherits from the classes:
+    1. Talk2Biomodels
+    2. Talk2KnowledgeGraphs
+    """

aiagents4pharma-1.25.1/aiagents4pharma/talk2aiagents4pharma/tests/test_main_agent.py

@@ -0,0 +1,112 @@
+'''
+Test Talk2AIAgents4Pharma supervisor agent.
+'''
+
+import pytest
+from langchain_core.messages import HumanMessage
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from ..agents.main_agent import get_app
+
+# Define the data path for the test files of Talk2KnowledgeGraphs agent
+DATA_PATH = "aiagents4pharma/talk2knowledgegraphs/tests/files"
+LLM_MODEL = ChatOpenAI(model="gpt-4o-mini", temperature=0.0)
+
+@pytest.fixture(name="input_dict")
+def input_dict_fixture():
+    """
+    Input dictionary fixture for Talk2AIAgents4Pharma agent,
+    which is partly inherited from the Talk2KnowledgeGraphs agent.
+    """
+    input_dict = {
+        "topk_nodes": 3,
+        "topk_edges": 3,
+        "uploaded_files": [],
+        "dic_source_graph": [
+            {
+                "name": "PrimeKG",
+                "kg_pyg_path": f"{DATA_PATH}/primekg_ibd_pyg_graph.pkl",
+                "kg_text_path": f"{DATA_PATH}/primekg_ibd_text_graph.pkl",
+            }
+        ],
+        "dic_extracted_graph": []
+    }
+
+    return input_dict
+
+def test_main_agent_invokes_t2kg(input_dict):
+    """
+    In the following test, we will ask the main agent (supervisor)
+    to list drugs that target the gene Interleukin-6. We will check
+    if the Talk2KnowledgeGraphs agent is invoked. We will do so by
+    checking the state of the Talk2AIAgents4Pharma agent, which is
+    partly inherited from the Talk2KnowledgeGraphs agent
+
+    Args:
+        input_dict: Input dictionary
+    """
+    # Prepare LLM and embedding model
+    input_dict["llm_model"] = LLM_MODEL
+    input_dict["embedding_model"] = OpenAIEmbeddings(model="text-embedding-3-small")
+
+    # Setup the app
+    unique_id = 12345
+    app = get_app(unique_id, llm_model=input_dict["llm_model"])
+    config = {"configurable": {"thread_id": unique_id}}
+    # Update state
+    app.update_state(
+        config,
+        input_dict,
+    )
+    prompt = "List drugs that target the gene Interleukin-6"
+
+    # Invoke the agent
+    response = app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
+
+    # Check assistant message
+    assistant_msg = response["messages"][-1].content
+    assert isinstance(assistant_msg, str)
+
+    # Check extracted subgraph dictionary
+    current_state = app.get_state(config)
+    dic_extracted_graph = current_state.values["dic_extracted_graph"][0]
+    assert isinstance(dic_extracted_graph, dict)
+    assert dic_extracted_graph["graph_source"] == "PrimeKG"
+    assert dic_extracted_graph["topk_nodes"] == 3
+    assert dic_extracted_graph["topk_edges"] == 3
+    assert isinstance(dic_extracted_graph["graph_dict"], dict)
+    assert len(dic_extracted_graph["graph_dict"]["nodes"]) > 0
+    assert len(dic_extracted_graph["graph_dict"]["edges"]) > 0
+    assert isinstance(dic_extracted_graph["graph_text"], str)
+    # Check summarized subgraph
+    assert isinstance(dic_extracted_graph["graph_summary"], str)
+
+def test_main_agent_invokes_t2b():
+    '''
+    In the following test, we will ask the main agent (supervisor)
+    to simulate a model. And we will check if the Talk2BioModels
+    agent is invoked. We will do so by checking the state of the
+    Talk2AIAgents4Pharma agent, which is partly inherited from the
+    Talk2BioModels agent.
+    '''
+    unique_id = 123
+    app = get_app(unique_id, llm_model=LLM_MODEL)
+    config = {"configurable": {"thread_id": unique_id}}
+    prompt = "Simulate model 64"
+    # Invoke the agent
+    app.invoke(
+        {"messages": [HumanMessage(content=prompt)]},
+        config=config
+    )
+    # Get the state of the Talk2AIAgents4Pharma agent
+    current_state = app.get_state(config)
+    # Check if the dic_simulated_data is in the state
+    dic_simulated_data = current_state.values["dic_simulated_data"]
+    # Check if the dic_simulated_data is a list
+    assert isinstance(dic_simulated_data, list)
+    # Check if the length of the dic_simulated_data is 1
+    assert len(dic_simulated_data) == 1
+    # Check if the source of the model is 64
+    assert dic_simulated_data[0]['source'] == 64
+    # Check if the data of the model contains
+    # '1,3-bisphosphoglycerate'
+    assert '1,3-bisphosphoglycerate' in dic_simulated_data[0]['data']

{aiagents4pharma-1.24.1 → aiagents4pharma-1.25.1}/aiagents4pharma/talk2biomodels/agents/t2b_agent.py

@@ -7,7 +7,6 @@ This is the agent file for the Talk2BioModels agent.
 import logging
 from typing import Annotated
 import hydra
-from langchain_openai import ChatOpenAI
 from langchain_core.language_models.chat_models import BaseChatModel
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import START, StateGraph
@@ -28,7 +27,7 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 def get_app(uniq_id,
-            llm_model: BaseChatModel = ChatOpenAI(model='gpt-4o-mini', temperature=0)):
+            llm_model: BaseChatModel):
     '''
     This function returns the langraph app.
     '''
@@ -65,7 +64,8 @@ def get_app(uniq_id,
                 llm_model,
                 tools=tools,
                 state_schema=Talk2Biomodels,
-                state_modifier=cfg.state_modifier,
+                prompt=cfg.state_modifier,
+                version='v2',
                 checkpointer=MemorySaver()
             )
 

aiagents4pharma-1.25.1/aiagents4pharma/talk2biomodels/tests/__init__.py

@@ -0,0 +1,3 @@
+'''
+This module contains the test cases.
+'''

{aiagents4pharma-1.24.1 → aiagents4pharma-1.25.1}/aiagents4pharma/talk2biomodels/tests/test_ask_question.py

@@ -11,7 +11,7 @@ def test_ask_question_tool():
     Test the ask_question tool without the simulation results.
     '''
     unique_id = 12345
-    app = get_app(unique_id)
+    app = get_app(unique_id, llm_model=ChatOpenAI(model='gpt-4o-mini', temperature=0))
     config = {"configurable": {"thread_id": unique_id}}
 
     ##########################################
@@ -20,9 +20,6 @@ def test_ask_question_tool():
     # simulation has not been run. In this
     # case, the tool should return an error
     ##########################################
-    # Update state
-    app.update_state(config,
-            {"llm_model": ChatOpenAI(model='gpt-4o-mini', temperature=0)})
     # Define the prompt
     prompt = "Call the ask_question tool to answer the "
     prompt += "question: What is the concentration of CRP "

{aiagents4pharma-1.24.1 → aiagents4pharma-1.25.1}/aiagents4pharma/talk2biomodels/tests/test_get_annotation.py

@@ -9,17 +9,20 @@ from langchain_openai import ChatOpenAI
 from ..agents.t2b_agent import get_app
 from ..tools.get_annotation import prepare_content_msg
 
+LLM_MODEL = ChatOpenAI(model='gpt-4o-mini', temperature=0)
+
 @pytest.fixture(name="make_graph")
 def make_graph_fixture():
     '''
     Create an instance of the talk2biomodels agent.
     '''
     unique_id = random.randint(1000, 9999)
-    graph = get_app(unique_id)
+    graph = get_app(unique_id, llm_model=LLM_MODEL)
     config = {"configurable": {"thread_id": unique_id}}
-    graph.update_state(config, {"llm_model": ChatOpenAI(model='gpt-4o-mini',
-                                    temperature=0)
-                                })
+    graph.update_state(
+            config,
+            {"llm_model": LLM_MODEL}
+        )
     return graph, config
 
 def test_no_model_provided(make_graph):

{aiagents4pharma-1.24.1 → aiagents4pharma-1.25.1}/aiagents4pharma/talk2biomodels/tests/test_getmodelinfo.py

@@ -3,14 +3,17 @@ Test cases for Talk2Biomodels get_modelinfo tool.
 '''
 
 from langchain_core.messages import HumanMessage, ToolMessage
+from langchain_openai import ChatOpenAI
 from ..agents.t2b_agent import get_app
 
+LLM_MODEL = ChatOpenAI(model='gpt-4o-mini',temperature=0)
+
 def test_get_modelinfo_tool():
     '''
     Test the get_modelinfo tool.
     '''
     unique_id = 12345
-    app = get_app(unique_id)
+    app = get_app(unique_id, LLM_MODEL)
     config = {"configurable": {"thread_id": unique_id}}
     # Update state
     app.update_state(config,
@@ -33,7 +36,7 @@ def test_model_with_no_species():
     This should raise a tool error.
     '''
     unique_id = 12345
-    app = get_app(unique_id)
+    app = get_app(unique_id, LLM_MODEL)
     config = {"configurable": {"thread_id": unique_id}}
     prompt = "Extract all species from model 20"
     # Test the tool get_modelinfo
@@ -64,7 +67,7 @@ def test_model_with_no_parameters():
     This should raise a tool error.
     '''
     unique_id = 12345
-    app = get_app(unique_id)
+    app = get_app(unique_id, LLM_MODEL)
     config = {"configurable": {"thread_id": unique_id}}
     prompt = "Extract all parameters from model 10"
     # Test the tool get_modelinfo

{aiagents4pharma-1.24.1 → aiagents4pharma-1.25.1}/aiagents4pharma/talk2biomodels/tests/test_integration.py

@@ -14,9 +14,8 @@ def test_integration():
     Test the integration of the tools.
     '''
     unique_id = 1234567
-    app = get_app(unique_id)
+    app = get_app(unique_id, llm_model=LLM_MODEL)
     config = {"configurable": {"thread_id": unique_id}}
-    app.update_state(config, {"llm_model": LLM_MODEL})
     # ##########################################
     # ## Test simulate_model tool
     # ##########################################