aiagents4pharma 1.17.1__py3-none-any.whl → 1.19.0__py3-none-any.whl

aiagents4pharma/talk2knowledgegraphs/tests/test_tools_subgraph_extraction.py

@@ -0,0 +1,174 @@
+"""
+Test cases for tools/subgraph_extraction.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
+DATA_PATH = "aiagents4pharma/talk2knowledgegraphs/tests/files"
+
+
+@pytest.fixture(name="input_dict")
+def input_dict_fixture():
+    """
+    Input dictionary fixture.
+    """
+    input_dict = {
+        "llm_model": None,  # TBA for each test case
+        "embedding_model": None,  # TBA for each test case
+        "uploaded_files": [],
+        "topk_nodes": 3,
+        "topk_edges": 3,
+        "dic_source_graph": [
+            {
+                "name": "PrimeKG",
+                "kg_pyg_path": f"{DATA_PATH}/primekg_ibd_pyg_graph.pkl",
+                "kg_text_path": f"{DATA_PATH}/primekg_ibd_text_graph.pkl",
+            }
+        ],
+    }
+
+    return input_dict
+
+
+def test_extract_subgraph_wo_docs(input_dict):
+    """
+    Test the subgraph extraction tool without any documents using OpenAI model.
+
+    Args:
+        input_dict: Input dictionary.
+    """
+    # Prepare LLM and embedding model
+    input_dict["llm_model"] = ChatOpenAI(model="gpt-4o-mini", temperature=0.0)
+    input_dict["embedding_model"] = OpenAIEmbeddings(model="text-embedding-3-small")
+
+    # Setup the app
+    unique_id = 12345
+    app = get_app(unique_id, llm_model=input_dict["llm_model"])
+    config = {"configurable": {"thread_id": unique_id}}
+    # Update state
+    app.update_state(
+        config,
+        input_dict,
+    )
+    prompt = """
+    Please directly invoke `subgraph_extraction` tool without calling any other tools 
+    to respond to the following prompt:
+
+    Extract all relevant information related to nodes of genes related to inflammatory bowel disease 
+    (IBD) that existed in the knowledge graph.
+    Please set the extraction name for this process as `subkg_12345`.
+    """
+
+    # Test the tool subgraph_extraction
+    response = app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
+
+    # Check assistant message
+    assistant_msg = response["messages"][-1].content
+    assert isinstance(assistant_msg, str)
+
+    # Check tool message
+    tool_msg = response["messages"][-2]
+    assert tool_msg.name == "subgraph_extraction"
+
+    # Check extracted subgraph dictionary
+    current_state = app.get_state(config)
+    dic_extracted_graph = current_state.values["dic_extracted_graph"][0]
+    assert isinstance(dic_extracted_graph, dict)
+    assert dic_extracted_graph["name"] == "subkg_12345"
+    assert dic_extracted_graph["graph_source"] == "PrimeKG"
+    assert dic_extracted_graph["topk_nodes"] == 3
+    assert dic_extracted_graph["topk_edges"] == 3
+    assert isinstance(dic_extracted_graph["graph_dict"], dict)
+    assert len(dic_extracted_graph["graph_dict"]["nodes"]) > 0
+    assert len(dic_extracted_graph["graph_dict"]["edges"]) > 0
+    assert isinstance(dic_extracted_graph["graph_text"], str)
+    # Check if the nodes are in the graph_text
+    assert all(
+        n[0] in dic_extracted_graph["graph_text"]
+        for n in dic_extracted_graph["graph_dict"]["nodes"]
+    )
+    # Check if the edges are in the graph_text
+    assert all(
+        ",".join([e[0], '"' + str(tuple(e[2]["relation"])) + '"', e[1]])
+        in dic_extracted_graph["graph_text"]
+        for e in dic_extracted_graph["graph_dict"]["edges"]
+    )
+
+
+def test_extract_subgraph_w_docs(input_dict):
+    """
+    Test the subgraph extraction tool with a document as reference (i.e., endotype document)
+    using OpenAI model.
+
+    Args:
+        input_dict: Input dictionary.
+    """
+    # Prepare LLM and embedding model
+    input_dict["llm_model"] = ChatOpenAI(model="gpt-4o-mini", temperature=0.0)
+    input_dict["embedding_model"] = OpenAIEmbeddings(model="text-embedding-3-small")
+
+    # Setup the app
+    unique_id = 12345
+    app = get_app(unique_id, llm_model=input_dict["llm_model"])
+    config = {"configurable": {"thread_id": unique_id}}
+    # Update state
+    input_dict["uploaded_files"] = [
+        {
+            "file_name": "DGE_human_Colon_UC-vs-Colon_Control.pdf",
+            "file_path": f"{DATA_PATH}/DGE_human_Colon_UC-vs-Colon_Control.pdf",
+            "file_type": "endotype",
+            "uploaded_by": "VPEUser",
+            "uploaded_timestamp": "2024-11-05 00:00:00",
+        }
+    ]
+    app.update_state(
+        config,
+        input_dict,
+    )
+    prompt = """
+    Please ONLY invoke `subgraph_extraction` tool without calling any other tools 
+    to respond to the following prompt:
+
+    Extract all relevant information related to nodes of genes related to inflammatory bowel disease 
+    (IBD) that existed in the knowledge graph.
+    Please set the extraction name for this process as `subkg_12345`.
+    """
+
+    # Test the tool subgraph_extraction
+    response = app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
+
+    # Check assistant message
+    assistant_msg = response["messages"][-1].content
+    assert isinstance(assistant_msg, str)
+
+    # Check tool message
+    tool_msg = response["messages"][-2]
+    assert tool_msg.name == "subgraph_extraction"
+
+    # Check extracted subgraph dictionary
+    current_state = app.get_state(config)
+    dic_extracted_graph = current_state.values["dic_extracted_graph"][0]
+    assert isinstance(dic_extracted_graph, dict)
+    assert dic_extracted_graph["name"] == "subkg_12345"
+    assert dic_extracted_graph["graph_source"] == "PrimeKG"
+    assert dic_extracted_graph["topk_nodes"] == 3
+    assert dic_extracted_graph["topk_edges"] == 3
+    assert isinstance(dic_extracted_graph["graph_dict"], dict)
+    assert len(dic_extracted_graph["graph_dict"]["nodes"]) > 0
+    assert len(dic_extracted_graph["graph_dict"]["edges"]) > 0
+    assert isinstance(dic_extracted_graph["graph_text"], str)
+    # Check if the nodes are in the graph_text
+    assert all(
+        n[0] in dic_extracted_graph["graph_text"]
+        for n in dic_extracted_graph["graph_dict"]["nodes"]
+    )
+    # Check if the edges are in the graph_text
+    assert all(
+        ",".join([e[0], '"' + str(tuple(e[2]["relation"])) + '"', e[1]])
+        in dic_extracted_graph["graph_text"]
+        for e in dic_extracted_graph["graph_dict"]["edges"]
+    )

aiagents4pharma/talk2knowledgegraphs/tests/test_tools_subgraph_summarization.py

@@ -0,0 +1,154 @@
+"""
+Test cases for tools/subgraph_summarization.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
+DATA_PATH = "aiagents4pharma/talk2knowledgegraphs/tests/files"
+
+
+@pytest.fixture(name="input_dict")
+def input_dict_fixture():
+    """
+    Input dictionary fixture.
+    """
+    input_dict = {
+        "llm_model": None,  # TBA for each test case
+        "embedding_model": None,  # TBA for each test case
+        "uploaded_files": [],
+        "topk_nodes": 3,
+        "topk_edges": 3,
+        "dic_source_graph": [
+            {
+                "name": "PrimeKG",
+                "kg_pyg_path": f"{DATA_PATH}/primekg_ibd_pyg_graph.pkl",
+                "kg_text_path": f"{DATA_PATH}/primekg_ibd_text_graph.pkl",
+            }
+        ],
+        "dic_extracted_graph": [
+            {
+                "name": "subkg_12345",
+                "tool_call_id": "tool_12345",
+                "graph_source": "PrimeKG",
+                "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']})]},
+                "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. 
+            [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 
+            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 
+            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.  
+
+            head_id,edge_type,tail_id
+            Crohn's colitis_(83770),"('disease', 'associated with', 'gene/protein')",
+            ATG16L1_(6661)
+            ATG16L1_(6661),"('gene/protein', 'ppi', 'gene/protein')",IKBKG_(3672)
+            IFNG_(3495),"('gene/protein', 'associated with', 'disease')",
+            inflammatory bowel disease_(28158)
+            IFNG_(3495),"('gene/protein', 'associated with', 'disease')",Crohn's colitis_(83770)
+            IFNG_(3495),"('gene/protein', 'associated with', 'disease')",
+            Crohn ileitis and jejunitis_(35814)
+            """,
+                "graph_summary": None,
+            }
+        ],
+    }
+
+    return input_dict
+
+
+def test_summarize_subgraph(input_dict):
+    """
+    Test the subgraph summarization tool without any documents using Ollama model.
+
+    Args:
+        input_dict: Input dictionary fixture.
+    """
+    # Prepare LLM and embedding model
+    input_dict["llm_model"] = ChatOpenAI(model="gpt-4o-mini", temperature=0.0)
+    input_dict["embedding_model"] = OpenAIEmbeddings(model="text-embedding-3-small")
+
+    # Setup the app
+    unique_id = 12345
+    app = get_app(unique_id, llm_model=input_dict["llm_model"])
+    config = {"configurable": {"thread_id": unique_id}}
+    # Update state
+    app.update_state(
+        config,
+        input_dict,
+    )
+    prompt = """
+    Please directly invoke `subgraph_summarization` tool without calling any other tools 
+    to respond to the following prompt:
+
+    You are given a subgraph in the forms of textualized subgraph representing
+    nodes and edges (triples) obtained from extraction_name `subkg_12345`.
+    Summarize the given subgraph and higlight the importance nodes and edges.
+    """
+
+    # Test the tool subgraph_summarization
+    response = app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
+
+    # Check assistant message
+    assistant_msg = response["messages"][-1].content
+    assert isinstance(assistant_msg, str)
+
+    # Check tool message
+    tool_msg = response["messages"][-2]
+    assert tool_msg.name == "subgraph_summarization"
+
+    # Check summarized subgraph
+    current_state = app.get_state(config)
+    dic_extracted_graph = current_state.values["dic_extracted_graph"][0]
+    assert isinstance(dic_extracted_graph["graph_summary"], str)

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_huggingface.py

@@ -31,7 +31,6 @@ def test_embedding_with_huggingface_embed_query(embedding_model):
     # 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

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_ollama.py

@@ -0,0 +1,56 @@
+"""
+Test cases for utils/embeddings/ollama.py
+"""
+
+import pytest
+import ollama
+from ..utils.embeddings.ollama import EmbeddingWithOllama
+
+@pytest.fixture(name="ollama_config")
+def fixture_ollama_config():
+    """Return a dictionary with Ollama configuration."""
+    return {
+        "model_name": "all-minilm", # Choose a small model
+    }
+
+def test_no_model_ollama(ollama_config):
+    """Test the case when the Ollama model is not available."""
+    cfg = ollama_config
+
+    # Delete the Ollama model
+    try:
+        ollama.delete(cfg["model_name"])
+    except ollama.ResponseError:
+        pass
+
+    # Check if the model is available
+    with pytest.raises(
+        ValueError, match=f"Error: Pulled {cfg["model_name"]} model and restarted Ollama server."
+    ):
+        EmbeddingWithOllama(model_name=cfg["model_name"])
+
+@pytest.fixture(name="embedding_model")
+def embedding_model_fixture(ollama_config):
+    """Return the configuration object for the Ollama embedding model and model object"""
+    cfg = ollama_config
+    return EmbeddingWithOllama(model_name=cfg["model_name"])
+
+def test_embedding_with_ollama_embed_documents(embedding_model):
+    """Test embedding documents using the EmbeddingWithOllama class."""
+    # Perform embedding
+    texts = ["Adalimumab", "Infliximab", "Vedolizumab"]
+    result = embedding_model.embed_documents(texts)
+    # Check the result
+    assert len(result) == 3
+    assert len(result[0]) == 384
+
+def test_embedding_with_ollama_embed_query(embedding_model):
+    """Test embedding a query using the EmbeddingWithOllama class."""
+    # Perform embedding
+    text = "Adalimumab"
+    result = embedding_model.embed_query(text)
+    # Check the result
+    assert len(result) == 384
+
+    # Delete the Ollama model so that it will not be cached afterward
+    ollama.delete(embedding_model.model_name)

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_ollama.py

@@ -6,20 +6,21 @@ import pytest
 import ollama
 from ..utils.enrichments.ollama import EnrichmentWithOllama
 
+
 @pytest.fixture(name="ollama_config")
 def fixture_ollama_config():
     """Return a dictionary with Ollama configuration."""
     return {
-        "model_name": "smollm2:360m",
+        "model_name": "llama3.2:1b",
         "prompt_enrichment": """
-            Given the input as a list of strings, please return the list of addditional information of
-            each input terms using your prior knowledge.
+            Given the input as a list of strings, please return the list of addditional information 
+            of each input terms using your prior knowledge.
 
             Example:
             Input: ['acetaminophen', 'aspirin']
-            Ouput: ['acetaminophen is a medication used to treat pain and fever', 
+            Ouput: ['acetaminophen is a medication used to treat pain and fever',
             'aspirin is a medication used to treat pain, fever, and inflammation']
-            
+
             Do not include any pretext as the output, only the list of strings enriched.
 
             Input: {input}
@@ -28,10 +29,11 @@ def fixture_ollama_config():
         "streaming": False,
     }
 
+
 def test_no_model_ollama(ollama_config):
     """Test the case when the Ollama model is not available."""
     cfg = ollama_config
-    cfg_model = "smollm2:135m" # Choose a small model
+    cfg_model = "smollm2:135m"  # Choose a small model
 
     # Delete the Ollama model
     try:
@@ -41,7 +43,8 @@ def test_no_model_ollama(ollama_config):
 
     # Check if the model is available
     with pytest.raises(
-        ValueError, match=f"Error: Pulled {cfg_model} model and restarted Ollama server."
+        ValueError,
+        match=f"Error: Pulled {cfg_model} model and restarted Ollama server.",
     ):
         EnrichmentWithOllama(
             model_name=cfg_model,
@@ -51,7 +54,8 @@ def test_no_model_ollama(ollama_config):
         )
     ollama.delete(cfg_model)
 
-def test_enrich_nodes_ollama(ollama_config):
+
+def test_enrich_ollama(ollama_config):
     """Test the Ollama textual enrichment class for node enrichment."""
     # Prepare enrichment model
     cfg = ollama_config
@@ -63,37 +67,11 @@ def test_enrich_nodes_ollama(ollama_config):
     )
 
     # Perform enrichment for nodes
-    nodes = ["Adalimumab", "Infliximab"]
+    nodes = ["acetaminophen"]
     enriched_nodes = enr_model.enrich_documents(nodes)
     # Check the enriched nodes
-    assert len(enriched_nodes) == 2
-    assert all(
-        enriched_nodes[i] != nodes[i] for i in range(len(nodes))
-    )
-
-
-def test_enrich_relations_ollama(ollama_config):
-    """Test the Ollama textual enrichment class for relation enrichment."""
-    # Prepare enrichment model
-    cfg = ollama_config
-    enr_model = EnrichmentWithOllama(
-        model_name=cfg["model_name"],
-        prompt_enrichment=cfg["prompt_enrichment"],
-        temperature=cfg["temperature"],
-        streaming=cfg["streaming"],
-    )
-    # Perform enrichment for relations
-    relations = [
-        "IL23R-gene causation disease-inflammatory bowel diseases",
-        "NOD2-gene causation disease-inflammatory bowel diseases",
-    ]
-    enriched_relations = enr_model.enrich_documents(relations)
-    # Check the enriched relations
-    assert len(enriched_relations) == 2
-    assert all(
-        enriched_relations[i] != relations[i]
-        for i in range(len(relations))
-    )
+    assert len(enriched_nodes) == 1
+    assert all(enriched_nodes[i] != nodes[i] for i in range(len(nodes)))
 
 
 def test_enrich_ollama_rag(ollama_config):
@@ -107,11 +85,9 @@ def test_enrich_ollama_rag(ollama_config):
         streaming=cfg["streaming"],
     )
     # Perform enrichment for nodes
-    nodes = ["Adalimumab", "Infliximab"]
+    nodes = ["acetaminophen"]
     docs = [r"\path\to\doc1", r"\path\to\doc2"]
     enriched_nodes = enr_model.enrich_documents_with_rag(nodes, docs)
     # Check the enriched nodes
-    assert len(enriched_nodes) == 2
-    assert all(
-        enriched_nodes[i] != nodes[i] for i in range(len(nodes))
-    )
+    assert len(enriched_nodes) == 1
+    assert all(enriched_nodes[i] != nodes[i] for i in range(len(nodes)))

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_kg_utils.py

@@ -0,0 +1,79 @@
+"""
+Test cases for utils/kg_utils.py
+"""
+
+import pytest
+import networkx as nx
+import pandas as pd
+from ..utils import kg_utils
+
+
+@pytest.fixture(name="sample_graph")
+def make_sample_graph():
+    """Return a sample graph"""
+    sg = nx.Graph()
+    sg.add_node(1, node_id=1, feature_id="A", feature_value="ValueA")
+    sg.add_node(2, node_id=2, feature_id="B", feature_value="ValueB")
+    sg.add_edge(1, 2, edge_id=1, feature_id="E", feature_value="EdgeValue")
+    return sg
+
+
+def test_kg_to_df_pandas(sample_graph):
+    """Test the kg_to_df_pandas function"""
+    df_nodes, df_edges = kg_utils.kg_to_df_pandas(sample_graph)
+    print(df_nodes)
+    expected_nodes_data = {
+        "node_id": [1, 2],
+        "feature_id": ["A", "B"],
+        "feature_value": ["ValueA", "ValueB"],
+    }
+    expected_nodes_df = pd.DataFrame(expected_nodes_data, index=[1, 2])
+    print(expected_nodes_df)
+    expected_edges_data = {
+        "node_source": [1],
+        "node_target": [2],
+        "edge_id": [1],
+        "feature_id": ["E"],
+        "feature_value": ["EdgeValue"],
+    }
+    expected_edges_df = pd.DataFrame(expected_edges_data)
+
+    # Assert that the dataframes are equal but the order of columns may be different
+    # Ignore the index of the dataframes
+    pd.testing.assert_frame_equal(df_nodes, expected_nodes_df, check_like=True)
+    pd.testing.assert_frame_equal(df_edges, expected_edges_df, check_like=True)
+
+
+def test_df_pandas_to_kg():
+    """Test the df_pandas_to_kg function"""
+    nodes_data = {
+        "node_id": [1, 2],
+        "feature_id": ["A", "B"],
+        "feature_value": ["ValueA", "ValueB"],
+    }
+    df_nodes_attrs = pd.DataFrame(nodes_data).set_index("node_id")
+
+    edges_data = {
+        "node_source": [1],
+        "node_target": [2],
+        "edge_id": [1],
+        "feature_id": ["E"],
+        "feature_value": ["EdgeValue"],
+    }
+    df_edges = pd.DataFrame(edges_data)
+
+    kg = kg_utils.df_pandas_to_kg(
+        df_edges, df_nodes_attrs, "node_source", "node_target"
+    )
+
+    assert len(kg.nodes) == 2
+    assert len(kg.edges) == 1
+
+    assert kg.nodes[1]["feature_id"] == "A"
+    assert kg.nodes[1]["feature_value"] == "ValueA"
+    assert kg.nodes[2]["feature_id"] == "B"
+    assert kg.nodes[2]["feature_value"] == "ValueB"
+
+    assert kg.edges[1, 2]["feature_id"] == "E"
+    assert kg.edges[1, 2]["feature_value"] == "EdgeValue"
+    assert kg.edges[1, 2]["edge_id"] == 1

aiagents4pharma/talk2knowledgegraphs/tools/__init__.py

@@ -0,0 +1,6 @@
+'''
+This file is used to import all the models in the package.
+'''
+from . import subgraph_extraction
+from . import subgraph_summarization
+from . import graphrag_reasoning