aiagents4pharma 1.39.4__py3-none-any.whl → 1.40.0__py3-none-any.whl

aiagents4pharma/talk2knowledgegraphs/tests/test_tools_milvus_multimodal_subgraph_extraction.py

@@ -0,0 +1,413 @@
+"""
+Test cases for tools/milvus_multimodal_subgraph_extraction.py
+"""
+
+import importlib
+import unittest
+from unittest.mock import patch, MagicMock
+import numpy as np
+import pandas as pd
+from ..tools.milvus_multimodal_subgraph_extraction import MultimodalSubgraphExtractionTool
+
+class TestMultimodalSubgraphExtractionTool(unittest.TestCase):
+    """
+    Test cases for MultimodalSubgraphExtractionTool (Milvus)
+    """
+    def setUp(self):
+        self.tool = MultimodalSubgraphExtractionTool()
+        self.state = {
+            "uploaded_files": [],
+            "embedding_model": MagicMock(),
+            "topk_nodes": 5,
+            "topk_edges": 5,
+            "dic_source_graph": [{"name": "TestGraph"}],
+        }
+        self.prompt = "Find subgraph for test"
+        self.arg_data = {"extraction_name": "subkg_12345"}
+        self.cfg_db = MagicMock()
+        self.cfg_db.milvus_db.database_name = "testdb"
+        self.cfg_db.milvus_db.alias = "default"
+        self.cfg = MagicMock()
+        self.cfg.cost_e = 1.0
+        self.cfg.c_const = 1.0
+        self.cfg.root = 0
+        self.cfg.num_clusters = 1
+        self.cfg.pruning = True
+        self.cfg.verbosity_level = 0
+        self.cfg.search_metric_type = "L2"
+        self.cfg.node_colors_dict = {"gene/protein": "red"}
+
+    @patch("aiagents4pharma.talk2knowledgegraphs.tools."
+           "milvus_multimodal_subgraph_extraction.Collection")
+    @patch("aiagents4pharma.talk2knowledgegraphs.tools."
+           "milvus_multimodal_subgraph_extraction.MultimodalPCSTPruning")
+    @patch("pymilvus.connections")
+    def test_extract_multimodal_subgraph_wo_doc(self,
+                                                mock_connections,
+                                                mock_pcst,
+                                                mock_collection):
+        """
+        Test the multimodal subgraph extraction tool for only text as modality.
+        """
+
+        # Mock Milvus connection utilities
+        mock_connections.has_connection.return_value = True
+
+        # No uploaded_files (no doc)
+        self.state["uploaded_files"] = []
+        self.state["embedding_model"].embed_query.return_value = [0.1, 0.2, 0.3]
+        self.state["selections"] = {}
+
+        # Mock Collection for nodes and edges
+        colls = {}
+        colls["nodes"] = MagicMock()
+        colls["nodes"] = MagicMock()
+        colls["nodes"].query.return_value = [
+            {"node_index": 0,
+             "node_id": "id1",
+             "node_name": "JAK1",
+             "node_type": "gene/protein",
+             "feat": "featA",
+             "feat_emb": [0.1, 0.2, 0.3],
+             "desc": "descA",
+             "desc_emb": [0.1, 0.2, 0.3]},
+            {"node_index": 1,
+             "node_id": "id2",
+             "node_name": "JAK2",
+             "node_type": "gene/protein",
+             "feat": "featB",
+             "feat_emb": [0.4, 0.5, 0.6],
+             "desc": "descB",
+             "desc_emb": [0.4, 0.5, 0.6]}
+        ]
+        colls["nodes"].load.return_value = None
+
+        colls["edges"] = MagicMock()
+        colls["edges"].query.return_value = [
+            {"triplet_index": 0,
+             "head_id": "id1",
+             "head_index": 0,
+             "tail_id": "id2",
+             "tail_index": 1,
+             "edge_type": "gene/protein,ppi,gene/protein",
+             "display_relation": "ppi",
+             "feat": "featC",
+             "feat_emb": [0.7, 0.8, 0.9]}
+        ]
+        colls["edges"].load.return_value = None
+
+        def collection_side_effect(name):
+            """
+            Mock side effect for Collection to return nodes or edges based on name.
+            """
+            if "nodes" in name:
+                return colls["nodes"]
+            if "edges" in name:
+                return colls["edges"]
+            return None
+        mock_collection.side_effect = collection_side_effect
+
+        # Mock MultimodalPCSTPruning
+        mock_pcst_instance = MagicMock()
+        mock_pcst_instance.extract_subgraph.return_value = {
+            "nodes": pd.Series([1, 2]),
+            "edges": pd.Series([0])
+        }
+        mock_pcst.return_value = mock_pcst_instance
+
+        # Patch hydra.compose to return config objects
+        with patch("aiagents4pharma.talk2knowledgegraphs.tools."
+                   "milvus_multimodal_subgraph_extraction.hydra.initialize"), \
+             patch("aiagents4pharma.talk2knowledgegraphs.tools."
+                   "milvus_multimodal_subgraph_extraction.hydra.compose") as mock_compose:
+            mock_compose.return_value = MagicMock()
+            mock_compose.return_value.app.frontend = self.cfg_db
+            mock_compose.return_value.tools.multimodal_subgraph_extraction = self.cfg
+
+            response = self.tool.invoke(
+                input={"prompt": self.prompt,
+                       "tool_call_id": "subgraph_extraction_tool",
+                       "state": self.state,
+                       "arg_data": self.arg_data}
+            )
+
+        # Check tool message
+        self.assertEqual(response.update["messages"][-1].tool_call_id, "subgraph_extraction_tool")
+
+        # Check extracted subgraph dictionary
+        dic_extracted_graph = response.update["dic_extracted_graph"][0]
+        self.assertIsInstance(dic_extracted_graph, dict)
+        self.assertEqual(dic_extracted_graph["name"], self.arg_data["extraction_name"])
+        self.assertEqual(dic_extracted_graph["graph_source"], "TestGraph")
+        self.assertEqual(dic_extracted_graph["topk_nodes"], 5)
+        self.assertEqual(dic_extracted_graph["topk_edges"], 5)
+        self.assertIsInstance(dic_extracted_graph["graph_dict"], dict)
+        self.assertGreater(len(dic_extracted_graph["graph_dict"]["nodes"]), 0)
+        self.assertGreater(len(dic_extracted_graph["graph_dict"]["edges"]), 0)
+        self.assertIsInstance(dic_extracted_graph["graph_text"], str)
+        # Check if the nodes are in the graph_text
+        self.assertTrue(all(
+            n[0] in dic_extracted_graph["graph_text"].replace('"', '')
+            for subgraph_nodes in dic_extracted_graph["graph_dict"]["nodes"]
+            for n in subgraph_nodes
+        ))
+        # Check if the edges are in the graph_text
+        self.assertTrue(all(
+            ",".join([str(e[0])] + str(e[2]['label'][0]).split(",") + [str(e[1])])
+            in dic_extracted_graph["graph_text"].replace('"', '').\
+                replace("[", "").replace("]", "").replace("'", "")
+            for subgraph_edges in dic_extracted_graph["graph_dict"]["edges"]
+            for e in subgraph_edges
+        ))
+
+        # Another test for unknown collection
+        result = collection_side_effect("unknown")
+        self.assertIsNone(result)
+
+    @patch("aiagents4pharma.talk2knowledgegraphs.tools."
+           "milvus_multimodal_subgraph_extraction.Collection")
+    @patch("aiagents4pharma.talk2knowledgegraphs.tools."
+           "milvus_multimodal_subgraph_extraction.pd.read_excel")
+    @patch("aiagents4pharma.talk2knowledgegraphs.tools."
+           "milvus_multimodal_subgraph_extraction.MultimodalPCSTPruning")
+    @patch("pymilvus.connections")
+    def test_extract_multimodal_subgraph_w_doc(self,
+                                               mock_connections,
+                                               mock_pcst,
+                                               mock_read_excel,
+                                               mock_collection):
+        """
+        Test the multimodal subgraph extraction tool for text as modality, plus genes.
+        """
+        # Mock Milvus connection utilities
+        mock_connections.has_connection.return_value = True
+
+        # With uploaded_files (with doc)
+        self.state["uploaded_files"] = [
+            {"file_type": "multimodal", "file_path": "dummy.xlsx"}
+        ]
+        self.state["embedding_model"].embed_query.return_value = [0.1, 0.2, 0.3]
+        self.state["selections"] = {"gene/protein": ["JAK1", "JAK2"]}
+
+        # Mock pd.read_excel to return a dict of DataFrames
+        df = pd.DataFrame({
+            "name": ["JAK1", "JAK2"],
+            "node_type": ["gene/protein", "gene/protein"]
+        })
+        mock_read_excel.return_value = {"gene/protein": df}
+
+        # Mock Collection for nodes and edges
+        colls = {}
+        colls["nodes"] = MagicMock()
+        colls["nodes"] = MagicMock()
+        colls["nodes"].query.return_value = [
+            {"node_index": 0,
+             "node_id": "id1",
+             "node_name": "JAK1",
+             "node_type": "gene/protein",
+             "feat": "featA",
+             "feat_emb": [0.1, 0.2, 0.3],
+             "desc": "descA",
+             "desc_emb": [0.1, 0.2, 0.3]},
+            {"node_index": 1,
+             "node_id": "id2",
+             "node_name": "JAK2",
+             "node_type": "gene/protein",
+             "feat": "featB",
+             "feat_emb": [0.4, 0.5, 0.6],
+             "desc": "descB",
+             "desc_emb": [0.4, 0.5, 0.6]}
+        ]
+        colls["nodes"].load.return_value = None
+
+        colls["edges"] = MagicMock()
+        colls["edges"].query.return_value = [
+            {"triplet_index": 0,
+             "head_id": "id1",
+             "head_index": 0,
+             "tail_id": "id2",
+             "tail_index": 1,
+             "edge_type": "gene/protein,ppi,gene/protein",
+             "display_relation": "ppi",
+             "feat": "featC",
+             "feat_emb": [0.7, 0.8, 0.9]}
+        ]
+        colls["edges"].load.return_value = None
+
+        def collection_side_effect(name):
+            """
+            Mock side effect for Collection to return nodes or edges based on name.
+            """
+            if "nodes" in name:
+                return colls["nodes"]
+            if "edges" in name:
+                return colls["edges"]
+            return None
+        mock_collection.side_effect = collection_side_effect
+
+        # Mock MultimodalPCSTPruning
+        mock_pcst_instance = MagicMock()
+        mock_pcst_instance.extract_subgraph.return_value = {
+            "nodes": pd.Series([1, 2]),
+            "edges": pd.Series([0])
+        }
+        mock_pcst.return_value = mock_pcst_instance
+
+        # Patch hydra.compose to return config objects
+        with patch("aiagents4pharma.talk2knowledgegraphs.tools."
+                   "milvus_multimodal_subgraph_extraction.hydra.initialize"), \
+             patch("aiagents4pharma.talk2knowledgegraphs.tools."
+                   "milvus_multimodal_subgraph_extraction.hydra.compose") as mock_compose:
+            mock_compose.return_value = MagicMock()
+            mock_compose.return_value.app.frontend = self.cfg_db
+            mock_compose.return_value.tools.multimodal_subgraph_extraction = self.cfg
+
+            response = self.tool.invoke(
+                input={"prompt": self.prompt,
+                       "tool_call_id": "subgraph_extraction_tool",
+                       "state": self.state,
+                       "arg_data": self.arg_data}
+            )
+
+                # Check tool message
+        self.assertEqual(response.update["messages"][-1].tool_call_id, "subgraph_extraction_tool")
+
+        # Check extracted subgraph dictionary
+        dic_extracted_graph = response.update["dic_extracted_graph"][0]
+        self.assertIsInstance(dic_extracted_graph, dict)
+        self.assertEqual(dic_extracted_graph["name"], self.arg_data["extraction_name"])
+        self.assertEqual(dic_extracted_graph["graph_source"], "TestGraph")
+        self.assertEqual(dic_extracted_graph["topk_nodes"], 5)
+        self.assertEqual(dic_extracted_graph["topk_edges"], 5)
+        self.assertIsInstance(dic_extracted_graph["graph_dict"], dict)
+        self.assertGreater(len(dic_extracted_graph["graph_dict"]["nodes"]), 0)
+        self.assertGreater(len(dic_extracted_graph["graph_dict"]["edges"]), 0)
+        self.assertIsInstance(dic_extracted_graph["graph_text"], str)
+        # Check if the nodes are in the graph_text
+        self.assertTrue(all(
+            n[0] in dic_extracted_graph["graph_text"].replace('"', '')
+            for subgraph_nodes in dic_extracted_graph["graph_dict"]["nodes"]
+            for n in subgraph_nodes
+        ))
+        # Check if the edges are in the graph_text
+        self.assertTrue(all(
+            ",".join([str(e[0])] + str(e[2]['label'][0]).split(",") + [str(e[1])])
+            in dic_extracted_graph["graph_text"].replace('"', '').\
+                replace("[", "").replace("]", "").replace("'", "")
+            for subgraph_edges in dic_extracted_graph["graph_dict"]["edges"]
+            for e in subgraph_edges
+        ))
+
+        # Another test for unknown collection
+        result = collection_side_effect("unknown")
+        self.assertIsNone(result)
+
+    def test_extract_multimodal_subgraph_wo_doc_gpu(self):
+        """
+        Test the multimodal subgraph extraction tool for only text as modality,
+        simulating GPU (cudf/cupy) environment.
+        """
+        module_name = "aiagents4pharma.talk2knowledgegraphs.tools."+\
+            "milvus_multimodal_subgraph_extraction"
+        with patch.dict("sys.modules", {"cupy": np, "cudf": pd}):
+            mod = importlib.reload(importlib.import_module(module_name))
+            # Patch Collection and MultimodalPCSTPruning after reload
+            with patch(f"{module_name}.Collection") as mock_collection, \
+                 patch(f"{module_name}.MultimodalPCSTPruning") as mock_pcst, \
+                 patch("pymilvus.connections") as mock_connections:
+                # Setup mocks as in the original test
+                mock_connections.has_connection.return_value = True
+                colls = {}
+                colls["nodes"] = MagicMock()
+                colls["nodes"].query.return_value = [
+                    {"node_index": 0,
+                     "node_id": "id1",
+                     "node_name": "JAK1",
+                     "node_type": "gene/protein",
+                     "feat": "featA",
+                     "feat_emb": [0.1, 0.2, 0.3],
+                     "desc": "descA",
+                     "desc_emb": [0.1, 0.2, 0.3]},
+                    {"node_index": 1,
+                     "node_id": "id2",
+                     "node_name": "JAK2",
+                     "node_type": "gene/protein",
+                     "feat": "featB",
+                     "feat_emb": [0.4, 0.5, 0.6],
+                     "desc": "descB",
+                     "desc_emb": [0.4, 0.5, 0.6]}
+                ]
+                colls["nodes"].load.return_value = None
+                colls["edges"] = MagicMock()
+                colls["edges"].query.return_value = [
+                    {"triplet_index": 0,
+                     "head_id": "id1",
+                     "head_index": 0,
+                     "tail_id": "id2",
+                     "tail_index": 1,
+                     "edge_type": "gene/protein,ppi,gene/protein",
+                     "display_relation": "ppi",
+                     "feat": "featC",
+                     "feat_emb": [0.7, 0.8, 0.9]}
+                ]
+                colls["edges"].load.return_value = None
+                def collection_side_effect(name):
+                    if "nodes" in name:
+                        return colls["nodes"]
+                    if "edges" in name:
+                        return colls["edges"]
+                    return None
+                mock_collection.side_effect = collection_side_effect
+                mock_pcst_instance = MagicMock()
+                mock_pcst_instance.extract_subgraph.return_value = {
+                    "nodes": pd.Series([1, 2]),
+                    "edges": pd.Series([0])
+                }
+                mock_pcst.return_value = mock_pcst_instance
+                # Setup config mocks
+                tool_cls = getattr(mod, "MultimodalSubgraphExtractionTool")
+                tool = tool_cls()
+
+                # Patch hydra.compose
+                with patch(f"{module_name}.hydra.initialize"), \
+                     patch(f"{module_name}.hydra.compose") as mock_compose:
+                    mock_compose.return_value = MagicMock()
+                    mock_compose.return_value.app.frontend = self.cfg_db
+                    mock_compose.return_value.tools.multimodal_subgraph_extraction = self.cfg
+                    self.state["embedding_model"].embed_query.return_value = [0.1, 0.2, 0.3]
+                    self.state["selections"] = {}
+                    response = tool.invoke(
+                        input={"prompt": self.prompt,
+                               "tool_call_id": "subgraph_extraction_tool",
+                               "state": self.state,
+                               "arg_data": self.arg_data}
+                    )
+                # Check tool message
+                self.assertEqual(response.update["messages"][-1].tool_call_id,
+                                 "subgraph_extraction_tool")
+                dic_extracted_graph = response.update["dic_extracted_graph"][0]
+                self.assertIsInstance(dic_extracted_graph, dict)
+                self.assertEqual(dic_extracted_graph["name"], self.arg_data["extraction_name"])
+                self.assertEqual(dic_extracted_graph["graph_source"], "TestGraph")
+                self.assertEqual(dic_extracted_graph["topk_nodes"], 5)
+                self.assertEqual(dic_extracted_graph["topk_edges"], 5)
+                self.assertIsInstance(dic_extracted_graph["graph_dict"], dict)
+                self.assertGreater(len(dic_extracted_graph["graph_dict"]["nodes"]), 0)
+                self.assertGreater(len(dic_extracted_graph["graph_dict"]["edges"]), 0)
+                self.assertIsInstance(dic_extracted_graph["graph_text"], str)
+                self.assertTrue(all(
+                    n[0] in dic_extracted_graph["graph_text"].replace('"', '')
+                    for subgraph_nodes in dic_extracted_graph["graph_dict"]["nodes"]
+                    for n in subgraph_nodes
+                ))
+                self.assertTrue(all(
+                    ",".join([str(e[0])] + str(e[2]['label'][0]).split(",") + [str(e[1])])
+                    in dic_extracted_graph["graph_text"].replace('"', '').\
+                        replace("[", "").replace("]", "").replace("'", "")
+                    for subgraph_edges in dic_extracted_graph["graph_dict"]["edges"]
+                    for e in subgraph_edges
+                ))
+
+                # Another test for unknown collection
+                result = collection_side_effect("unknown")
+                self.assertIsNone(result)

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_ols.py

