aiagents4pharma 1.8.2__tar.gz → 1.8.3__tar.gz

{aiagents4pharma-1.8.2 → aiagents4pharma-1.8.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: aiagents4pharma
-Version: 1.8.2
+Version: 1.8.3
 Summary: AI Agents for drug discovery, drug development, and other pharmaceutical R&D
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License

{aiagents4pharma-1.8.2 → aiagents4pharma-1.8.3}/aiagents4pharma/__init__.py

@@ -5,3 +5,4 @@ This file is used to import aiagents4pharma modules.
 from . import talk2biomodels
 from . import talk2cells
 from . import talk2knowledgegraphs
+from . import configs

aiagents4pharma-1.8.3/aiagents4pharma/configs/__init__.py

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

aiagents4pharma-1.8.3/aiagents4pharma/configs/config.yaml

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

aiagents4pharma-1.8.3/aiagents4pharma/configs/talk2biomodels/__init__.py

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

aiagents4pharma-1.8.3/aiagents4pharma/configs/talk2biomodels/agents/__init__.py

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

aiagents4pharma-1.8.3/aiagents4pharma/configs/talk2biomodels/agents/t2b_agent/__init__.py

@@ -0,0 +1,3 @@
+'''
+Import all the modules in the package
+'''

aiagents4pharma-1.8.3/aiagents4pharma/configs/talk2biomodels/agents/t2b_agent/default.yaml

@@ -0,0 +1,6 @@
+_target_: talk2biomodels.agents.t2b_agent.get_app
+state_modifier: >
+  You are Talk2BioModels agent.
+  If the user asks for the uploaded model, 
+  then pass the use_uploaded_model argument
+    as True.

{aiagents4pharma-1.8.2 → aiagents4pharma-1.8.3}/aiagents4pharma/talk2biomodels/agents/t2b_agent.py

@@ -52,10 +52,10 @@ def get_app(uniq_id, llm_model='gpt-4o-mini'):
     llm = ChatOpenAI(model=llm_model, temperature=0)
     # Load hydra configuration
     logger.log(logging.INFO, "Load Hydra configuration for Talk2BioModels agent.")
-    with hydra.initialize(version_base=None, config_path="../../../configs"):
+    with hydra.initialize(version_base=None, config_path="../../configs"):
         cfg = hydra.compose(config_name='config',
-                            overrides=['aiagents4pharma/talk2biomodels/agents/t2b_agent=default'])
-        cfg = cfg.aiagents4pharma.talk2biomodels.agents.t2b_agent
+                            overrides=['talk2biomodels/agents/t2b_agent=default'])
+        cfg = cfg.talk2biomodels.agents.t2b_agent
     logger.log(logging.INFO, "state_modifier: %s", cfg.state_modifier)
     # Create the agent
     model = create_react_agent(

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

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

aiagents4pharma-1.8.3/aiagents4pharma/talk2biomodels/tests/test_basico_model.py

@@ -0,0 +1,55 @@
+'''
+A test BasicoModel class for pytest unit testing.
+'''
+
+import pandas as pd
+import pytest
+import basico
+from ..models.basico_model import BasicoModel
+
+@pytest.fixture(name="model")
+def model_fixture():
+    """
+    A fixture for the BasicoModel class.
+    """
+    return BasicoModel(biomodel_id=64, species={"Pyruvate": 100}, duration=2, interval=2)
+
+def test_with_biomodel_id(model):
+    """
+    Test initialization of BasicoModel with biomodel_id.
+    """
+    assert model.biomodel_id == 64
+    # check if the simulation results are a pandas DataFrame object
+    assert isinstance(model.simulate(parameters={'Pyruvate': 0.5, 'KmPFKF6P': 1.5},
+                                     duration=2,
+                                     interval=2),
+                    pd.DataFrame)
+    assert isinstance(model.simulate(parameters={None: None}, duration=2, interval=2),
+                    pd.DataFrame)
+    assert model.description == basico.biomodels.get_model_info(model.biomodel_id)["description"]
+
+def test_with_sbml_file():
+    """
+    Test initialization of BasicoModel with sbml_file_path.
+    """
+    model_object = BasicoModel(sbml_file_path="./BIOMD0000000064_url.xml")
+    assert model_object.sbml_file_path == "./BIOMD0000000064_url.xml"
+    assert isinstance(model_object.simulate(duration=2, interval=2), pd.DataFrame)
+    assert isinstance(model_object.simulate(parameters={'NADH': 0.5}, duration=2, interval=2),
+                      pd.DataFrame)
+
+def test_check_biomodel_id_or_sbml_file_path():
+    '''
+    Test the check_biomodel_id_or_sbml_file_path method of the BioModel class.
+    '''
+    with pytest.raises(ValueError):
+        BasicoModel(species={"Pyruvate": 100}, duration=2, interval=2)
+
+def test_get_model_metadata():
+    """
+    Test the get_model_metadata method of the BasicoModel class.
+    """
+    model = BasicoModel(biomodel_id=64)
+    metadata = model.get_model_metadata()
+    assert metadata["Model Type"] == "SBML Model (COPASI)"
+    assert metadata["Parameter Count"] == len(basico.get_parameters())

aiagents4pharma-1.8.3/aiagents4pharma/talk2biomodels/tests/test_langgraph.py

@@ -0,0 +1,189 @@
+'''
+Test cases
+'''
+
+from langchain_core.messages import HumanMessage, ToolMessage
+from ..agents.t2b_agent import get_app
+
+def test_get_modelinfo_tool():
+    '''
+    Test the get_modelinfo tool.
+    '''
+    unique_id = 12345
+    app = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    # Update state
+    app.update_state(config,{"sbml_file_path": ["BIOMD0000000537.xml"]})
+    prompt = "Extract all relevant information from the uploaded model."
+    # Test the tool get_modelinfo
+    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)
+
+def test_search_models_tool():
+    '''
+    Test the search_models tool.
+    '''
+    unique_id = 12345
+    app = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    # Update state
+    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    prompt = "Search for models on Crohn's disease."
+    # Test the tool get_modelinfo
+    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)
+    # Check if the assistant message contains the
+    # biomodel id BIO0000000537
+    assert "BIOMD0000000537" in assistant_msg
+
+def test_ask_question_tool():
+    '''
+    Test the ask_question tool without the simulation results.
+    '''
+    unique_id = 12345
+    app = get_app(unique_id, llm_model='gpt-4o-mini')
+    config = {"configurable": {"thread_id": unique_id}}
+
+    ##########################################
+    # Test ask_question tool when simulation
+    # results are not available
+    ##########################################
+    # Update state
+    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    prompt = "Call the ask_question tool to answer the "
+    prompt += "question: What is the concentration of CRP "
+    prompt += "in serum at 1000 hours?"
+
+    # Test the tool get_modelinfo
+    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)
+
+def test_simulate_model_tool():
+    '''
+    Test the simulate_model tool.
+    '''
+    unique_id = 123
+    app = get_app(unique_id)
+    config = {"configurable": {"thread_id": unique_id}}
+    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    # ##########################################
+    # ## Test simulate_model tool
+    # ##########################################
+    prompt = "Simulate the model 537 for 2016 hours and intervals"
+    prompt += " 2016 with an initial concentration of `DoseQ2W` "
+    prompt += "set to 300 and `Dose` set to 0. Reset the concentration"
+    prompt += " of `NAD` to 100 every 500 hours."
+    # Test the tool get_modelinfo
+    response = app.invoke(
+                        {"messages": [HumanMessage(content=prompt)]},
+                        config=config
+                    )
+    assistant_msg = response["messages"][-1].content
+    print (assistant_msg)
+    # Check if the assistant message is a string
+    assert isinstance(assistant_msg, str)
+    ##########################################
+    # Test ask_question tool when simulation
+    # results are available
+    ##########################################
+    # Update state
+    app.update_state(config, {"llm_model": "gpt-4o-mini"})
+    prompt = "What is the concentration of CRP in serum at 1000 hours? "
+    # prompt += "Show only the concentration, rounded to 1 decimal place."
+    # prompt += "For example, if the concentration is 0.123456, "
+    # prompt += "your response should be `0.1`. Do not return any other information."
+    # Test the tool get_modelinfo
+    response = app.invoke(
+                        {"messages": [HumanMessage(content=prompt)]},
+                        config=config
+                    )
+    assistant_msg = response["messages"][-1].content
+    # print (assistant_msg)
+    # Check if the assistant message is a string
+    assert "1.7" in assistant_msg
+
+    ##########################################
+    # Test custom_plotter tool when the
+    # simulation results are available
+    ##########################################
+    prompt = "Plot only CRP related species."
+
+    # Update state
+    app.update_state(config, {"llm_model": "gpt-4o-mini"}
+                    )
+    # Test the tool get_modelinfo
+    response = app.invoke(
+                        {"messages": [HumanMessage(content=prompt)]},
+                        config=config
+                    )
+    assistant_msg = response["messages"][-1].content
+    current_state = app.get_state(config)
+    # Get the messages from the current state
+    # and reverse the order
+    reversed_messages = current_state.values["messages"][::-1]
+    # Loop through the reversed messages
+    # until a ToolMessage is found.
+    expected_artifact = ['CRP[serum]', 'CRPExtracellular']
+    expected_artifact += ['CRP Suppression (%)', 'CRP (% of baseline)']
+    expected_artifact += ['CRP[liver]']
+    predicted_artifact = []
+    for msg in reversed_messages:
+        if isinstance(msg, ToolMessage):
+            # Work on the message if it is a ToolMessage
+            # These may contain additional visuals that
+            # need to be displayed to the user.
+            if msg.name == "custom_plotter":
+                predicted_artifact = msg.artifact
+                break
+    # Check if the two artifacts are equal
+    # assert expected_artifact in predicted_artifact
+    assert set(expected_artifact).issubset(set(predicted_artifact))
+    ##########################################
+    # Test custom_plotter tool when the
+    # simulation results are available but
+    # the species is not available
+    ##########################################
+    prompt = "Plot the species `TP53`."
+
+    # Update state
+    app.update_state(config, {"llm_model": "gpt-4o-mini"}
+                    )
+    # Test the tool get_modelinfo
+    response = app.invoke(
+                        {"messages": [HumanMessage(content=prompt)]},
+                        config=config
+                    )
+    assistant_msg = response["messages"][-1].content
+    # print (response["messages"])
+    current_state = app.get_state(config)
+    # Get the messages from the current state
+    # and reverse the order
+    reversed_messages = current_state.values["messages"][::-1]
+    # Loop through the reversed messages until a
+    # ToolMessage is found.
+    predicted_artifact = []
+    for msg in reversed_messages:
+        if isinstance(msg, ToolMessage):
+            # Work on the message if it is a ToolMessage
+            # These may contain additional visuals that
+            # need to be displayed to the user.
+            if msg.name == "custom_plotter":
+                predicted_artifact = msg.artifact
+                break
+    # Check if the the predicted artifact is `None`
+    assert predicted_artifact is None

