aiagents4pharma 1.8.2__py3-none-any.whl → 1.9.0__py3-none-any.whl

aiagents4pharma/talk2biomodels/tools/simulate_model.py

@@ -52,10 +52,10 @@ class TimeSpeciesNameConcentration:
 class RecurringData:
     """
     Dataclass for storing the species and time data 
-    on recurring basis.
+    on reocurring basis.
     """
     data: List[TimeSpeciesNameConcentration] = Field(
-        description="species and time data on recurring basis",
+        description="species and time data on reocurring basis",
         default=None)
 
 @dataclass
@@ -68,12 +68,15 @@ class ArgumentData:
         description="species name and initial concentration data",
         default=None)
     recurring_data: RecurringData = Field(
-        description="species and time data on recurring basis",
+        description="species and time data on reocurring basis",
         default=None)
+    simulation_name: str = Field(
+        description="""An AI assigned `_` separated name of
+        the simulation based on human query""")
 
 def add_rec_events(model_object, recurring_data):
     """
-    Add recurring events to the model.
+    Add reocurring events to the model.
     """
     for row in recurring_data.data:
         tp, sn, sc = row.time, row.species_name, row.species_concentration
@@ -86,9 +89,12 @@ 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 recurring data",
-                                        default=None)
+    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)
     tool_call_id: Annotated[str, InjectedToolCallId]
     state: Annotated[dict, InjectedState]
 
@@ -153,12 +159,20 @@ class SimulateModelTool(BaseTool):
             interval=interval
             )
 
+        dic_simulated_data = {
+            'name': arg_data.simulation_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 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():
+            "model_id": [sys_bio_model.biomodel_id],
+            "sbml_file_path": [sbml_file_path],
+            "dic_simulated_data": [dic_simulated_data],
+            }.items():
             if value:
                 dic_updated_state_for_model[key] = value
 