@@ -42,11 +42,11 @@ def test_enrich_documents(enrich_obj):
         "XYZ_0000000",
     ]
     descriptions = enrich_obj.enrich_documents(ols_terms)
-    assert descriptions[0].startswith(CL_DESC)
-    assert descriptions[1].startswith(GO_DESC)
-    assert descriptions[2].startswith(UBERON_DESC)
-    assert descriptions[3].startswith(HP_DESC)
-    assert descriptions[4].startswith(MONDO_DESC)
+    assert CL_DESC in descriptions[0]
+    assert GO_DESC in descriptions[1]
+    assert UBERON_DESC in descriptions[2]
+    assert HP_DESC in descriptions[3]
+    assert MONDO_DESC in descriptions[4]
     assert descriptions[5] is None
 
 
@@ -61,9 +61,9 @@ def test_enrich_documents_with_rag(enrich_obj):
         "XYZ_0000000",
     ]
     descriptions = enrich_obj.enrich_documents_with_rag(ols_terms, None)
-    assert descriptions[0].startswith(CL_DESC)
-    assert descriptions[1].startswith(GO_DESC)
-    assert descriptions[2].startswith(UBERON_DESC)
-    assert descriptions[3].startswith(HP_DESC)
-    assert descriptions[4].startswith(MONDO_DESC)
+    assert CL_DESC in descriptions[0]
+    assert GO_DESC in descriptions[1]
+    assert UBERON_DESC in descriptions[2]
+    assert HP_DESC in descriptions[3]
+    assert MONDO_DESC in descriptions[4]
     assert descriptions[5] is None

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_extractions_milvus_multimodal_pcst.py