aiagents4pharma-1.8.3/aiagents4pharma/talk2biomodels/tests/test_sys_bio_model.py

@@ -0,0 +1,57 @@
+'''
+This file contains the unit tests for the BioModel class.
+'''
+
+from typing import List, Dict, Union, Optional
+from pydantic import Field
+import pytest
+from ..models.sys_bio_model import SysBioModel
+
+class TestBioModel(SysBioModel):
+    '''
+    A test BioModel class for unit testing.
+    '''
+
+    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")
+
+    def get_model_metadata(self) -> Dict[str, Union[str, int]]:
+        '''
+        Get the metadata of the model.
+        '''
+        return self.biomodel_id
+
+    def simulate(self,
+                 parameters: Dict[str, Union[float, int]],
+                 duration: Union[int, float]) -> List[float]:
+        '''
+        Simulate the model.
+        '''
+        param1 = parameters.get('param1', 0.0)
+        param2 = parameters.get('param2', 0.0)
+        return [param1 + param2 * t for t in range(int(duration))]
+
+def test_get_model_metadata():
+    '''
+    Test the get_model_metadata method of the BioModel class.
+    '''
+    model = TestBioModel(biomodel_id=123, name="Test Model", description="A test model")
+    metadata = model.get_model_metadata()
+    assert metadata == 123
+
+def test_check_biomodel_id_or_sbml_file_path():
+    '''
+    Test the check_biomodel_id_or_sbml_file_path method of the BioModel class.
+    '''
+    with pytest.raises(ValueError):
+        TestBioModel(name="Test Model", description="A test model")
+
+def test_simulate():
+    '''
+    Test the simulate method of the BioModel class.
+    '''
+    model = TestBioModel(biomodel_id=123, name="Test Model", description="A test model")
+    results = model.simulate(parameters={'param1': 1.0, 'param2': 2.0}, duration=4.0)
+    assert results == [1.0, 3.0, 5.0, 7.0]