@@ -166,11 +180,11 @@ class SimulateModelTool(BaseTool):
         return Command(
                 update=dic_updated_state_for_model|{
                 # update the state keys
-                "dic_simulated_data": df.to_dict(),
+                # "dic_simulated_data": df.to_dict(),
                 # update the message history
                 "messages": [
                     ToolMessage(
-                        content="Simulation results are ready.",
+                        content=f"Simulation results of {arg_data.simulation_name}",
                         tool_call_id=tool_call_id
                         )
                     ],

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/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/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

aiagents4pharma/talk2knowledgegraphs/tests/test_datasets_primekg.py

@@ -0,0 +1,73 @@
+"""
+Test cases for datasets/primekg_loader.py
+"""
+
+import os
+import shutil
+import pytest
+from ..datasets.primekg import PrimeKG
+
+# Remove the data folder for testing if it exists
+LOCAL_DIR = "../data/primekg_test/"
+shutil.rmtree(LOCAL_DIR, ignore_errors=True)
+
+@pytest.fixture(name="primekg")
+def primekg_fixture():
+    """
+    Fixture for creating an instance of PrimeKG.
+    """
+    return PrimeKG(local_dir=LOCAL_DIR)
+
+def test_download_primekg(primekg):
+    """
+    Test the loading method of the PrimeKG class by downloading PrimeKG from server.
+    """
+    # Load PrimeKG data
+    primekg.load_data()
+    primekg_nodes = primekg.get_nodes()
+    primekg_edges = primekg.get_edges()
+
+    # Check if the local directory exists
+    assert os.path.exists(primekg.local_dir)
+    # Check if downloaded and processed files exist
+    files = ["nodes.tab", f"{primekg.name}_nodes.tsv.gz",
+             "edges.csv", f"{primekg.name}_edges.tsv.gz"]
+    for file in files:
+        path = f"{primekg.local_dir}/{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
+
+def test_load_existing_primekg(primekg):
+    """
+    Test the loading method of the PrimeKG class by loading existing PrimeKG in local.
+    """
+    # Load PrimeKG data
+    primekg.load_data()
+    primekg_nodes = primekg.get_nodes()
+    primekg_edges = primekg.get_edges()
+
+    # Check if the local directory exists
+    assert os.path.exists(primekg.local_dir)
+    # Check if downloaded and processed files exist
+    files = ["nodes.tab", f"{primekg.name}_nodes.tsv.gz",
+             "edges.csv", f"{primekg.name}_edges.tsv.gz"]
+    for file in files:
+        path = f"{primekg.local_dir}/{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

aiagents4pharma/talk2knowledgegraphs/tests/test_datasets_starkqa_primekg.py

@@ -0,0 +1,116 @@
+"""
+Test cases for datasets/starkqa_primekg_loader.py
+"""
+
+import os
+import shutil
+import pytest
+from ..datasets.starkqa_primekg import StarkQAPrimeKG
+
+# Remove the data folder for testing if it exists
+LOCAL_DIR = "../data/starkqa_primekg_test/"
+shutil.rmtree(LOCAL_DIR, ignore_errors=True)
+
+@pytest.fixture(name="starkqa_primekg")
+def starkqa_primekg_fixture():
+    """
+    Fixture for creating an instance of StarkQAPrimeKGData.
+    """
+    return StarkQAPrimeKG(local_dir=LOCAL_DIR)
+
+def test_download_starkqa_primekg(starkqa_primekg):
+    """
+    Test the loading method of the StarkQAPrimeKGLoaderTool class by downloading files
+    from HuggingFace Hub.
+    """
+    # Load StarkQA PrimeKG data
+    starkqa_primekg.load_data()
+    starkqa_df = starkqa_primekg.get_starkqa()
+    primekg_node_info = starkqa_primekg.get_starkqa_node_info()
+    split_idx = starkqa_primekg.get_starkqa_split_indicies()
+    query_embeddings = starkqa_primekg.get_query_embeddings()
+    node_embeddings = starkqa_primekg.get_node_embeddings()
+
+    # Check if the local directory exists
+    assert os.path.exists(starkqa_primekg.local_dir)
+    # Check if downloaded files exist in the local directory
+    files = ['qa/prime/split/test-0.1.index',
+             'qa/prime/split/test.index',
+             'qa/prime/split/train.index',
+             'qa/prime/split/val.index',
+             'qa/prime/stark_qa/stark_qa.csv',
+             'qa/prime/stark_qa/stark_qa_human_generated_eval.csv',
+             'skb/prime/processed.zip']
+    for file in files:
+        path = f"{starkqa_primekg.local_dir}/{file}"
+        assert os.path.exists(path)
+    # Check dataframe
+    assert starkqa_df is not None
+    assert len(starkqa_df) > 0
+    assert starkqa_df.shape[0] == 11204
+    # Check node information
+    assert primekg_node_info is not None
+    assert len(primekg_node_info) == 129375
+    # Check split indices
+    assert list(split_idx.keys()) == ['train', 'val', 'test', 'test-0.1']
+    assert len(split_idx['train']) == 6162
+    assert len(split_idx['val']) == 2241
+    assert len(split_idx['test']) == 2801
+    assert len(split_idx['test-0.1']) == 280
+    # Check query embeddings
+    assert query_embeddings is not None
+    assert len(query_embeddings) == 11204
+    assert query_embeddings[0].shape[1] == 1536
+    # Check node embeddings
+    assert node_embeddings is not None
+    assert len(node_embeddings) == 129375
+    assert node_embeddings[0].shape[1] == 1536
+
+def test_load_existing_starkqa_primekg(starkqa_primekg):
+    """
+
+    Test the loading method of the StarkQAPrimeKGLoaderTool class by loading existing files
+    in the local directory.
+    """
+    # Load StarkQA PrimeKG data
+    starkqa_primekg.load_data()
+    starkqa_df = starkqa_primekg.get_starkqa()
+    primekg_node_info = starkqa_primekg.get_starkqa_node_info()
+    split_idx = starkqa_primekg.get_starkqa_split_indicies()
+    query_embeddings = starkqa_primekg.get_query_embeddings()
+    node_embeddings = starkqa_primekg.get_node_embeddings()
+
+    # Check if the local directory exists
+    assert os.path.exists(starkqa_primekg.local_dir)
+    # Check if downloaded and processed files exist
+    files = ['qa/prime/split/test-0.1.index',
+             'qa/prime/split/test.index',
+             'qa/prime/split/train.index',
+             'qa/prime/split/val.index',
+             'qa/prime/stark_qa/stark_qa.csv',
+             'qa/prime/stark_qa/stark_qa_human_generated_eval.csv',
+             'skb/prime/processed.zip']
+    for file in files:
+        path = f"{starkqa_primekg.local_dir}/{file}"
+        assert os.path.exists(path)
+    # Check dataframe
+    assert starkqa_df is not None
+    assert len(starkqa_df) > 0
+    assert starkqa_df.shape[0] == 11204
+    # Check node information
+    assert primekg_node_info is not None
+    assert len(primekg_node_info) == 129375
+    # Check split indices
+    assert list(split_idx.keys()) == ['train', 'val', 'test', 'test-0.1']
+    assert len(split_idx['train']) == 6162
+    assert len(split_idx['val']) == 2241
+    assert len(split_idx['test']) == 2801
+    assert len(split_idx['test-0.1']) == 280
+    # Check query embeddings
+    assert query_embeddings is not None
+    assert len(query_embeddings) == 11204
+    assert query_embeddings[0].shape[1] == 1536
+    # Check node embeddings
+    assert node_embeddings is not None
+    assert len(node_embeddings) == 129375
+    assert node_embeddings[0].shape[1] == 1536

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_embeddings.py

@@ -0,0 +1,47 @@
+"""
+Test cases for utils/embeddings/embeddings.py
+"""
+
+import pytest
+from ..utils.embeddings.embeddings import Embeddings
+
+class TestEmbeddings(Embeddings):
+    """Test implementation of the Embeddings interface for testing purposes."""
+
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        return [[0.1, 0.2, 0.3] for _ in texts]
+
+    def embed_query(self, text: str) -> list[float]:
+        return [0.1, 0.2, 0.3]
+
+def test_embed_documents():
+    """Test embedding documents using the Embeddings interface."""
+    embeddings = TestEmbeddings()
+    texts = ["text1", "text2"]
+    result = embeddings.embed_documents(texts)
+    assert result == [[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]
+
+
+def test_embed_query():
+    """Test embedding a query using the Embeddings interface."""
+    embeddings = TestEmbeddings()
+    text = "query"
+    result = embeddings.embed_query(text)
+    assert result == [0.1, 0.2, 0.3]
+
+@pytest.mark.asyncio
+async def test_aembed_documents():
+    """Test asynchronous embedding of documents using the Embeddings interface."""
+    embeddings = TestEmbeddings()
+    texts = ["text1", "text2"]
+    result = await embeddings.aembed_documents(texts)
+    assert result == [[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]
+
+
+@pytest.mark.asyncio
+async def test_aembed_query():
+    """Test asynchronous embedding of a query using the Embeddings interface."""
+    embeddings = TestEmbeddings()
+    text = "query"
+    result = await embeddings.aembed_query(text)
+    assert result == [0.1, 0.2, 0.3]

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_huggingface.py

@@ -0,0 +1,45 @@
+"""
+Test cases for utils/embeddings/huggingface.py
+"""
+
+import pytest
+from ..utils.embeddings.huggingface import EmbeddingWithHuggingFace
+
+@pytest.fixture(name="embedding_model")
+def embedding_model_fixture():
+    """Return the configuration object for the HuggingFace embedding model and model object"""
+    return EmbeddingWithHuggingFace(
+        model_name="NeuML/pubmedbert-base-embeddings",
+        model_cache_dir="../../cache",
+        truncation=True,
+    )
+
+def test_embedding_with_huggingface_embed_documents(embedding_model):
+    """Test embedding documents using the EmbeddingWithHuggingFace class."""
+    # Perform embedding
+    texts = ["Adalimumab", "Infliximab", "Vedolizumab"]
+    result = embedding_model.embed_documents(texts)
+    # Check the result
+    assert len(result) == 3
+    assert len(result[0]) == 768
+
+def test_embedding_with_huggingface_embed_query(embedding_model):
+    """Test embedding a query using the EmbeddingWithHuggingFace class."""
+    # Perform embedding
+    text = "Adalimumab"
+    result = embedding_model.embed_query(text)
+    # Check the result
+    assert len(result) == 768
+
+
+def test_embedding_with_huggingface_failed():
+    """Test embedding documents using the EmbeddingWithHuggingFace class."""
+    # Check if the model is available on HuggingFace Hub
+    model_name = "aiagents4pharma/embeddings"
+    err_msg = f"Model {model_name} is not available on HuggingFace Hub."
+    with pytest.raises(ValueError, match=err_msg):
+        EmbeddingWithHuggingFace(
+            model_name=model_name,
+            model_cache_dir="../../cache",
+            truncation=True,
+        )