aiagents4pharma 1.43.0__py3-none-any.whl → 1.45.0__py3-none-any.whl

aiagents4pharma/talk2knowledgegraphs/tests/test_datasets_dataset.py

@@ -4,16 +4,19 @@ 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.
@@ -21,6 +24,7 @@ def test_dataset_setup():
     dataset = MockDataset()
     assert dataset.setup() is None
 
+
 def test_dataset_load_data():
     """
     Test the load_data method of the Dataset class.

aiagents4pharma/talk2knowledgegraphs/tests/test_datasets_primekg.py

@@ -4,13 +4,16 @@ 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():
     """
@@ -18,6 +21,7 @@ def primekg_fixture():
     """
     return PrimeKG(local_dir=LOCAL_DIR)
 
+
 def test_download_primekg(primekg):
     """
     Test the loading method of the PrimeKG class by downloading PrimeKG from server.
@@ -30,8 +34,12 @@ def test_download_primekg(primekg):
     # 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"]
+    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)
@@ -45,6 +53,7 @@ def test_download_primekg(primekg):
     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.
@@ -57,8 +66,12 @@ def test_load_existing_primekg(primekg):
     # 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"]
+    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)

aiagents4pharma/talk2knowledgegraphs/tests/test_datasets_starkqa_primekg.py

@@ -4,13 +4,16 @@ 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():
     """
@@ -18,6 +21,7 @@ def starkqa_primekg_fixture():
     """
     return StarkQAPrimeKG(local_dir=LOCAL_DIR)
 
+
 def test_download_starkqa_primekg(starkqa_primekg):
     """
     Test the loading method of the StarkQAPrimeKGLoaderTool class by downloading files
@@ -34,13 +38,15 @@ def test_download_starkqa_primekg(starkqa_primekg):
     # 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']
+    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)
@@ -52,11 +58,11 @@ def test_download_starkqa_primekg(starkqa_primekg):
     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
+    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
@@ -66,6 +72,7 @@ def test_download_starkqa_primekg(starkqa_primekg):
     assert len(node_embeddings) == 129375
     assert node_embeddings[0].shape[1] == 1536
 
+
 def test_load_existing_starkqa_primekg(starkqa_primekg):
     """
 
@@ -83,13 +90,15 @@ def test_load_existing_starkqa_primekg(starkqa_primekg):
     # 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']
+    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)
@@ -101,11 +110,11 @@ def test_load_existing_starkqa_primekg(starkqa_primekg):
     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
+    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

aiagents4pharma/talk2knowledgegraphs/tests/test_tools_graphrag_reasoning.py

@@ -5,6 +5,7 @@ Test cases for tools/graphrag_reasoning.py
 import pytest
 from langchain_core.messages import HumanMessage
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+
 from ..agents.t2kg_agent import get_app
 
 # Define the data path
