aiagents4pharma 1.36.0__py3-none-any.whl → 1.37.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

aiagents4pharma/talk2knowledgegraphs/tools/multimodal_subgraph_extraction.py

@@ -0,0 +1,374 @@
+"""
+Tool for performing multimodal subgraph extraction.
+"""
+
+from typing import Type, Annotated
+import logging
+import pickle
+import numpy as np
+import pandas as pd
+import hydra
+import networkx as nx
+from pydantic import BaseModel, Field
+from langchain_core.tools import BaseTool
+from langchain_core.messages import ToolMessage
+from langchain_core.tools.base import InjectedToolCallId
+from langgraph.types import Command
+from langgraph.prebuilt import InjectedState
+import torch
+from torch_geometric.data import Data
+from ..utils.extractions.multimodal_pcst import MultimodalPCSTPruning
+from ..utils.embeddings.ollama import EmbeddingWithOllama
+from .load_arguments import ArgumentData
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class MultimodalSubgraphExtractionInput(BaseModel):
+    """
+    MultimodalSubgraphExtractionInput is a Pydantic model representing an input
+    for extracting a subgraph.
+
+    Args:
+        prompt: Prompt to interact with the backend.
+        tool_call_id: Tool call ID.
+        state: Injected state.
+        arg_data: Argument for analytical process over graph data.
+    """
+
+    tool_call_id: Annotated[str, InjectedToolCallId] = Field(
+        description="Tool call ID."
+    )
+    state: Annotated[dict, InjectedState] = Field(description="Injected state.")
+    prompt: str = Field(description="Prompt to interact with the backend.")
+    arg_data: ArgumentData = Field(
+        description="Experiment over graph data.", default=None
+    )
+
+
+class MultimodalSubgraphExtractionTool(BaseTool):
+    """
+    This tool performs subgraph extraction based on user's prompt by taking into account
+    the top-k nodes and edges.
+    """
+
+    name: str = "subgraph_extraction"
+    description: str = "A tool for subgraph extraction based on user's prompt."
+    args_schema: Type[BaseModel] = MultimodalSubgraphExtractionInput
+
+    def _prepare_query_modalities(self,
+                                  prompt_emb: list,
+                                  state: Annotated[dict, InjectedState],
+                                  pyg_graph: Data) -> pd.DataFrame:
+        """
+        Prepare the modality-specific query for subgraph extraction.
+
+        Args:
+            prompt_emb: The embedding of the user prompt in a list.
+            state: The injected state for the tool.
+            pyg_graph: The PyTorch Geometric graph Data.
+
+        Returns:
+            A DataFrame containing the query embeddings and modalities.
+        """
+        # Initialize dataframes
+        multimodal_df = pd.DataFrame({"name": []})
+        query_df = pd.DataFrame({"node_id": [],
+                                 "node_type": [],
+                                 "x": [],
+                                 "desc_x": [],
+                                 "use_description": []})
+
+        # Loop over the uploaded files and find multimodal files
+        for i in range(len(state["uploaded_files"])):
+            # Check if multimodal file is uploaded
+            if state["uploaded_files"][i]["file_type"] == "multimodal":
+                # Read the Excel file
+                multimodal_df = pd.read_excel(state["uploaded_files"][i]["file_path"],
+                                              sheet_name=None)
+
+        # Check if the multimodal_df is empty
+        if len(multimodal_df) > 0:
+            # Merge all obtained dataframes into a single dataframe
+            multimodal_df = pd.concat(multimodal_df).reset_index()
+            multimodal_df.drop(columns=["level_1"], inplace=True)
+            multimodal_df.rename(columns={"level_0": "q_node_type",
+                                        "name": "q_node_name"}, inplace=True)
+            # Since an excel sheet name could not contain a `/`,
+            # but the node type can be 'gene/protein' as exists in the PrimeKG
+            multimodal_df["q_node_type"] = multimodal_df.q_node_type.apply(
+                lambda x: x.replace('-', '/')
+            )
+
+            # Convert PyG graph to a DataFrame for easier filtering
+            graph_df = pd.DataFrame({
+                "node_id": pyg_graph.node_id,
+                "node_name": pyg_graph.node_name,
+                "node_type": pyg_graph.node_type,
+                "x": pyg_graph.x,
+                "desc_x": pyg_graph.desc_x.tolist(),
+            })
+
+            # Make a query dataframe by merging the graph_df and multimodal_df
+            query_df = graph_df.merge(multimodal_df, how='cross')
+            query_df = query_df[
+                query_df.apply(
+                    lambda x:
+                    (x['q_node_name'].lower() in x['node_name'].lower()) & # node name
+                    (x['node_type'] == x['q_node_type']), # node type
+                    axis=1
+                )
+            ]
+            query_df = query_df[['node_id', 'node_type', 'x', 'desc_x']].reset_index(drop=True)
+            query_df['use_description'] = False # set to False for modal-specific embeddings
+
+            # Update the state by adding the the selected node IDs
+            state["selections"] = query_df.groupby("node_type")["node_id"].apply(list).to_dict()
+
+        # Append a user prompt to the query dataframe
+        query_df = pd.concat([
+            query_df,
+            pd.DataFrame({
+                'node_id': 'user_prompt',
+                'node_type': 'prompt',
+                'x': prompt_emb,
+                'desc_x': prompt_emb,
+                'use_description': True # set to True for user prompt embedding
+            })
+        ]).reset_index(drop=True)
+
+        return query_df
+
+    def _perform_subgraph_extraction(self,
+                                     state: Annotated[dict, InjectedState],
+                                     cfg: dict,
+                                     pyg_graph: Data,
+                                     query_df: pd.DataFrame) -> dict:
+        """
+        Perform multimodal subgraph extraction based on modal-specific embeddings.
+
+        Args:
+            state: The injected state for the tool.
+            cfg: The configuration dictionary.
+            pyg_graph: The PyTorch Geometric graph Data.
+            query_df: The DataFrame containing the query embeddings and modalities.
+
+        Returns:
+            A dictionary containing the extracted subgraph with nodes and edges.
+        """
+        # Initialize the subgraph dictionary
+        subgraphs = {}
+        subgraphs["nodes"] = []
+        subgraphs["edges"] = []
+
+        # Loop over query embeddings and modalities
+        for q in query_df.iterrows():
+            # Prepare the PCSTPruning object and extract the subgraph
+            # Parameters were set in the configuration file obtained from Hydra
+            subgraph = MultimodalPCSTPruning(
+                topk=state["topk_nodes"],
+                topk_e=state["topk_edges"],
+                cost_e=cfg.cost_e,
+                c_const=cfg.c_const,
+                root=cfg.root,
+                num_clusters=cfg.num_clusters,
+                pruning=cfg.pruning,
+                verbosity_level=cfg.verbosity_level,
+                use_description=q[1]['use_description'],
+            ).extract_subgraph(pyg_graph,
+                               torch.tensor(q[1]['desc_x']), # description embedding
+                               torch.tensor(q[1]['x']), # modal-specific embedding
+                               q[1]['node_type'])
+
+            # Append the extracted subgraph to the dictionary
+            subgraphs["nodes"].append(subgraph["nodes"].tolist())
+            subgraphs["edges"].append(subgraph["edges"].tolist())
+
+        # Concatenate and get unique node and edge indices
+        subgraphs["nodes"] = np.unique(
+            np.concatenate([np.array(list_) for list_ in subgraphs["nodes"]])
+        )
+        subgraphs["edges"] = np.unique(
+            np.concatenate([np.array(list_) for list_ in subgraphs["edges"]])
+        )
+
+        return subgraphs
+
+    def _prepare_final_subgraph(self,
+                               state:Annotated[dict, InjectedState],
+                               subgraph: dict,
+                               graph: dict,
+                               cfg) -> dict:
+        """
+        Prepare the subgraph based on the extracted subgraph.
+
+        Args:
+            state: The injected state for the tool.
+            subgraph: The extracted subgraph.
+            graph: The initial graph containing PyG and textualized graph.
+            cfg: The configuration dictionary.
+
+        Returns:
+            A dictionary containing the PyG graph, NetworkX graph, and textualized graph.
+        """
+        # print(subgraph)
+        # Prepare the PyTorch Geometric graph
+        mapping = {n: i for i, n in enumerate(subgraph["nodes"].tolist())}
+        pyg_graph = Data(
+            # Node features
+            # x=pyg_graph.x[subgraph["nodes"]],
+            x=[graph["pyg"].x[i] for i in subgraph["nodes"]],
+            node_id=np.array(graph["pyg"].node_id)[subgraph["nodes"]].tolist(),
+            node_name=np.array(graph["pyg"].node_id)[subgraph["nodes"]].tolist(),
+            enriched_node=np.array(graph["pyg"].enriched_node)[subgraph["nodes"]].tolist(),
+            num_nodes=len(subgraph["nodes"]),
+            # Edge features
+            edge_index=torch.LongTensor(
+                [
+                    [
+                        mapping[i]
+                        for i in graph["pyg"].edge_index[:, subgraph["edges"]][0].tolist()
+                    ],
+                    [
+                        mapping[i]
+                        for i in graph["pyg"].edge_index[:, subgraph["edges"]][1].tolist()
+                    ],
+                ]
+            ),
+            edge_attr=graph["pyg"].edge_attr[subgraph["edges"]],
+            edge_type=np.array(graph["pyg"].edge_type)[subgraph["edges"]].tolist(),
+            relation=np.array(graph["pyg"].edge_type)[subgraph["edges"]].tolist(),
+            label=np.array(graph["pyg"].edge_type)[subgraph["edges"]].tolist(),
+            enriched_edge=np.array(graph["pyg"].enriched_edge)[subgraph["edges"]].tolist(),
+        )
+
+        # Networkx DiGraph construction to be visualized in the frontend
+        nx_graph = nx.DiGraph()
+        # Add nodes with attributes
+        node_colors = {n: cfg.node_colors_dict[k]
+                       for k, v in state["selections"].items() for n in v}
+        for n in pyg_graph.node_name:
+            nx_graph.add_node(n, color=node_colors.get(n, None))
+
+        # Add edges with attributes
+        edges = zip(
+            pyg_graph.edge_index[0].tolist(),
+            pyg_graph.edge_index[1].tolist(),
+            pyg_graph.edge_type
+        )
+        for src, dst, edge_type in edges:
+            nx_graph.add_edge(
+                pyg_graph.node_name[src],
+                pyg_graph.node_name[dst],
+                relation=edge_type,
+                label=edge_type,
+            )
+
+        # Prepare the textualized subgraph
+        textualized_graph = (
+            graph["text"]["nodes"].iloc[subgraph["nodes"]].to_csv(index=False)
+            + "\n"
+            + graph["text"]["edges"].iloc[subgraph["edges"]].to_csv(index=False)
+        )
+
+        return {
+            "graph_pyg": pyg_graph,
+            "graph_nx": nx_graph,
+            "graph_text": textualized_graph,
+        }
+
+    def _run(
+        self,
+        tool_call_id: Annotated[str, InjectedToolCallId],
+        state: Annotated[dict, InjectedState],
+        prompt: str,
+        arg_data: ArgumentData = None,
+    ) -> Command:
+        """
+        Run the subgraph extraction tool.
+
+        Args:
+            tool_call_id: The tool call ID for the tool.
+            state: Injected state for the tool.
+            prompt: The prompt to interact with the backend.
+            arg_data (ArgumentData): The argument data.
+
+        Returns:
+            Command: The command to be executed.
+        """
+        logger.log(logging.INFO, "Invoking subgraph_extraction tool")
+
+        # Load hydra configuration
+        with hydra.initialize(version_base=None, config_path="../configs"):
+            cfg = hydra.compose(
+                config_name="config", overrides=["tools/multimodal_subgraph_extraction=default"]
+            )
+            cfg = cfg.tools.multimodal_subgraph_extraction
+
+        # Retrieve source graph from the state
+        initial_graph = {}
+        initial_graph["source"] = state["dic_source_graph"][-1]  # The last source graph as of now
+        # logger.log(logging.INFO, "Source graph: %s", source_graph)
+
+        # Load the knowledge graph
+        with open(initial_graph["source"]["kg_pyg_path"], "rb") as f:
+            initial_graph["pyg"] = pickle.load(f)
+        with open(initial_graph["source"]["kg_text_path"], "rb") as f:
+            initial_graph["text"] = pickle.load(f)
+
+        # Prepare the query embeddings and modalities
+        query_df = self._prepare_query_modalities(
+            [EmbeddingWithOllama(model_name=cfg.ollama_embeddings[0]).embed_query(prompt)],
+            state,
+            initial_graph["pyg"]
+        )
+
+        # Perform subgraph extraction
+        subgraphs = self._perform_subgraph_extraction(state,
+                                                      cfg,
+                                                      initial_graph["pyg"],
+                                                      query_df)
+
+        # Prepare subgraph as a NetworkX graph and textualized graph
+        final_subgraph = self._prepare_final_subgraph(state,
+                                                      subgraphs,
+                                                      initial_graph,
+                                                      cfg)
+
+        # Prepare the dictionary of extracted graph
+        dic_extracted_graph = {
+            "name": arg_data.extraction_name,
+            "tool_call_id": tool_call_id,
+            "graph_source": initial_graph["source"]["name"],
+            "topk_nodes": state["topk_nodes"],
+            "topk_edges": state["topk_edges"],
+            "graph_dict": {
+                "nodes": list(final_subgraph["graph_nx"].nodes(data=True)),
+                "edges": list(final_subgraph["graph_nx"].edges(data=True)),
+            },
+            "graph_text": final_subgraph["graph_text"],
+            "graph_summary": None,
+        }
+
+        # Prepare the dictionary of updated state
+        dic_updated_state_for_model = {}
+        for key, value in {
+            "dic_extracted_graph": [dic_extracted_graph],
+        }.items():
+            if value:
+                dic_updated_state_for_model[key] = value
+
+        # Return the updated state of the tool
+        return Command(
+            update=dic_updated_state_for_model | {
+                # update the message history
+                "messages": [
+                    ToolMessage(
+                        content=f"Subgraph Extraction Result of {arg_data.extraction_name}",
+                        tool_call_id=tool_call_id,
+                    )
+                ],
+            }
+        )

