aiagents4pharma 1.35.0__tar.gz → 1.37.0__tar.gz

{aiagents4pharma-1.35.0 → aiagents4pharma-1.37.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aiagents4pharma
-Version: 1.35.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)
@@ -90,6 +91,17 @@ Our toolkit currently consists of the following agents:
 
 ![AIAgents4Pharma](docs/assets/AIAgents4Pharma.png)
 
+## News
+
+- T2B and T2KG accepted at the MLGenX workshop during ICLR #2025 in Singapore. [Read More](https://openreview.net/forum?id=av4QhBNeZo)
+
+<div align="center">
+  <strong>Watch the presentation:</strong><br><br>
+  <a href="https://www.youtube.com/watch?v=3cU_OxY4HiE">
+    <img src="https://img.youtube.com/vi/3cU_OxY4HiE/0.jpg" alt="Watch the presentation" width="480">
+  </a>
+</div>
+
 ## Getting Started
 
 ### Installation
@@ -250,15 +262,19 @@ If you are using docker on Windows, please follow these [Windows Setup Notes](#n
 ![Python Version from PEP 621 TOML](https://img.shields.io/python/required-version-toml?tomlFilePath=https%3A%2F%2Fraw.githubusercontent.com%2FVirtualPatientEngine%2FAIAgents4Pharma%2Frefs%2Fheads%2Fmain%2Fpyproject.toml)
 
 1. **Clone the repository:**
-   ```sh
-   git clone https://github.com/VirtualPatientEngine/AIAgents4Pharma
-   cd AIAgents4Pharma
-   ```
+
+```sh
+git clone https://github.com/VirtualPatientEngine/AIAgents4Pharma
+cd AIAgents4Pharma
+```
+
 2. **Install dependencies:**
 
-   ```python
-   pip install -r requirements.txt
-   ```
+We use Conda as our environment manager, Follow the official [Quickstart](https://www.anaconda.com/docs/getting-started/miniconda/install#quickstart-install-instructions) install instructions provided by anaconda/miniconda.
+
+```python
+conda create --name AIAgents4Pharma python=3.12 -y && conda activate AIAgents4Pharma && pip install --upgrade pip && pip install -r requirements.txt
+```
 
 3. **Initialize API Keys**
 

{aiagents4pharma-1.35.0 → aiagents4pharma-1.37.0}/README.md

@@ -24,6 +24,17 @@ Our toolkit currently consists of the following agents:
 
 ![AIAgents4Pharma](docs/assets/AIAgents4Pharma.png)
 
+## News
+
+- T2B and T2KG accepted at the MLGenX workshop during ICLR #2025 in Singapore. [Read More](https://openreview.net/forum?id=av4QhBNeZo)
+
+<div align="center">
+  <strong>Watch the presentation:</strong><br><br>
+  <a href="https://www.youtube.com/watch?v=3cU_OxY4HiE">
+    <img src="https://img.youtube.com/vi/3cU_OxY4HiE/0.jpg" alt="Watch the presentation" width="480">
+  </a>
+</div>
+
 ## Getting Started
 
 ### Installation
@@ -184,15 +195,19 @@ If you are using docker on Windows, please follow these [Windows Setup Notes](#n
 ![Python Version from PEP 621 TOML](https://img.shields.io/python/required-version-toml?tomlFilePath=https%3A%2F%2Fraw.githubusercontent.com%2FVirtualPatientEngine%2FAIAgents4Pharma%2Frefs%2Fheads%2Fmain%2Fpyproject.toml)
 
 1. **Clone the repository:**
-   ```sh
-   git clone https://github.com/VirtualPatientEngine/AIAgents4Pharma
-   cd AIAgents4Pharma
-   ```
+
+```sh
+git clone https://github.com/VirtualPatientEngine/AIAgents4Pharma
+cd AIAgents4Pharma
+```
+
 2. **Install dependencies:**
 
-   ```python
-   pip install -r requirements.txt
-   ```
+We use Conda as our environment manager, Follow the official [Quickstart](https://www.anaconda.com/docs/getting-started/miniconda/install#quickstart-install-instructions) install instructions provided by anaconda/miniconda.
+
+```python
+conda create --name AIAgents4Pharma python=3.12 -y && conda activate AIAgents4Pharma && pip install --upgrade pip && pip install -r requirements.txt
+```
 
 3. **Initialize API Keys**
 

{aiagents4pharma-1.35.0 → aiagents4pharma-1.37.0}/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-1.35.0 → aiagents4pharma-1.37.0}/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-1.35.0 → aiagents4pharma-1.37.0}/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-1.35.0 → aiagents4pharma-1.37.0}/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-1.35.0 → aiagents4pharma-1.37.0}/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-1.35.0 → aiagents4pharma-1.37.0}/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-1.37.0/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-1.35.0 → aiagents4pharma-1.37.0}/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-1.35.0 → aiagents4pharma-1.37.0}/aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_sentencetransformer.py

@@ -2,17 +2,23 @@
 Test cases for utils/embeddings/sentence_transformer.py
 """
 
-import pytest
 import numpy as np
+import pytest
+
 from ..utils.embeddings.sentence_transformer import EmbeddingWithSentenceTransformer
 
+
 @pytest.fixture(name="embedding_model")
 def embedding_model_fixture():
     """
     Fixture for creating an instance of EmbeddingWithSentenceTransformer.
     """
     model_name = "sentence-transformers/all-MiniLM-L6-v1"  # Small model for testing
-    return EmbeddingWithSentenceTransformer(model_name=model_name)
+    embedding_model = EmbeddingWithSentenceTransformer(model_name=model_name)
+    # Move underlying model to CPU for testing
+    embedding_model.model.to("cpu")
+    return embedding_model
+
 
 def test_embed_documents(embedding_model):
     """
@@ -27,6 +33,7 @@ def test_embed_documents(embedding_model):
     assert len(embeddings[0]) == 384
     assert embeddings.dtype == np.float32
 
+
 def test_embed_query(embedding_model):
     """
     Test the embed_query method of EmbeddingWithSentenceTransformer class.