@@ -52,58 +53,104 @@ def input_dict_fixture():
                 "topk_nodes": 3,
                 "topk_edges": 3,
                 "graph_dict": {
-                    'nodes': [('IFNG_(3495)', {}), 
-                              ('IKBKG_(3672)', {}),
-                              ('ATG16L1_(6661)', {}),
-                              ('inflammatory bowel disease_(28158)', {}),
-                              ('Crohn ileitis and jejunitis_(35814)', {}),
-                              ("Crohn's colitis_(83770)", {})],
-                    'edges': [('IFNG_(3495)', 'inflammatory bowel disease_(28158)', 
-                               {'relation': ['gene/protein', 'associated with', 'disease'],
-                                'label': ['gene/protein', 'associated with', 'disease']}), 
-                              ('IFNG_(3495)', "Crohn's colitis_(83770)",
-                               {'relation': ['gene/protein', 'associated with', 'disease'],
-                                'label': ['gene/protein', 'associated with', 'disease']}),
-                              ('IFNG_(3495)', 'Crohn ileitis and jejunitis_(35814)',
-                               {'relation': ['gene/protein', 'associated with', 'disease'],
-                                'label': ['gene/protein', 'associated with', 'disease']}), 
-                              ('ATG16L1_(6661)', 'IKBKG_(3672)',
-                               {'relation': ['gene/protein', 'ppi', 'gene/protein'],
-                                'label': ['gene/protein', 'ppi', 'gene/protein']}), 
-                              ("Crohn's colitis_(83770)", 'ATG16L1_(6661)',
-                               {'relation': ['disease', 'associated with', 'gene/protein'],
-                                'label': ['disease', 'associated with', 'gene/protein']})]},
+                    "nodes": [
+                        ("IFNG_(3495)", {}),
+                        ("IKBKG_(3672)", {}),
+                        ("ATG16L1_(6661)", {}),
+                        ("inflammatory bowel disease_(28158)", {}),
+                        ("Crohn ileitis and jejunitis_(35814)", {}),
+                        ("Crohn's colitis_(83770)", {}),
+                    ],
+                    "edges": [
+                        (
+                            "IFNG_(3495)",
+                            "inflammatory bowel disease_(28158)",
+                            {
+                                "relation": [
+                                    "gene/protein",
+                                    "associated with",
+                                    "disease",
+                                ],
+                                "label": ["gene/protein", "associated with", "disease"],
+                            },
+                        ),
+                        (
+                            "IFNG_(3495)",
+                            "Crohn's colitis_(83770)",
+                            {
+                                "relation": [
+                                    "gene/protein",
+                                    "associated with",
+                                    "disease",
+                                ],
+                                "label": ["gene/protein", "associated with", "disease"],
+                            },
+                        ),
+                        (
+                            "IFNG_(3495)",
+                            "Crohn ileitis and jejunitis_(35814)",
+                            {
+                                "relation": [
+                                    "gene/protein",
+                                    "associated with",
+                                    "disease",
+                                ],
+                                "label": ["gene/protein", "associated with", "disease"],
+                            },
+                        ),
+                        (
+                            "ATG16L1_(6661)",
+                            "IKBKG_(3672)",
+                            {
+                                "relation": ["gene/protein", "ppi", "gene/protein"],
+                                "label": ["gene/protein", "ppi", "gene/protein"],
+                            },
+                        ),
+                        (
+                            "Crohn's colitis_(83770)",
+                            "ATG16L1_(6661)",
+                            {
+                                "relation": [
+                                    "disease",
+                                    "associated with",
+                                    "gene/protein",
+                                ],
+                                "label": ["disease", "associated with", "gene/protein"],
+                            },
+                        ),
+                    ],
+                },
                 "graph_text": """
             node_id,node_attr
-            IFNG_(3495),"IFNG belongs to gene/protein category. 
-            This gene encodes a soluble cytokine that is a member of the type II interferon class. 
-            The encoded protein is secreted by cells of both the innate and adaptive immune systems. 
-            The active protein is a homodimer that binds to the interferon gamma receptor 
-            which triggers a cellular response to viral and microbial infections. 
-            Mutations in this gene are associated with an increased susceptibility to viral, 
-            bacterial and parasitic infections and to several autoimmune diseases. 
+            IFNG_(3495),"IFNG belongs to gene/protein category.
+            This gene encodes a soluble cytokine that is a member of the type II interferon class.
+            The encoded protein is secreted by cells of both the innate and adaptive immune systems.
+            The active protein is a homodimer that binds to the interferon gamma receptor
+            which triggers a cellular response to viral and microbial infections.
+            Mutations in this gene are associated with an increased susceptibility to viral,
+            bacterial and parasitic infections and to several autoimmune diseases.
             [provided by RefSeq, Dec 2015]."
-            IKBKG_(3672),"IKBKG belongs to gene/protein category. This gene encodes the regulatory 
-            subunit of the inhibitor of kappaB kinase (IKK) complex, which activates NF-kappaB 
-            resulting in activation of genes involved in inflammation, immunity, cell survival, 
-            and other pathways. Mutations in this gene result in incontinentia pigmenti, 
-            hypohidrotic ectodermal dysplasia, and several other types of immunodeficiencies. 
-            A pseudogene highly similar to this locus is located in an adjacent region of the 
+            IKBKG_(3672),"IKBKG belongs to gene/protein category. This gene encodes the regulatory
+            subunit of the inhibitor of kappaB kinase (IKK) complex, which activates NF-kappaB
+            resulting in activation of genes involved in inflammation, immunity, cell survival,
+            and other pathways. Mutations in this gene result in incontinentia pigmenti,
+            hypohidrotic ectodermal dysplasia, and several other types of immunodeficiencies.
+            A pseudogene highly similar to this locus is located in an adjacent region of the
             X chromosome. [provided by RefSeq, Mar 2016]."
-            ATG16L1_(6661),"ATG16L1 belongs to gene/protein category. The protein encoded 
-            by this gene is part of a large protein complex that is necessary for autophagy, 
-            the major process by which intracellular components are targeted to lysosomes 
-            for degradation. Defects in this gene are a cause of susceptibility to inflammatory 
-            bowel disease type 10 (IBD10). Several transcript variants encoding different 
+            ATG16L1_(6661),"ATG16L1 belongs to gene/protein category. The protein encoded
+            by this gene is part of a large protein complex that is necessary for autophagy,
+            the major process by which intracellular components are targeted to lysosomes
+            for degradation. Defects in this gene are a cause of susceptibility to inflammatory
+            bowel disease type 10 (IBD10). Several transcript variants encoding different
             isoforms have been found for this gene.[provided by RefSeq, Jun 2010]."
-            inflammatory bowel disease_(28158),inflammatory bowel disease belongs to disease 
-            category. Any inflammatory bowel disease in which the cause of the disease 
-            is a mutation in the NOD2 gene.  
-            Crohn ileitis and jejunitis_(35814),Crohn ileitis and jejunitis belongs to 
-            disease category. An Crohn disease involving a pathogenic inflammatory 
-            response in the ileum.  
-            Crohn's colitis_(83770),Crohn's colitis belongs to disease category. 
-            Crohn's disease affecting the colon.  
+            inflammatory bowel disease_(28158),inflammatory bowel disease belongs to disease
+            category. Any inflammatory bowel disease in which the cause of the disease
+            is a mutation in the NOD2 gene.
+            Crohn ileitis and jejunitis_(35814),Crohn ileitis and jejunitis belongs to
+            disease category. An Crohn disease involving a pathogenic inflammatory
+            response in the ileum.
+            Crohn's colitis_(83770),Crohn's colitis belongs to disease category.
+            Crohn's disease affecting the colon.
 
             head_id,edge_type,tail_id
             Crohn's colitis_(83770),"('disease', 'associated with', 'gene/protein')",
@@ -116,45 +163,45 @@ def input_dict_fixture():
             Crohn ileitis and jejunitis_(35814)
             """,
                 "graph_summary": """
-            The subgraph extracted from `subkg_12345` includes several important genes and 
+            The subgraph extracted from `subkg_12345` includes several important genes and
             their associations with inflammatory bowel diseases, particularly Crohn's disease.
 
             Key Nodes:
-            1. **IFNG (Interferon gamma)**: This gene encodes a cytokine that plays a crucial 
-            role in immune response. It is associated with several diseases, including 
-            inflammatory bowel disease and specifically Crohn's colitis and Crohn ileitis and 
-            jejunitis. Mutations in IFNG can lead to increased susceptibility to infections 
+            1. **IFNG (Interferon gamma)**: This gene encodes a cytokine that plays a crucial
+            role in immune response. It is associated with several diseases, including
+            inflammatory bowel disease and specifically Crohn's colitis and Crohn ileitis and
+            jejunitis. Mutations in IFNG can lead to increased susceptibility to infections
             and autoimmune diseases.
 
-            2. **IKBKG (Inhibitor of kappaB kinase gamma)**: This gene is involved in the 
-            regulation of NF-kappaB, which is critical for inflammation and immune responses. 
+            2. **IKBKG (Inhibitor of kappaB kinase gamma)**: This gene is involved in the
+            regulation of NF-kappaB, which is critical for inflammation and immune responses.
             Mutations can lead to immunodeficiencies and other disorders.
 
-            3. **ATG16L1**: This gene is essential for autophagy, a process that helps in 
-            degrading intracellular components. Defects in ATG16L1 are linked to inflammatory 
+            3. **ATG16L1**: This gene is essential for autophagy, a process that helps in
+            degrading intracellular components. Defects in ATG16L1 are linked to inflammatory
             bowel disease type 10 (IBD10) and are associated with Crohn's colitis.
 
-            4. **Inflammatory Bowel Disease**: A category of diseases characterized by 
-            chronic inflammation of the gastrointestinal tract, with specific mention of 
+            4. **Inflammatory Bowel Disease**: A category of diseases characterized by
+            chronic inflammation of the gastrointestinal tract, with specific mention of
             mutations in the NOD2 gene as a cause.
 
-            5. **Crohn's Colitis**: A specific type of Crohn's disease affecting the colon, 
+            5. **Crohn's Colitis**: A specific type of Crohn's disease affecting the colon,
             indicating a pathogenic inflammatory response.
 
-            6. **Crohn Ileitis and Jejunitis**: Another form of Crohn's disease that involves 
+            6. **Crohn Ileitis and Jejunitis**: Another form of Crohn's disease that involves
             inflammation in the ileum.
 
             Key Edges:
-            - **IFNG is associated with inflammatory bowel disease, Crohn's colitis, and 
+            - **IFNG is associated with inflammatory bowel disease, Crohn's colitis, and
             Crohn ileitis and jejunitis**: This highlights the role of IFNG in these diseases.
-            - **ATG16L1 is associated with Crohn's colitis**: This indicates a direct link 
+            - **ATG16L1 is associated with Crohn's colitis**: This indicates a direct link
             between the gene and the disease.
-            - **ATG16L1 interacts with IKBKG**: This protein-protein interaction suggests a 
-            functional relationship between these two genes in the context of immune response 
+            - **ATG16L1 interacts with IKBKG**: This protein-protein interaction suggests a
+            functional relationship between these two genes in the context of immune response
             and inflammation.
 
-            In summary, the subgraph illustrates the connections between key genes 
-            (IFNG, IKBKG, ATG16L1) and their associations with inflammatory bowel diseases, 
+            In summary, the subgraph illustrates the connections between key genes
+            (IFNG, IKBKG, ATG16L1) and their associations with inflammatory bowel diseases,
             particularly Crohn's disease, emphasizing the genetic underpinnings of these conditions.
             """,
             }
@@ -186,11 +233,11 @@ def test_graphrag_reasoning_openai(input_dict):
     )
     prompt = """
     Without extracting a new subgraph, based on subgraph extracted from `subkg_12345`
-    perform Graph RAG reasoning to get insights related to nodes of genes 
-    mentioned in the knowledge graph related to Adalimumab. 
+    perform Graph RAG reasoning to get insights related to nodes of genes
+    mentioned in the knowledge graph related to Adalimumab.
 
     Here is an additional context:
-    Adalimumab is a fully human monoclonal antibody (IgG1) 
+    Adalimumab is a fully human monoclonal antibody (IgG1)
     that specifically binds to tumor necrosis factor-alpha (TNF-α), a pro-inflammatory cytokine.
     """