aiagents4pharma 1.36.0__py3-none-any.whl → 1.38.0__py3-none-any.whl

aiagents4pharma/talk2aiagents4pharma/tests/test_main_agent.py

@@ -20,12 +20,20 @@ def input_dict_fixture():
     input_dict = {
         "topk_nodes": 3,
         "topk_edges": 3,
+        "selections": {
+            "gene/protein": [],
+            "molecular_function": [],
+            "cellular_component": [],
+            "biological_process": [],
+            "drug": [],
+            "disease": []
+        },
         "uploaded_files": [],
         "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",
+                "name": "BioBridge",
+                "kg_pyg_path": f"{DATA_PATH}/biobridge_multimodal_pyg_graph.pkl",
+                "kg_text_path": f"{DATA_PATH}/biobridge_multimodal_text_graph.pkl",
             }
         ],
         "dic_extracted_graph": []
@@ -70,7 +78,7 @@ def test_main_agent_invokes_t2kg(input_dict):
     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["graph_source"] == "PrimeKG"
+    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)

aiagents4pharma/talk2knowledgegraphs/agents/t2kg_agent.py

@@ -9,7 +9,7 @@ from langchain_core.language_models.chat_models import BaseChatModel
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import START, StateGraph
 from langgraph.prebuilt import create_react_agent, ToolNode, InjectedState
-from ..tools.subgraph_extraction import SubgraphExtractionTool
+from ..tools.multimodal_subgraph_extraction import MultimodalSubgraphExtractionTool
 from ..tools.subgraph_summarization import SubgraphSummarizationTool
 from ..tools.graphrag_reasoning import GraphRAGReasoningTool
 from ..states.state_talk2knowledgegraphs import Talk2KnowledgeGraphs
@@ -39,7 +39,7 @@ def get_app(uniq_id, llm_model: BaseChatModel):
         cfg = cfg.agents.t2kg_agent
 
     # Define the tools
