aiagents4pharma 0.0.0__py3-none-any.whl

aiagents4pharma/talk2knowledgegraphs/tests/test_tools_multimodal_subgraph_extraction.py

@@ -0,0 +1,159 @@
+"""
+Test cases for tools/subgraph_extraction.py
+"""
+
+import pytest
+
+# from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from ..tools.multimodal_subgraph_extraction import MultimodalSubgraphExtractionTool
+
+# Define the data path
+DATA_PATH = "aiagents4pharma/talk2knowledgegraphs/tests/files"
+
+
+@pytest.fixture(name="agent_state")
+def agent_state_fixture():
+    """
+    Agent state fixture.
+    """
+    agent_state = {
+        # "llm_model": ChatOpenAI(model="gpt-4o-mini", temperature=0.0),
+        # "embedding_model": OpenAIEmbeddings(model="text-embedding-3-small"),
+        "selections": {
+            "gene/protein": [],
+            "molecular_function": [],
+            "cellular_component": [],
+            "biological_process": [],
+            "drug": [],
+            "disease": [],
+        },
+        "uploaded_files": [],
+        "topk_nodes": 3,
+        "topk_edges": 3,
+        "dic_source_graph": [
+            {
+                "name": "BioBridge",
+                "kg_pyg_path": f"{DATA_PATH}/biobridge_multimodal_pyg_graph.pkl",
+                "kg_text_path": f"{DATA_PATH}/biobridge_multimodal_text_graph.pkl",
+            }
+        ],
+    }
+
+    return agent_state
+
+
+def test_extract_multimodal_subgraph_wo_doc(agent_state):
+    """
+    Test the multimodal subgraph extraction tool for only text as modality.
+
+    Args:
+        agent_state: Agent state in the form of a dictionary.
+    """
+    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`.
+    """
+
+    # Instantiate the MultimodalSubgraphExtractionTool
+    subgraph_extraction_tool = MultimodalSubgraphExtractionTool()
+
+    # Invoking the subgraph_extraction_tool
+    response = subgraph_extraction_tool.invoke(
+        input={
+            "prompt": prompt,
+            "tool_call_id": "subgraph_extraction_tool",
+            "state": agent_state,
+            "arg_data": {"extraction_name": "subkg_12345"},
+        }
+    )
+
+    # Check tool message
+    assert response.update["messages"][-1].tool_call_id == "subgraph_extraction_tool"
+
+    # Check extracted subgraph dictionary
+    dic_extracted_graph = response.update["dic_extracted_graph"][0]
+    assert isinstance(dic_extracted_graph, dict)
+    assert dic_extracted_graph["name"] == "subkg_12345"
+    assert dic_extracted_graph["graph_source"] == "BioBridge"
+    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"].replace('"', "")
+        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"].replace('"', "")
+        for e in dic_extracted_graph["graph_dict"]["edges"]
+    )
+
+
+def test_extract_multimodal_subgraph_w_doc(agent_state):
+    """
+    Test the multimodal subgraph extraction tool for text as modality, plus genes.
+
+    Args:
+        agent_state: Agent state in the form of a dictionary.
+    """
+    # Update state
+    agent_state["uploaded_files"] = [
+        {
+            "file_name": "multimodal-analysis.xlsx",
+            "file_path": f"{DATA_PATH}/multimodal-analysis.xlsx",
+            "file_type": "multimodal",
+            "uploaded_by": "VPEUser",
+            "uploaded_timestamp": "2025-05-12 00:00:00",
+        }
+    ]
+
+    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`.
+    """
+
+    # Instantiate the SubgraphExtractionTool
+    subgraph_extraction_tool = MultimodalSubgraphExtractionTool()
+
+    # Invoking the subgraph_extraction_tool
+    response = subgraph_extraction_tool.invoke(
+        input={
+            "prompt": prompt,
+            "tool_call_id": "subgraph_extraction_tool",
+            "state": agent_state,
+            "arg_data": {"extraction_name": "subkg_12345"},
+        }
+    )
+
+    # Check tool message
+    assert response.update["messages"][-1].tool_call_id == "subgraph_extraction_tool"
+
+    # Check extracted subgraph dictionary
+    dic_extracted_graph = response.update["dic_extracted_graph"][0]
+    assert isinstance(dic_extracted_graph, dict)
+    assert dic_extracted_graph["name"] == "subkg_12345"
+    assert dic_extracted_graph["graph_source"] == "BioBridge"
+    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"].replace('"', "")
+        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"].replace('"', "")
+        for e in dic_extracted_graph["graph_dict"]["edges"]
+    )

aiagents4pharma/talk2knowledgegraphs/tests/test_tools_subgraph_extraction.py

@@ -0,0 +1,152 @@
+"""
+Test cases for tools/subgraph_extraction.py
+"""
+
+import pytest
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+
+from ..tools.subgraph_extraction import SubgraphExtractionTool
+
+# Define the data path
+DATA_PATH = "aiagents4pharma/talk2knowledgegraphs/tests/files"
+
+
+@pytest.fixture(name="agent_state")
+def agent_state_fixture():
+    """
+    Agent state fixture.
+    """
+    agent_state = {
+        "llm_model": ChatOpenAI(model="gpt-4o-mini", temperature=0.0),
+        "embedding_model": OpenAIEmbeddings(model="text-embedding-3-small"),
+        "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 agent_state
+
+
+def test_extract_subgraph_wo_docs(agent_state):
+    """
+    Test the subgraph extraction tool without any documents using OpenAI model.
+
+    Args:
+        agent_state: Agent state in the form of a dictionary.
+    """
+    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`.
+    """
+
+    # Instantiate the SubgraphExtractionTool
+    subgraph_extraction_tool = SubgraphExtractionTool()
+
+    # Invoking the subgraph_extraction_tool
+    response = subgraph_extraction_tool.invoke(
+        input={
+            "prompt": prompt,
+            "tool_call_id": "subgraph_extraction_tool",
+            "state": agent_state,
+            "arg_data": {"extraction_name": "subkg_12345"},
+        }
+    )
+
+    # Check tool message
+    assert response.update["messages"][-1].tool_call_id == "subgraph_extraction_tool"
+
+    # Check extracted subgraph dictionary
+    dic_extracted_graph = response.update["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(agent_state):
+    """
+    As a knowledge graph agent, I would like you to call a tool called `subgraph_extraction`.
+    After calling the tool, restrain yourself to call any other tool.
+
+    Args:
+        agent_state: Agent state in the form of a dictionary.
+    """
+    # Update state
+    agent_state["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",
+        }
+    ]
+
+    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`.
+    """
+
+    # Instantiate the SubgraphExtractionTool
+    subgraph_extraction_tool = SubgraphExtractionTool()
+
+    # Invoking the subgraph_extraction_tool
+    response = subgraph_extraction_tool.invoke(
+        input={
+            "prompt": prompt,
+            "tool_call_id": "subgraph_extraction_tool",
+            "state": agent_state,
+            "arg_data": {"extraction_name": "subkg_12345"},
+        }
+    )
+
+    # Check tool message
+    assert response.update["messages"][-1].tool_call_id == "subgraph_extraction_tool"
+
+    # Check extracted subgraph dictionary
+    dic_extracted_graph = response.update["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,201 @@
+"""
+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)