aiagents4pharma-1.8.3/aiagents4pharma/talk2cells/tests/scp_agent/test_scp_agent.py

@@ -0,0 +1,23 @@
+'''
+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
+
+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
+                    )
+    assistant_msg = response["messages"][-1].content
+    # Check if the assistant message is a string
+    assert isinstance(assistant_msg, str)

aiagents4pharma-1.8.3/aiagents4pharma/talk2knowledgegraphs/tests/test_datasets_biobridge_primekg.py

@@ -0,0 +1,242 @@
+"""
+Test cases for datasets/primekg_loader.py
+"""
+
+import os
+import shutil
+import pytest
+from ..datasets.biobridge_primekg import BioBridgePrimeKG
+
+# Remove the data folder for testing if it exists
+PRIMEKG_LOCAL_DIR = "../data/primekg_test/"
+LOCAL_DIR = "../data/biobridge_primekg_test/"
+shutil.rmtree(LOCAL_DIR, ignore_errors=True)
+
+@pytest.fixture(name="biobridge_primekg")
+def biobridge_primekg_fixture():
+    """
+    Fixture for creating an instance of PrimeKG.
+    """
+    return BioBridgePrimeKG(primekg_dir=PRIMEKG_LOCAL_DIR,
+                            local_dir=LOCAL_DIR)
+
+def test_download_primekg(biobridge_primekg):
+    """
+    Test the loading method of the BioBridge-PrimeKG class by downloading data from repository.
+    """
+    # Load BioBridge-PrimeKG data
+    biobridge_primekg.load_data()
+    primekg_nodes = biobridge_primekg.get_primekg().get_nodes()
+    primekg_edges = biobridge_primekg.get_primekg().get_edges()
+    biobridge_data_config = biobridge_primekg.get_data_config()
+    biobridge_emb_dict = biobridge_primekg.get_node_embeddings()
+    biobridge_triplets = biobridge_primekg.get_primekg_triplets()
+    biobridge_splits = biobridge_primekg.get_train_test_split()
+    biobridge_node_info = biobridge_primekg.get_node_info_dict()
+
+    # Check if the local directories exists
+    assert os.path.exists(biobridge_primekg.primekg_dir)
+    assert os.path.exists(biobridge_primekg.local_dir)
+    # Check if downloaded and processed files exist
+    # PrimeKG files
+    files = ["nodes.tab", "primekg_nodes.tsv.gz",
+             "edges.csv", "primekg_edges.tsv.gz"]
+    for file in files:
+        path = f"{biobridge_primekg.primekg_dir}/{file}"
+        assert os.path.exists(path)
+    # BioBridge data config
+    assert os.path.exists(f"{biobridge_primekg.local_dir}/data_config.json")
+    # BioBridge embeddings
+    files = [
+        "protein.pkl",
+        "mf.pkl",
+        "cc.pkl",
+        "bp.pkl",
+        "drug.pkl",
+        "disease.pkl",
+        "embedding_dict.pkl"
+    ]
+    for file in files:
+        path = f"{biobridge_primekg.local_dir}/embeddings/{file}"
+        assert os.path.exists(path)
+    # BioBridge processed files
+    files = [
+        "protein.csv",
+        "mf.csv",
+        "cc.csv",
+        "bp.csv",
+        "drug.csv",
+        "disease.csv",
+        "triplet_full.tsv.gz",
+        "triplet_full_altered.tsv.gz",
+        "node_train.tsv.gz",
+        "triplet_train.tsv.gz",
+        "node_test.tsv.gz",
+        "triplet_test.tsv.gz",
+    ]
+    for file in files:
+        path = f"{biobridge_primekg.local_dir}/processed/{file}"
+        assert os.path.exists(path)
+    # Check processed PrimeKG dataframes
+    # Nodes
+    assert primekg_nodes is not None
+    assert len(primekg_nodes) > 0
+    assert primekg_nodes.shape[0] == 129375
+    # Edges
+    assert primekg_edges is not None
+    assert len(primekg_edges) > 0
+    assert primekg_edges.shape[0] == 8100498
+    # Check processed BioBridge data config
+    assert biobridge_data_config is not None
+    assert len(biobridge_data_config) > 0
+    assert len(biobridge_data_config['node_type']) == 10
+    assert len(biobridge_data_config['relation_type']) == 18
+    assert len(biobridge_data_config['emb_dim']) == 6
+    # Check processed BioBridge embeddings
+    assert biobridge_emb_dict is not None
+    assert len(biobridge_emb_dict) > 0
+    assert len(biobridge_emb_dict) == 85466
+    # Check processed BioBridge triplets
+    assert biobridge_triplets is not None
+    assert len(biobridge_triplets) > 0
+    assert biobridge_triplets.shape[0] == 3904610
+    assert list(biobridge_splits.keys()) == ['train', 'node_train', 'test', 'node_test']
+    assert len(biobridge_splits['train']) == 3510930
+    assert len(biobridge_splits['node_train']) == 76486
+    assert len(biobridge_splits['test']) == 393680
+    assert len(biobridge_splits['node_test']) == 8495
+    # Check node info dictionary
+    assert list(biobridge_node_info.keys()) == ['gene/protein',
+                                                'molecular_function',
+                                                'cellular_component',
+                                                'biological_process',
+                                                'drug',
+                                                'disease']
+    assert len(biobridge_node_info['gene/protein']) == 19162
+    assert len(biobridge_node_info['molecular_function']) == 10966
+    assert len(biobridge_node_info['cellular_component']) == 4013
+    assert len(biobridge_node_info['biological_process']) == 27478
+    assert len(biobridge_node_info['drug']) == 6948
+    assert len(biobridge_node_info['disease']) == 44133
+
+
+def test_load_existing_primekg(biobridge_primekg):
+    """
+    Test the loading method of the BioBridge-PrimeKG class by loading existing data in local.
+    """
+    # Load BioBridge-PrimeKG data
+    biobridge_primekg.load_data()
+    primekg_nodes = biobridge_primekg.get_primekg().get_nodes()
+    primekg_edges = biobridge_primekg.get_primekg().get_edges()
+    biobridge_data_config = biobridge_primekg.get_data_config()
+    biobridge_emb_dict = biobridge_primekg.get_node_embeddings()
+    biobridge_triplets = biobridge_primekg.get_primekg_triplets()
+    biobridge_splits = biobridge_primekg.get_train_test_split()
+    biobridge_node_info = biobridge_primekg.get_node_info_dict()
+
+    # Check if the local directories exists
+    assert os.path.exists(biobridge_primekg.primekg_dir)
+    assert os.path.exists(biobridge_primekg.local_dir)
+    # Check if downloaded and processed files exist
+    # PrimeKG files
+    files = ["nodes.tab", "primekg_nodes.tsv.gz",
+             "edges.csv", "primekg_edges.tsv.gz"]
+    for file in files:
+        path = f"{biobridge_primekg.primekg_dir}/{file}"
+        assert os.path.exists(path)
+    # BioBridge data config
+    assert os.path.exists(f"{biobridge_primekg.local_dir}/data_config.json")
+    # BioBridge embeddings
+    files = [
+        "protein.pkl",
+        "mf.pkl",
+        "cc.pkl",
+        "bp.pkl",
+        "drug.pkl",
+        "disease.pkl",
+        "embedding_dict.pkl"
+    ]
+    for file in files:
+        path = f"{biobridge_primekg.local_dir}/embeddings/{file}"
+        assert os.path.exists(path)
+    # BioBridge processed files
+    files = [
+        "protein.csv",
+        "mf.csv",
+        "cc.csv",
+        "bp.csv",
+        "drug.csv",
+        "disease.csv",
+        "triplet_full.tsv.gz",
+        "triplet_full_altered.tsv.gz",
+        "node_train.tsv.gz",
+        "triplet_train.tsv.gz",
+        "node_test.tsv.gz",
+        "triplet_test.tsv.gz",
+    ]
+    for file in files:
+        path = f"{biobridge_primekg.local_dir}/processed/{file}"
+        assert os.path.exists(path)
+    # Check processed PrimeKG dataframes
+    # Nodes
+    assert primekg_nodes is not None
+    assert len(primekg_nodes) > 0
+    assert primekg_nodes.shape[0] == 129375
+    # Edges
+    assert primekg_edges is not None
+    assert len(primekg_edges) > 0
+    assert primekg_edges.shape[0] == 8100498
+    # Check processed BioBridge data config
+    assert biobridge_data_config is not None
+    assert len(biobridge_data_config) > 0
+    assert len(biobridge_data_config['node_type']) == 10
+    assert len(biobridge_data_config['relation_type']) == 18
+    assert len(biobridge_data_config['emb_dim']) == 6
+    # Check processed BioBridge embeddings
+    assert biobridge_emb_dict is not None
+    assert len(biobridge_emb_dict) > 0
+    assert len(biobridge_emb_dict) == 85466
+    # Check processed BioBridge triplets
+    assert biobridge_triplets is not None
+    assert len(biobridge_triplets) > 0
+    assert biobridge_triplets.shape[0] == 3904610
+    assert list(biobridge_splits.keys()) == ['train', 'node_train', 'test', 'node_test']
+    assert len(biobridge_splits['train']) == 3510930
+    assert len(biobridge_splits['node_train']) == 76486
+    assert len(biobridge_splits['test']) == 393680
+    assert len(biobridge_splits['node_test']) == 8495
+    # Check node info dictionary
+    assert list(biobridge_node_info.keys()) == ['gene/protein',
+                                                'molecular_function',
+                                                'cellular_component',
+                                                'biological_process',
+                                                'drug',
+                                                'disease']
+    assert len(biobridge_node_info['gene/protein']) == 19162
+    assert len(biobridge_node_info['molecular_function']) == 10966
+    assert len(biobridge_node_info['cellular_component']) == 4013
+    assert len(biobridge_node_info['biological_process']) == 27478
+    assert len(biobridge_node_info['drug']) == 6948
+    assert len(biobridge_node_info['disease']) == 44133
+
+# def test_load_existing_primekg_with_negative_triplets(biobridge_primekg):
+#     """
+#     Test the loading method of the BioBridge-PrimeKG class by loading existing data in local.
+#     In addition, it builds negative triplets for training data.
+#     """
+#     # Load BioBridge-PrimeKG data
+#     # Using 1 negative sample per positive triplet
+#     biobridge_primekg.load_data(build_neg_triplest=True, n_neg_samples=1)
+#     biobridge_neg_triplets = biobridge_primekg.get_primekg_triplets_negative()
+
+#     # Check if the local directories exists
+#     assert os.path.exists(biobridge_primekg.primekg_dir)
+#     assert os.path.exists(biobridge_primekg.local_dir)
+#     # Check if downloaded and processed files exist
+#     path = f"{biobridge_primekg.local_dir}/processed/triplet_train_negative.tsv.gz"
+#     assert os.path.exists(path)
+#     # Check processed BioBridge triplets
+#     assert biobridge_neg_triplets is not None
+#     assert len(biobridge_neg_triplets) > 0
+#     assert biobridge_neg_triplets.shape[0] == 3510930
+#     assert len(biobridge_neg_triplets.negative_tail_index[0]) == 1

aiagents4pharma-1.8.3/aiagents4pharma/talk2knowledgegraphs/tests/test_datasets_dataset.py

@@ -0,0 +1,29 @@
+"""
+Test cases for datasets/dataset.py
+"""
+
+from ..datasets.dataset import Dataset
+
+class MockDataset(Dataset):
+    """
+    Mock dataset class for testing purposes.
+    """
+    def setup(self):
+        pass
+
+    def load_data(self):
+        pass
+
+def test_dataset_setup():
+    """
+    Test the setup method of the Dataset class.
+    """
+    dataset = MockDataset()
+    assert dataset.setup() is None
+
+def test_dataset_load_data():
+    """
+    Test the load_data method of the Dataset class.
+    """
+    dataset = MockDataset()
+    assert dataset.load_data() is None