-    subgraph_extraction = SubgraphExtractionTool()
+    subgraph_extraction = MultimodalSubgraphExtractionTool()
     subgraph_summarization = SubgraphSummarizationTool()
     graphrag_reasoning = GraphRAGReasoningTool()
     tools = ToolNode([

aiagents4pharma/talk2knowledgegraphs/configs/app/frontend/default.yaml

@@ -2,12 +2,13 @@ _target_: app.frontend.streamlit_app_talk2knowledgegraphs
 default_user: "talk2kg_user"
 data_package_allowed_file_types:
   - "pdf"
-endotype_allowed_file_types:
-  - "pdf"
+multimodal_allowed_file_types:
+  - "xls"
+  - "xlsx"
 upload_data_dir: "../files"
 kg_name: "PrimeKG"
-kg_pyg_path: "aiagents4pharma/talk2knowledgegraphs/tests/files/primekg_ibd_pyg_graph.pkl"
-kg_text_path: "aiagents4pharma/talk2knowledgegraphs/tests/files/primekg_ibd_text_graph.pkl"
+kg_pyg_path: "aiagents4pharma/talk2knowledgegraphs/tests/files/biobridge_multimodal_pyg_graph.pkl"
+kg_text_path: "aiagents4pharma/talk2knowledgegraphs/tests/files/biobridge_multimodal_text_graph.pkl"
 openai_api_key: ${oc.env:OPENAI_API_KEY}
 openai_llms:
   - "gpt-4o-mini"
@@ -23,9 +24,9 @@ ollama_embeddings:
   - "nomic-embed-text"
 temperature: 0.1
 streaming: False
-reasoning_subgraph_topk_nodes: 10
+reasoning_subgraph_topk_nodes: 5
 reasoning_subgraph_topk_nodes_min: 1
 reasoning_subgraph_topk_nodes_max: 100
-reasoning_subgraph_topk_edges: 10
+reasoning_subgraph_topk_edges: 5
 reasoning_subgraph_topk_edges_min: 1
 reasoning_subgraph_topk_edges_max: 100

aiagents4pharma/talk2knowledgegraphs/configs/config.yaml

@@ -2,6 +2,7 @@ defaults:
   - _self_
   - agents/t2kg_agent: default
   - tools/subgraph_extraction: default
+  - tools/multimodal_subgraph_extraction: default
   - tools/subgraph_summarization: default
   - tools/graphrag_reasoning: default
   - utils/pubchem_utils: default

aiagents4pharma/talk2knowledgegraphs/states/state_talk2knowledgegraphs.py

@@ -31,6 +31,7 @@ class Talk2KnowledgeGraphs(AgentState):
 
     llm_model: BaseChatModel
     embedding_model: Embeddings
+    selections: dict
     uploaded_files: list
     topk_nodes: int
     topk_edges: int

aiagents4pharma/talk2knowledgegraphs/tests/test_agents_t2kg_agent.py

@@ -19,6 +19,14 @@ def input_dict_fixture():
     input_dict = {
         "llm_model": None,  # TBA for each test case
         "embedding_model": None,  # TBA for each test case
+        "selections": {
+            "gene/protein": [],
+            "molecular_function": [],
+            "cellular_component": [],
+            "biological_process": [],
+            "drug": [],
+            "disease": []
+        },
         "uploaded_files": [
             {
                 "file_name": "adalimumab.pdf",
@@ -27,21 +35,14 @@ def input_dict_fixture():
                 "uploaded_by": "VPEUser",
                 "uploaded_timestamp": "2024-11-05 00:00:00",
             },
-            {
-                "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",
-            },
         ],
         "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",
+                "name": "BioBridge",
+                "kg_pyg_path": f"{DATA_PATH}/biobridge_multimodal_pyg_graph.pkl",
+                "kg_text_path": f"{DATA_PATH}/biobridge_multimodal_text_graph.pkl",
             }
         ],
         "dic_extracted_graph": []
@@ -96,7 +97,7 @@ def test_t2kg_agent_openai(input_dict):
     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["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)

aiagents4pharma/talk2knowledgegraphs/tests/test_tools_multimodal_subgraph_extraction.py

@@ -0,0 +1,152 @@
+"""
+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

@@ -3,22 +3,21 @@ 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
+from ..tools.subgraph_extraction import SubgraphExtractionTool
 
 # Define the data path
 DATA_PATH = "aiagents4pharma/talk2knowledgegraphs/tests/files"
 
 
-@pytest.fixture(name="input_dict")
-def input_dict_fixture():
+@pytest.fixture(name="agent_state")
+def agent_state_fixture():
     """
-    Input dictionary fixture.
+    Agent state fixture.
     """
-    input_dict = {
-        "llm_model": None,  # TBA for each test case
-        "embedding_model": None,  # TBA for each test case
+    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,
@@ -31,52 +30,37 @@ def input_dict_fixture():
         ],
     }
 
-    return input_dict
+    return agent_state
 
 
-def test_extract_subgraph_wo_docs(input_dict):
+def test_extract_subgraph_wo_docs(agent_state):
     """
     Test the subgraph extraction tool without any documents using OpenAI model.
 
     Args:
-        input_dict: Input dictionary.
+        agent_state: Agent state in the form of a 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)
+    # Instantiate the SubgraphExtractionTool
+    subgraph_extraction_tool = SubgraphExtractionTool()
 
-    # Check assistant message
-    assistant_msg = response["messages"][-1].content
-    assert isinstance(assistant_msg, str)
+    # 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
-    tool_msg = response["messages"][-2]
-    assert tool_msg.name == "subgraph_extraction"
+    assert response.update["messages"][-1].tool_call_id  == "subgraph_extraction_tool"
 
     # Check extracted subgraph dictionary
-    current_state = app.get_state(config)
-    dic_extracted_graph = current_state.values["dic_extracted_graph"][0]
+    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"
@@ -99,24 +83,16 @@ def test_extract_subgraph_wo_docs(input_dict):
     )
 
 
-def test_extract_subgraph_w_docs(input_dict):
+def test_extract_subgraph_w_docs(agent_state):
     """
-    Test the subgraph extraction tool with a document as reference (i.e., endotype document)
-    using OpenAI model.
+    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:
-        input_dict: Input dictionary.
+        agent_state: Agent state in the form of a 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"] = [
+    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",
@@ -125,33 +101,28 @@ def test_extract_subgraph_w_docs(input_dict):
             "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:
 
+    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)
+    # Instantiate the SubgraphExtractionTool
+    subgraph_extraction_tool = SubgraphExtractionTool()
 
-    # Check assistant message
-    assistant_msg = response["messages"][-1].content
-    assert isinstance(assistant_msg, str)
+    # 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
-    tool_msg = response["messages"][-2]
-    assert tool_msg.name == "subgraph_extraction"
+    assert response.update["messages"][-1].tool_call_id  == "subgraph_extraction_tool"
 
     # Check extracted subgraph dictionary
-    current_state = app.get_state(config)
-    dic_extracted_graph = current_state.values["dic_extracted_graph"][0]
+    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"

aiagents4pharma/talk2knowledgegraphs/tools/__init__.py

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