aiagents4pharma/talk2knowledgegraphs/utils/extractions/__init__.py

@@ -2,3 +2,4 @@
 This file is used to import all the models in the package.
 '''
 from . import pcst
+from . import multimodal_pcst

aiagents4pharma/talk2knowledgegraphs/utils/extractions/multimodal_pcst.py

@@ -0,0 +1,292 @@
+"""
+Exctraction of multimodal subgraph using Prize-Collecting Steiner Tree (PCST) algorithm.
+"""
+
+from typing import Tuple, NamedTuple
+import numpy as np
+import pandas as pd
+import torch
+import pcst_fast
+from torch_geometric.data.data import Data
+
+class MultimodalPCSTPruning(NamedTuple):
+    """
+    Prize-Collecting Steiner Tree (PCST) pruning algorithm implementation inspired by G-Retriever
+    (He et al., 'G-Retriever: Retrieval-Augmented Generation for Textual Graph Understanding and
+    Question Answering', NeurIPS 2024) paper.
+    https://arxiv.org/abs/2402.07630
+    https://github.com/XiaoxinHe/G-Retriever/blob/main/src/dataset/utils/retrieval.py
+
+    Args:
+        topk: The number of top nodes to consider.
+        topk_e: The number of top edges to consider.
+        cost_e: The cost of the edges.
+        c_const: The constant value for the cost of the edges computation.
+        root: The root node of the subgraph, -1 for unrooted.
+        num_clusters: The number of clusters.
+        pruning: The pruning strategy to use.
+        verbosity_level: The verbosity level.
+    """
+    topk: int = 3
+    topk_e: int = 3
+    cost_e: float = 0.5
+    c_const: float = 0.01
+    root: int = -1
+    num_clusters: int = 1
+    pruning: str = "gw"
+    verbosity_level: int = 0
+    use_description: bool = False
+
+    def _compute_node_prizes(self,
+                             graph: Data,
+                             query_emb: torch.Tensor,
+                             modality: str) :
+        """
+        Compute the node prizes based on the cosine similarity between the query and nodes.
+
+        Args:
+            graph: The knowledge graph in PyTorch Geometric Data format.
+            query_emb: The query embedding in PyTorch Tensor format. This can be an embedding of
+                a prompt, sequence, or any other feature to be used for the subgraph extraction.
+            modality: The modality to use for the subgraph extraction based on the node type.
+
+        Returns:
+            The prizes of the nodes.
+        """
+        # Convert PyG graph to a DataFrame
+        graph_df = pd.DataFrame({
+            "node_type": graph.node_type,
+            "desc_x": [x.tolist() for x in graph.desc_x],
+            "x": [list(x) for x in graph.x],
+            "score": [0.0 for _ in range(len(graph.node_id))],
+        })
+
+        # Calculate cosine similarity for text features and update the score
+        if self.use_description:
+            graph_df.loc[:, "score"] = torch.nn.CosineSimilarity(dim=-1)(
+                    query_emb,
+                    torch.tensor(list(graph_df.desc_x.values)) # Using textual description features
+                ).tolist()
+        else:
+            graph_df.loc[graph_df["node_type"] == modality,
+                         "score"] = torch.nn.CosineSimilarity(dim=-1)(
+                    query_emb,
+                    torch.tensor(list(graph_df[graph_df["node_type"]== modality].x.values))
+                ).tolist()
+
+        # Set the prizes for nodes based on the similarity scores
+        n_prizes = torch.tensor(graph_df.score.values, dtype=torch.float32)
+        # n_prizes = torch.nn.CosineSimilarity(dim=-1)(query_emb, graph.x)
+        topk = min(self.topk, graph.num_nodes)
+        _, topk_n_indices = torch.topk(n_prizes, topk, largest=True)
+        n_prizes = torch.zeros_like(n_prizes)
+        n_prizes[topk_n_indices] = torch.arange(topk, 0, -1).float()
+
+        return n_prizes
+
+    def _compute_edge_prizes(self,
+                             graph: Data,
+                             text_emb: torch.Tensor) :
+        """
+        Compute the node prizes based on the cosine similarity between the query and nodes.
+
+        Args:
+            graph: The knowledge graph in PyTorch Geometric Data format.
+            text_emb: The textual description embedding in PyTorch Tensor format.
+
+        Returns:
+            The prizes of the nodes.
+        """
+        # Note that as of now, the edge features are based on textual features
+        # Compute prizes for edges
+        e_prizes = torch.nn.CosineSimilarity(dim=-1)(text_emb, graph.edge_attr)
+        unique_prizes, inverse_indices = e_prizes.unique(return_inverse=True)
+        topk_e = min(self.topk_e, unique_prizes.size(0))
+        topk_e_values, _ = torch.topk(unique_prizes, topk_e, largest=True)
+        e_prizes[e_prizes < topk_e_values[-1]] = 0.0
+        last_topk_e_value = topk_e
+        for k in range(topk_e):
+            indices = inverse_indices == (
+                unique_prizes == topk_e_values[k]
+            ).nonzero(as_tuple=True)[0]
+            value = min((topk_e - k) / indices.sum().item(), last_topk_e_value)
+            e_prizes[indices] = value
+            last_topk_e_value = value * (1 - self.c_const)
+
+        return e_prizes
+
+    def compute_prizes(self,
+                       graph: Data,
+                       text_emb: torch.Tensor,
+                       query_emb: torch.Tensor,
+                       modality: str):
+        """
+        Compute the node prizes based on the cosine similarity between the query and nodes,
+        as well as the edge prizes based on the cosine similarity between the query and edges.
+        Note that the node and edge embeddings shall use the same embedding model and dimensions
+        with the query.
+
+        Args:
+            graph: The knowledge graph in PyTorch Geometric Data format.
+            text_emb: The textual description embedding in PyTorch Tensor format.
+            query_emb: The query embedding in PyTorch Tensor format. This can be an embedding of
+                a prompt, sequence, or any other feature to be used for the subgraph extraction.
+            modality: The modality to use for the subgraph extraction based on node type.
+
+        Returns:
+            The prizes of the nodes and edges.
+        """
+        # Compute prizes for nodes
+        n_prizes = self._compute_node_prizes(graph, query_emb, modality)
+
+        # Compute prizes for edges
+        e_prizes = self._compute_edge_prizes(graph, text_emb)
+
+        return {"nodes": n_prizes, "edges": e_prizes}
+
+    def compute_subgraph_costs(self,
+                               graph: Data,
+                               prizes: dict) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """
+        Compute the costs in constructing the subgraph proposed by G-Retriever paper.
+
+        Args:
+            graph: The knowledge graph in PyTorch Geometric Data format.
+            prizes: The prizes of the nodes and the edges.
+
+        Returns:
+            edges: The edges of the subgraph, consisting of edges and number of edges without
+                virtual edges.
+            prizes: The prizes of the subgraph.
+            costs: The costs of the subgraph.
+        """
+        # Logic to reduce the cost of the edges such that at least one edge is selected
+        updated_cost_e = min(
+            self.cost_e,
+            prizes["edges"].max().item() * (1 - self.c_const / 2),
+        )
+
+        # Initialize variables
+        edges = []
+        costs = []
+        virtual = {
+            "n_prizes": [],
+            "edges": [],
+            "costs": [],
+        }
+        mapping = {"nodes": {}, "edges": {}}
+
+        # Compute the costs, edges, and virtual variables based on the prizes
+        for i, (src, dst) in enumerate(graph.edge_index.T.numpy()):
+            prize_e = prizes["edges"][i]
+            if prize_e <= updated_cost_e:
+                mapping["edges"][len(edges)] = i
+                edges.append((src, dst))
+                costs.append(updated_cost_e - prize_e)
+            else:
+                virtual_node_id = graph.num_nodes + len(virtual["n_prizes"])
+                mapping["nodes"][virtual_node_id] = i
+                virtual["edges"].append((src, virtual_node_id))
+                virtual["edges"].append((virtual_node_id, dst))
+                virtual["costs"].append(0)
+                virtual["costs"].append(0)
+                virtual["n_prizes"].append(prize_e - updated_cost_e)
+        prizes = np.concatenate([prizes["nodes"], np.array(virtual["n_prizes"])])
+        edges_dict = {}
+        edges_dict["edges"] = edges
+        edges_dict["num_prior_edges"] = len(edges)
+        # Final computation of the costs and edges based on the virtual costs and virtual edges
+        if len(virtual["costs"]) > 0:
+            costs = np.array(costs + virtual["costs"])
+            edges = np.array(edges + virtual["edges"])
+            edges_dict["edges"] = edges
+
+        return edges_dict, prizes, costs, mapping
+
+    def get_subgraph_nodes_edges(
+        self, graph: Data, vertices: np.ndarray, edges_dict: dict, mapping: dict,
+    ) -> dict:
+        """
+        Get the selected nodes and edges of the subgraph based on the vertices and edges computed
+        by the PCST algorithm.
+
+        Args:
+            graph: The knowledge graph in PyTorch Geometric Data format.
+            vertices: The vertices of the subgraph computed by the PCST algorithm.
+            edges_dict: The dictionary of edges of the subgraph computed by the PCST algorithm,
+                and the number of prior edges (without virtual edges).
+            mapping: The mapping dictionary of the nodes and edges.
+            num_prior_edges: The number of edges before adding virtual edges.
+
+        Returns:
+            The selected nodes and edges of the extracted subgraph.
+        """
+        # Get edges information
+        edges = edges_dict["edges"]
+        num_prior_edges = edges_dict["num_prior_edges"]
+        # Retrieve the selected nodes and edges based on the given vertices and edges
+        subgraph_nodes = vertices[vertices < graph.num_nodes]
+        subgraph_edges = [mapping["edges"][e] for e in edges if e < num_prior_edges]
+        virtual_vertices = vertices[vertices >= graph.num_nodes]
+        if len(virtual_vertices) > 0:
+            virtual_vertices = vertices[vertices >= graph.num_nodes]
+            virtual_edges = [mapping["nodes"][i] for i in virtual_vertices]
+            subgraph_edges = np.array(subgraph_edges + virtual_edges)
+        edge_index = graph.edge_index[:, subgraph_edges]
+        subgraph_nodes = np.unique(
+            np.concatenate(
+                [subgraph_nodes, edge_index[0].numpy(), edge_index[1].numpy()]
+            )
+        )
+
+        return {"nodes": subgraph_nodes, "edges": subgraph_edges}
+
+    def extract_subgraph(self,
+                         graph: Data,
+                         text_emb: torch.Tensor,
+                         query_emb: torch.Tensor,
+                         modality: str) -> dict:
+        """
+        Perform the Prize-Collecting Steiner Tree (PCST) algorithm to extract the subgraph.
+
+        Args:
+            graph: The knowledge graph in PyTorch Geometric Data format.
+            text_emb: The textual description embedding in PyTorch Tensor format.
+            query_emb: The query embedding in PyTorch Tensor format. This can be an embedding of
+                a prompt, sequence, or any other feature to be used for the subgraph extraction.
+            modality: The modality to use for the subgraph extraction
+                (e.g., "text", "sequence", "smiles").
+
+        Returns:
+            The selected nodes and edges of the subgraph.
+        """
+        # Assert the topk and topk_e values for subgraph retrieval
+        assert self.topk > 0, "topk must be greater than or equal to 0"
+        assert self.topk_e > 0, "topk_e must be greater than or equal to 0"
+
+        # Retrieve the top-k nodes and edges based on the query embedding
+        prizes = self.compute_prizes(graph, text_emb, query_emb, modality)
+
+        # Compute costs in constructing the subgraph
+        edges_dict, prizes, costs, mapping = self.compute_subgraph_costs(
+            graph, prizes
+        )
+
+        # Retrieve the subgraph using the PCST algorithm
+        result_vertices, result_edges = pcst_fast.pcst_fast(
+            edges_dict["edges"],
+            prizes,
+            costs,
+            self.root,
+            self.num_clusters,
+            self.pruning,
+            self.verbosity_level,
+        )
+
+        subgraph = self.get_subgraph_nodes_edges(
+            graph,
+            result_vertices,
+            {"edges": result_edges, "num_prior_edges": edges_dict["num_prior_edges"]},
+            mapping)
+
+        return subgraph

{aiagents4pharma-1.36.0.dist-info → aiagents4pharma-1.37.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aiagents4pharma
-Version: 1.36.0
+Version: 1.37.0
 Summary: AI Agents for drug discovery, drug development, and other pharmaceutical R&D.
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
@@ -62,6 +62,7 @@ Requires-Dist: umap-learn==0.5.7
 Requires-Dist: plotly-express==0.4.1
 Requires-Dist: seaborn==0.13.2
 Requires-Dist: scanpy==1.11.0
+Requires-Dist: openpyxl==3.1.5
 Dynamic: license-file
 
 [![Talk2BioModels](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2biomodels.yml/badge.svg)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2biomodels.yml)

{aiagents4pharma-1.36.0.dist-info → aiagents4pharma-1.37.0.dist-info}/RECORD

@@ -9,7 +9,7 @@ aiagents4pharma/talk2aiagents4pharma/configs/agents/main_agent/default.yaml,sha2
 aiagents4pharma/talk2aiagents4pharma/states/__init__.py,sha256=3wSvCpM29oqvVjhbhabm7FNm9Zt0rHO5tEn63YW6doc,108
 aiagents4pharma/talk2aiagents4pharma/states/state_talk2aiagents4pharma.py,sha256=NxujEBDKubvpV9UG7ERTDRB6psr0XnObCNHyztLAhgo,485
 aiagents4pharma/talk2aiagents4pharma/tests/__init__.py,sha256=Jbw5tJxSrjGoaK5IX3pJWDCNzhrVQ10lkYq2oQ_KQD8,45
-aiagents4pharma/talk2aiagents4pharma/tests/test_main_agent.py,sha256=ndHQT4ycWy-uKRAND7JmX_SuNg4g9hJw4UCW0CbKSp0,4165
+aiagents4pharma/talk2aiagents4pharma/tests/test_main_agent.py,sha256=_zUm8i8vrBbcDgpwExa1sVGr1A9FgZFuwoLS395RnhU,4418
 aiagents4pharma/talk2biomodels/__init__.py,sha256=1cq1HX2xoi_a0nDPuXYoSTrnL26OHQBW3zXNwwwjFO0,181
 aiagents4pharma/talk2biomodels/agents/__init__.py,sha256=sn5-fREjMdEvb-OUan3iOqrgYGjplNx3J8hYOaW0Po8,128
 aiagents4pharma/talk2biomodels/agents/t2b_agent.py,sha256=g0DIW5P-dtJoVyG4weFdDgTrJPL_Dx1MMbTWextJDZ4,3455
@@ -73,17 +73,18 @@ aiagents4pharma/talk2cells/tools/scp_agent/display_studies.py,sha256=6q59gh_NQai
 aiagents4pharma/talk2cells/tools/scp_agent/search_studies.py,sha256=MLe-twtFnOu-P8P9diYq7jvHBHbWFRRCZLcfpUzqPMg,2806
 aiagents4pharma/talk2knowledgegraphs/__init__.py,sha256=Z0Eo7LTiKk0STsr8VI7wkCLq7PHrK1vYlH4I1hSNLiA,165
 aiagents4pharma/talk2knowledgegraphs/agents/__init__.py,sha256=iOAzuy_8A03tQDFtSBhC9dldUo62z5gfxcVtXAdLOJs,92
-aiagents4pharma/talk2knowledgegraphs/agents/t2kg_agent.py,sha256=IcXSZ2qQA1m-gS-o0Pj_g1oar8uPdhsbaovloUFka3Q,3058
+aiagents4pharma/talk2knowledgegraphs/agents/t2kg_agent.py,sha256=w4wSSQ9gw_fzpcHZ2Bnqok17YDkFeQ3d72JenvQm6Oc,3089
 aiagents4pharma/talk2knowledgegraphs/configs/__init__.py,sha256=4_DVdpahaJ55yPl0aZotlFA_MYWLFF2cubWyKtBVI_Q,126
-aiagents4pharma/talk2knowledgegraphs/configs/config.yaml,sha256=X91262b-wkygiH4HrEr0bIzHxHDuDWwuxLQAmdUe-E4,367
+aiagents4pharma/talk2knowledgegraphs/configs/config.yaml,sha256=-AJXKnR2z5ig0SK_3vLL9JFjNRri7q7blHYFWxoTDl0,417
 aiagents4pharma/talk2knowledgegraphs/configs/agents/t2kg_agent/__init__.py,sha256=-fAORvyFmG2iSvFOFDixmt9OTQRR58y89uhhu2EgbA8,46
 aiagents4pharma/talk2knowledgegraphs/configs/agents/t2kg_agent/default.yaml,sha256=ENCGROwYFpR6g4QD518h73sshdn3vPVpotBMk1QJcpU,4830
 aiagents4pharma/talk2knowledgegraphs/configs/app/__init__.py,sha256=fKfc3FR7g5KjY9b6jzrU6cwKTVVpkoVZQS3dvUowu34,69
 aiagents4pharma/talk2knowledgegraphs/configs/app/frontend/__init__.py,sha256=-fAORvyFmG2iSvFOFDixmt9OTQRR58y89uhhu2EgbA8,46
-aiagents4pharma/talk2knowledgegraphs/configs/app/frontend/default.yaml,sha256=4azC4cH-_-zt-bRVgNjkFM24mjNke6Rgn9pNl7XWrPQ,912
+aiagents4pharma/talk2knowledgegraphs/configs/app/frontend/default.yaml,sha256=WJgd2ZU7_WQ1qlcTfkFlM8u23sH6eU2KgAm0E4kqqfs,941
 aiagents4pharma/talk2knowledgegraphs/configs/tools/__init__.py,sha256=C1yyRZW8hqWw46p_bh1vAJp2z9aVvn4HpKjKkjlWIqY,150
 aiagents4pharma/talk2knowledgegraphs/configs/tools/graphrag_reasoning/__init__.py,sha256=-fAORvyFmG2iSvFOFDixmt9OTQRR58y89uhhu2EgbA8,46
 aiagents4pharma/talk2knowledgegraphs/configs/tools/graphrag_reasoning/default.yaml,sha256=Ua99yECXiwp4ZCUDgsDskYbKzcJrv7roQuLj31Zky4c,1037
+aiagents4pharma/talk2knowledgegraphs/configs/tools/multimodal_subgraph_extraction/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aiagents4pharma/talk2knowledgegraphs/configs/tools/subgraph_extraction/__init__.py,sha256=-fAORvyFmG2iSvFOFDixmt9OTQRR58y89uhhu2EgbA8,46
 aiagents4pharma/talk2knowledgegraphs/configs/tools/subgraph_extraction/default.yaml,sha256=U8HvMsYbaOwDwQPATj7EFvLtTy7XZEplE5WMoNjgYYc,1469
 aiagents4pharma/talk2knowledgegraphs/configs/tools/subgraph_summarization/__init__.py,sha256=-fAORvyFmG2iSvFOFDixmt9OTQRR58y89uhhu2EgbA8,46
@@ -94,15 +95,16 @@ aiagents4pharma/talk2knowledgegraphs/datasets/dataset.py,sha256=-LaPLse8BkALqwFe
 aiagents4pharma/talk2knowledgegraphs/datasets/primekg.py,sha256=KBMhCJ7yjMWqQJJctFYdpjYAlwv48Jl6i1dddXP4f08,7599
 aiagents4pharma/talk2knowledgegraphs/datasets/starkqa_primekg.py,sha256=Y-6-nORsnBJlU6rH0skyfr9S9J4PfTWK-af_p5UuknQ,7483
 aiagents4pharma/talk2knowledgegraphs/states/__init__.py,sha256=XaqorSvx634dWRRlXUdzlisHtYMyqgJ2q7TanzsKlhw,108
-aiagents4pharma/talk2knowledgegraphs/states/state_talk2knowledgegraphs.py,sha256=6HqGo-awqoyNJG0igm5so5A4Tq8RPkCsjPg8Go38csE,1066
+aiagents4pharma/talk2knowledgegraphs/states/state_talk2knowledgegraphs.py,sha256=y5bp6yObN-AQtTq-m8ml7UnZaeKYUiPV_yjskAzBJaI,1087
 aiagents4pharma/talk2knowledgegraphs/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aiagents4pharma/talk2knowledgegraphs/tests/test_agents_t2kg_agent.py,sha256=CCN6cyhEaiXSvIC-4y3ueDSzjDCBYDsmSmOor-DMeF4,3928
+aiagents4pharma/talk2knowledgegraphs/tests/test_agents_t2kg_agent.py,sha256=PPfHKnfqMbUOBKU7q4VbQvHQymX1M_zTYdysQgVxKCs,3851
 aiagents4pharma/talk2knowledgegraphs/tests/test_datasets_biobridge_primekg.py,sha256=crH0eFA3P8P6IYzi1UWNa4YvRVrtlBzoScf9NaE1lDk,9827
 aiagents4pharma/talk2knowledgegraphs/tests/test_datasets_dataset.py,sha256=NFUlsZvhfIrkF4YenWfahrLK93Xhm5UYEGG_uYN2LVM,566
 aiagents4pharma/talk2knowledgegraphs/tests/test_datasets_primekg.py,sha256=Pvu0r93CpnhjkfMxc-EiVLpAJ04FdW9iTamCnetu654,2272
 aiagents4pharma/talk2knowledgegraphs/tests/test_datasets_starkqa_primekg.py,sha256=TuIsqcN1Mww3DTqGk6ebgJBWzUWdMWEq2yRQuYSFqvA,4416
 aiagents4pharma/talk2knowledgegraphs/tests/test_tools_graphrag_reasoning.py,sha256=aOKHTber2Cg3mjNjfIa6RZU7XdFj5C2ps1YEUXw76CI,10650
-aiagents4pharma/talk2knowledgegraphs/tests/test_tools_subgraph_extraction.py,sha256=zRi2j9Dm3VFywhhrPjVoJ7z_zJpAEM74MJRXapnhwVE,6246
+aiagents4pharma/talk2knowledgegraphs/tests/test_tools_multimodal_subgraph_extraction.py,sha256=Da-hXcu41_5Ge4DPlOoY6OqBwYnXPc58Q89wuywqVJM,5806
+aiagents4pharma/talk2knowledgegraphs/tests/test_tools_subgraph_extraction.py,sha256=C2HzmAG1XCeV1hwZzz3-9_2dm_84-i1BvTNWA1pqUwM,5393
 aiagents4pharma/talk2knowledgegraphs/tests/test_tools_subgraph_summarization.py,sha256=oBqfspXXOxH04OQuPb8BCW0liIQTGKXtaPNSrPpQtFc,7597
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_embeddings.py,sha256=uYFoE_6zeU10_1mLLAHUr5c4S2XZMSc0Q_860o-KWEw,1517
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_huggingface.py,sha256=hzX84pheZdEsTtikF2KtBFiH44_xPjYXxLA6p4Ax1CY,1623
@@ -117,9 +119,10 @@ aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_reactome.py,sh
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_uniprot.py,sha256=G13Diw7cA5TGINUNO1CDnN4rM6KbepxRXNjuzY578DI,1611
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_kg_utils.py,sha256=pal76wi7WgQWUNk56BrzfFV8jKpbDaHHdbwtgx_gXLI,2410
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_pubchem_utils.py,sha256=K1Y6QM0MDP1IrAdcWkigl8R-O-i-lsL4NCyOrWewhdM,1246
-aiagents4pharma/talk2knowledgegraphs/tools/__init__.py,sha256=zpD4h7EYtyq0QNOqLd6bkxrPlPb2XN64ceI9ncgESrA,171
+aiagents4pharma/talk2knowledgegraphs/tools/__init__.py,sha256=uleTEbhgvlYw4fOqV4NmoFvxGTon2Oim7jTQ5qPmYoU,216
 aiagents4pharma/talk2knowledgegraphs/tools/graphrag_reasoning.py,sha256=OEuOFncDRdb7TQEGq4rkT5On-jI-R7Nt8K5EBzaND8w,5338
 aiagents4pharma/talk2knowledgegraphs/tools/load_arguments.py,sha256=zhmsRp-8vjB5rRekqTA07d3yb-42HWqng9dDMkvK6hM,623
+aiagents4pharma/talk2knowledgegraphs/tools/multimodal_subgraph_extraction.py,sha256=Qjl8hXG8Gv5jQ4pBX8me0pGGakqRZmcDfTGgdEHD9pc,15394
 aiagents4pharma/talk2knowledgegraphs/tools/subgraph_extraction.py,sha256=te06QMFQfgJWrjaGrqpcOYeaV38jwm0KY_rXVSMHkeI,11468
 aiagents4pharma/talk2knowledgegraphs/tools/subgraph_summarization.py,sha256=mDSBOxopDfNhEJeU8fVI8b5lXTYrRzcc97aLbFgYSy4,4413
 aiagents4pharma/talk2knowledgegraphs/utils/__init__.py,sha256=cZqb3LZLmBnmyAtWFv2Z-4uJvQmx0M4zKsfiWrlM3Pk,195
@@ -138,7 +141,8 @@ aiagents4pharma/talk2knowledgegraphs/utils/enrichments/ols_terms.py,sha256=xSPP-
 aiagents4pharma/talk2knowledgegraphs/utils/enrichments/pubchem_strings.py,sha256=CQEGQ6Qsex2T91Vw7zTrclJBbSGGhxeWaVJb8tnURAQ,1691
 aiagents4pharma/talk2knowledgegraphs/utils/enrichments/reactome_pathways.py,sha256=I0cD0Fk2Uk27_4jEaIhpoGhoMh_RphY1VtkMnk4dkPg,2011
 aiagents4pharma/talk2knowledgegraphs/utils/enrichments/uniprot_proteins.py,sha256=z0Jb3tt8VzRjzqI9oVcUvRlPPg6BUdmslfKDIEFE_h8,3013
-aiagents4pharma/talk2knowledgegraphs/utils/extractions/__init__.py,sha256=7gwwtfzKhB8GuOBD47XRi0NprwEXkOzwNl5eeu-hDTI,86
+aiagents4pharma/talk2knowledgegraphs/utils/extractions/__init__.py,sha256=5bt3H6gGSAwN2K-IG7AHwG2lC4yQeMd2_jbhu2z5XKg,116
+aiagents4pharma/talk2knowledgegraphs/utils/extractions/multimodal_pcst.py,sha256=Irh5JXEhaLZ6Rxv3h5Anif_rGNItyLOGDWg1RACmoDA,12628
 aiagents4pharma/talk2knowledgegraphs/utils/extractions/pcst.py,sha256=m5p0yoJb7I19ua5yeQfXPf7c4r6S1XPwttsrM7Qoy94,9336
 aiagents4pharma/talk2scholars/__init__.py,sha256=NOZxTklAH1j1ggu97Ib8Xn9LCKudEWt-8dx8w7yxVD8,180
 aiagents4pharma/talk2scholars/agents/__init__.py,sha256=c_0Pk85bt-RfK5RMyALM3MXo3qXVMoYS7BOqM9wuFME,317
@@ -225,8 +229,8 @@ aiagents4pharma/talk2scholars/tools/zotero/utils/read_helper.py,sha256=lyrfpx8NH
 aiagents4pharma/talk2scholars/tools/zotero/utils/review_helper.py,sha256=IPD1V9yrBYaDnRe7sR6PrpwR82OBJbA2P_Tc6RbxAbM,2748
 aiagents4pharma/talk2scholars/tools/zotero/utils/write_helper.py,sha256=ALwLecy1QVebbsmXJiDj1GhGmyhq2R2tZlAyEl1vfhw,7410
 aiagents4pharma/talk2scholars/tools/zotero/utils/zotero_path.py,sha256=oIrfbOySgts50ksHKyjcWjRkPRIS88g3Lc0v9mBkU8w,6375
-aiagents4pharma-1.36.0.dist-info/licenses/LICENSE,sha256=IcIbyB1Hyk5ZDah03VNQvJkbNk2hkBCDqQ8qtnCvB4Q,1077
-aiagents4pharma-1.36.0.dist-info/METADATA,sha256=4S4eCTvL7mAxQUmSDp4SIyn2WAYuOivLpcCdL-j5dGQ,16757
-aiagents4pharma-1.36.0.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
-aiagents4pharma-1.36.0.dist-info/top_level.txt,sha256=-AH8rMmrSnJtq7HaAObS78UU-cTCwvX660dSxeM7a0A,16
-aiagents4pharma-1.36.0.dist-info/RECORD,,
+aiagents4pharma-1.37.0.dist-info/licenses/LICENSE,sha256=IcIbyB1Hyk5ZDah03VNQvJkbNk2hkBCDqQ8qtnCvB4Q,1077
+aiagents4pharma-1.37.0.dist-info/METADATA,sha256=F-uncJSmjQ9bOlTHKuLJMa311nbF90UL7aJXwn2zVe0,16788
+aiagents4pharma-1.37.0.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
+aiagents4pharma-1.37.0.dist-info/top_level.txt,sha256=-AH8rMmrSnJtq7HaAObS78UU-cTCwvX660dSxeM7a0A,16
+aiagents4pharma-1.37.0.dist-info/RECORD,,