@@ -0,0 +1,175 @@
+"""
+Test cases for tools/utils/extractions/milvus_multimodal_pcst.py
+"""
+
+import importlib
+import sys
+import unittest
+from unittest.mock import patch, MagicMock, mock_open
+import numpy as np
+import pandas as pd
+from ..utils.extractions.milvus_multimodal_pcst import MultimodalPCSTPruning
+
+class TestMultimodalPCSTPruning(unittest.TestCase):
+    """
+    Test cases for MultimodalPCSTPruning class (Milvus-based PCST pruning).
+    """
+    def setUp(self):
+        # Patch cupy and cudf to simulate GPU environment
+        patcher_cupy = patch.dict('sys.modules', {'cupy': MagicMock(), 'cudf': MagicMock()})
+        patcher_cupy.start()
+        self.addCleanup(patcher_cupy.stop)
+
+        # Patch pcst_fast
+        self.pcst_fast_patcher = patch("aiagents4pharma.talk2knowledgegraphs.utils."
+                                       "extractions.milvus_multimodal_pcst.pcst_fast")
+        self.mock_pcst_fast = self.pcst_fast_patcher.start()
+        self.addCleanup(self.pcst_fast_patcher.stop)
+        self.mock_pcst_fast.pcst_fast.return_value = ([0, 1], [0])
+
+        # Patch Collection
+        self.collection_patcher = patch("aiagents4pharma.talk2knowledgegraphs.utils."
+                                        "extractions.milvus_multimodal_pcst.Collection")
+        self.mock_collection = self.collection_patcher.start()
+        self.addCleanup(self.collection_patcher.stop)
+
+        # Patch open for cache_edge_index_path
+        self.open_patcher = patch('builtins.open', mock_open(read_data='[[0,1],[1,2]]'))
+        self.mock_open = self.open_patcher.start()
+        self.addCleanup(self.open_patcher.stop)
+
+        # Patch pickle.load to return a numpy array for edge_index
+        self.pickle_patcher = patch("aiagents4pharma.talk2knowledgegraphs.utils."
+                                    "extractions.milvus_multimodal_pcst.pickle")
+        self.mock_pickle = self.pickle_patcher.start()
+        self.addCleanup(self.pickle_patcher.stop)
+        self.mock_pickle.load.return_value = np.array([[0, 1], [1, 2]])
+
+        # Setup config mock
+        self.cfg = MagicMock()
+        self.cfg.milvus_db.database_name = "testdb"
+        self.cfg.milvus_db.cache_edge_index_path = "dummy_cache.pkl"
+
+        # Setup Collection mocks
+        node_coll = MagicMock()
+        node_coll.num_entities = 2
+        node_coll.search.return_value = [[MagicMock(id=0), MagicMock(id=1)]]
+        edge_coll = MagicMock()
+        edge_coll.num_entities = 2
+        edge_coll.search.return_value = [[MagicMock(id=0, score=1.0), MagicMock(id=1, score=0.5)]]
+        self.mock_collection.side_effect = lambda name: node_coll if "nodes" in name else edge_coll
+
+    def test_extract_subgraph_use_description_true(self):
+        """
+        Test the extract_subgraph method of MultimodalPCSTPruning with use_description=True.
+        """
+        # Create instance
+        pcst = MultimodalPCSTPruning(
+            topk=3, topk_e=3, cost_e=0.5, c_const=0.01, root=-1,
+            num_clusters=1, pruning="gw", verbosity_level=0, use_description=True, metric_type="IP"
+        )
+        # Dummy embeddings
+        text_emb = [0.1, 0.2, 0.3]
+        query_emb = [0.1, 0.2, 0.3]
+        modality = "gene/protein"
+
+        # Call extract_subgraph
+        result = pcst.extract_subgraph(text_emb, query_emb, modality, self.cfg)
+
+        # Assertions
+        self.assertIn("nodes", result)
+        self.assertIn("edges", result)
+        self.assertGreaterEqual(len(result["nodes"]), 0)
+        self.assertGreaterEqual(len(result["edges"]), 0)
+
+    def test_extract_subgraph_use_description_false(self):
+        """
+        Test the extract_subgraph method of MultimodalPCSTPruning with use_description=False.
+        """
+        # Create instance
+        pcst = MultimodalPCSTPruning(
+            topk=3, topk_e=3, cost_e=0.5, c_const=0.01, root=-1,
+            num_clusters=1, pruning="gw", verbosity_level=0, use_description=False, metric_type="IP"
+        )
+        # Dummy embeddings
+        text_emb = [0.1, 0.2, 0.3]
+        query_emb = [0.1, 0.2, 0.3]
+        modality = "gene/protein"
+
+        # Call extract_subgraph
+        result = pcst.extract_subgraph(text_emb, query_emb, modality, self.cfg)
+
+        # Assertions
+        self.assertIn("nodes", result)
+        self.assertIn("edges", result)
+        self.assertGreaterEqual(len(result["nodes"]), 0)
+        self.assertGreaterEqual(len(result["edges"]), 0)
+
+    def test_extract_subgraph_with_virtual_vertices(self):
+        """
+        Test get_subgraph_nodes_edges with virtual vertices present (len(virtual_vertices) > 0).
+        """
+        pcst = MultimodalPCSTPruning(
+            topk=3, topk_e=3, cost_e=0.5, c_const=0.01, root=-1,
+            num_clusters=1, pruning="gw", verbosity_level=0, use_description=True, metric_type="IP"
+        )
+        # Simulate num_nodes = 2, vertices contains [0, 1, 2, 3] (2 and 3 are virtual)
+        num_nodes = 2
+        # vertices: [0, 1, 2, 3] (2 and 3 are virtual)
+        vertices = np.array([0, 1, 2, 3])
+        # edges_dict simulates prior edges and edge_index
+        edges_dict = {
+            "edges": np.array([0, 1, 2]),
+            "num_prior_edges": 2,
+            "edge_index": np.array([[0, 1, 2, 3], [1, 2, 3, 4]])
+        }
+        # mapping simulates mapping for edges and nodes
+        mapping = {
+            "edges": {0: 0, 1: 1},
+            "nodes": {2: 2, 3: 3}
+        }
+
+        # Call extract_subgraph
+        result = pcst.get_subgraph_nodes_edges(num_nodes, vertices, edges_dict, mapping)
+
+        # Assertions
+        self.assertIn("nodes", result)
+        self.assertIn("edges", result)
+        self.assertGreaterEqual(len(result["nodes"]), 0)
+        self.assertGreaterEqual(len(result["edges"]), 0)
+        # Check that virtual edges are included
+        self.assertTrue(any(e in [2, 3] for e in result["edges"]))
+
+    def test_gpu_import_branch(self):
+        """
+        Test coverage for GPU import branch by patching sys.modules to mock cupy and
+        cudf as numpy and pandas.
+        """
+        module_name = "aiagents4pharma.talk2knowledgegraphs.utils" + \
+            ".extractions.milvus_multimodal_pcst"
+        with patch.dict("sys.modules", {"cupy": np, "cudf": pd}):
+            # Reload the module to trigger the GPU branch
+            mod = importlib.reload(sys.modules[module_name])
+                # Patch Collection, pcst_fast, and pickle after reload
+            with patch(f"{module_name}.Collection", self.mock_collection), \
+                patch(f"{module_name}.pcst_fast", self.mock_pcst_fast), \
+                patch(f"{module_name}.pickle", self.mock_pickle):
+                pcst_pruning_cls = getattr(mod, "MultimodalPCSTPruning")
+                pcst = pcst_pruning_cls(
+                    topk=3, topk_e=3, cost_e=0.5, c_const=0.01, root=-1,
+                    num_clusters=1, pruning="gw", verbosity_level=0, use_description=True,
+                    metric_type="IP"
+                )
+                # Dummy embeddings
+                text_emb = [0.1, 0.2, 0.3]
+                query_emb = [0.1, 0.2, 0.3]
+                modality = "gene/protein"
+
+                # Call extract_subgraph
+                result = pcst.extract_subgraph(text_emb, query_emb, modality, self.cfg)
+
+                # Assertions
+                self.assertIn("nodes", result)
+                self.assertIn("edges", result)
+                self.assertGreaterEqual(len(result["nodes"]), 0)
+                self.assertGreaterEqual(len(result["edges"]), 0)

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_pubchem_utils.py

@@ -4,6 +4,17 @@ Test cases for utils/pubchem_utils.py
 
 from ..utils import pubchem_utils
 
+def test_cas_rn2pubchem_cid():
+    """
+    Test the casRN2pubchem_cid function.
+
+    The CAS RN for ethyl cabonate is 105-58-8.
+    The PubChem CID for ethyl cabonate is 7766.
+    """
+    casrn = "105-58-8"
+    pubchem_cid = pubchem_utils.cas_rn2pubchem_cid(casrn)
+    assert pubchem_cid == 7766
+
 def test_external_id2pubchem_cid():
     """
     Test the external_id2pubchem_cid function.

aiagents4pharma/talk2knowledgegraphs/tools/__init__.py

@@ -3,5 +3,6 @@ This file is used to import all the models in the package.
 '''
 from . import subgraph_extraction
 from . import multimodal_subgraph_extraction
+from . import milvus_multimodal_subgraph_extraction
 from . import subgraph_summarization
 from . import graphrag_reasoning