aiagents4pharma 0.0.0__py3-none-any.whl

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_kg_utils.py

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

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_pubchem_utils.py

@@ -0,0 +1,123 @@
+"""
+Test cases for utils/pubchem_utils.py
+
+These tests mock PubChem HTTP calls to avoid network dependency and verify
+parsing logic without modifying business logic.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+
+from ..utils import pubchem_utils
+
+
+def _resp(json_data):
+    """Create a minimal response-like object exposing .json()."""
+    return SimpleNamespace(json=lambda: json_data)
+
+
+@pytest.fixture(autouse=True)
+def mock_requests_get(monkeypatch):
+    """Mock requests.get used by pubchem_utils to return deterministic JSON."""
+
+    def _mock_get(url, _timeout=60, **_kwargs):
+        # CAS RN lookup: ethyl carbonate 105-58-8 -> CID 7766
+        if "substance/xref/RN/105-58-8/record/JSON" in url:
+            return _resp(
+                {"PC_Substances": [{"compound": [{"id": {"type": 1, "id": {"cid": 7766}}}]}]}
+            )
+
+        # DrugBank: DB00240 (Alclometasone) -> CID 5311000
+        if (
+            "substance/sourceid//drugbank/DB00240/JSON" in url
+            or "substance/sourceid/drugbank/DB00240/JSON" in url
+        ):
+            return _resp(
+                {"PC_Substances": [{"compound": [{"id": {"type": 1, "id": {"cid": 5311000}}}]}]}
+            )
+
+        # CTD: D002083 (Butylated Hydroxyanisole) -> CID 24667
+        if (
+            "substance/sourceid//comparative toxicogenomics database/D002083/JSON" in url
+            or "substance/sourceid/comparative toxicogenomics database/D002083/JSON" in url
+        ):
+            return _resp(
+                {"PC_Substances": [{"compound": [{"id": {"type": 1, "id": {"cid": 24667}}}]}]}
+            )
+
+        # CID description for 5311000
+        if (
+            "compound/cid//5311000/description/JSON" in url
+            or "compound/cid/5311000/description/JSON" in url
+        ):
+            return _resp(
+                {
+                    "InformationList": {
+                        "Information": [
+                            {
+                                "Description": (
+                                    "Alclometasone is a prednisolone compound having an "
+                                    "alpha-chloro substituent at the 9alpha-position."
+                                ),
+                            }
+                        ]
+                    }
+                }
+            )
+
+        # Default empty response
+        return _resp({})
+
+    monkeypatch.setattr(pubchem_utils.requests, "get", _mock_get)
+    return _mock_get
+
+
+def test_cas_rn2pubchem_cid():
+    """
+    Test the casRN2pubchem_cid function.
+
+    The CAS RN for ethyl carbonate is 105-58-8.
+    The PubChem CID for ethyl carbonate is 7766.
+    """
+    casrn = "105-58-8"
+    pubchem_cid = pubchem_utils.cas_rn2pubchem_cid(casrn)
+    assert pubchem_cid == 7766
+
+
+def test_external_id2pubchem_cid():
+    """
+    Test the external_id2pubchem_cid function.
+
+    The DrugBank ID for Alclometasone is DB00240 -> CID 5311000.
+    The CTD ID for Butylated Hydroxyanisole is D002083 -> CID 24667.
+    """
+    drugbank_id = "DB00240"
+    pubchem_cid = pubchem_utils.external_id2pubchem_cid("drugbank", drugbank_id)
+    assert pubchem_cid == 5311000
+
+    ctd_id = "D002083"
+    pubchem_cid = pubchem_utils.external_id2pubchem_cid(
+        "comparative toxicogenomics database", ctd_id
+    )
+    assert pubchem_cid == 24667
+
+
+def test_pubchem_cid_description():
+    """
+    Test the pubchem_cid_description function.
+
+    The PubChem CID for Alclometasone is 5311000.
+    The description starts with the expected prefix.
+    """
+    pubchem_cid = 5311000
+    description = pubchem_utils.pubchem_cid_description(pubchem_cid)
+    assert description.startswith(
+        "Alclometasone is a prednisolone compound having an alpha-chloro substituent"
+    )
+
+
+def test_mock_fallback_path_coverage():
+    """Exercise the default branch of the mocked requests.get for coverage."""
+    resp = pubchem_utils.requests.get("unknown://unmatched/url", timeout=1)
+    assert resp.json() == {}

aiagents4pharma/talk2knowledgegraphs/tools/__init__.py

@@ -0,0 +1,11 @@
+"""
+This file is used to import all the models in the package.
+"""
+
+from . import (
+    graphrag_reasoning,
+    milvus_multimodal_subgraph_extraction,
+    multimodal_subgraph_extraction,
+    subgraph_extraction,
+    subgraph_summarization,
+)

aiagents4pharma/talk2knowledgegraphs/tools/graphrag_reasoning.py

@@ -0,0 +1,138 @@
+"""
+Tool for performing Graph RAG reasoning.
+"""
+
+import logging
+from typing import Annotated
+
+import hydra
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain.chains.retrieval import create_retrieval_chain
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_core.messages import ToolMessage
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.tools import BaseTool
+from langchain_core.tools.base import InjectedToolCallId
+from langchain_core.vectorstores import InMemoryVectorStore
+from langgraph.prebuilt import InjectedState
+from langgraph.types import Command
+from pydantic import BaseModel, Field
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class GraphRAGReasoningInput(BaseModel):
+    """
+    GraphRAGReasoningInput is a Pydantic model representing an input for Graph RAG reasoning.
+
+    Args:
+        state: Injected state.
+        prompt: Prompt to interact with the backend.
+        extraction_name: Name assigned to the subgraph extraction process
+    """
+
+    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.")
+    extraction_name: str = Field(
+        description="""Name assigned to the subgraph extraction process
+                    when the subgraph_extraction tool is invoked."""
+    )
+
+
+class GraphRAGReasoningTool(BaseTool):
+    """
+    This tool performs reasoning using a Graph Retrieval-Augmented Generation (RAG) approach
+    over user's request by considering textualized subgraph context and document context.
+    """
+
+    name: str = "graphrag_reasoning"
+    description: str = """A tool to perform reasoning using a Graph RAG approach
+                        by considering textualized subgraph context and document context."""
+    args_schema: type[BaseModel] = GraphRAGReasoningInput
+
+    def _run(
+        self,
+        tool_call_id: Annotated[str, InjectedToolCallId],
+        state: Annotated[dict, InjectedState],
+        prompt: str,
+        extraction_name: str,
+    ):
+        """
+        Run the Graph RAG reasoning tool.
+
+        Args:
+            tool_call_id: The tool call ID.
+            state: The injected state.
+            prompt: The prompt to interact with the backend.
+            extraction_name: The name assigned to the subgraph extraction process.
+        """
+        logger.log(logging.INFO, "Invoking graphrag_reasoning tool for %s", extraction_name)
+
+        # Load Hydra configuration
+        with hydra.initialize(version_base=None, config_path="../configs"):
+            cfg = hydra.compose(
+                config_name="config", overrides=["tools/graphrag_reasoning=default"]
+            )
+            cfg = cfg.tools.graphrag_reasoning
+
+        # Prepare documents
+        all_docs = []
+        if len(state["uploaded_files"]) != 0:
+            for uploaded_file in state["uploaded_files"]:
+                if uploaded_file["file_type"] == "drug_data":
+                    # Load documents
+                    raw_documents = PyPDFLoader(file_path=uploaded_file["file_path"]).load()
+
+                    # Split documents
+                    # May need to find an optimal chunk size and overlap configuration
+                    documents = RecursiveCharacterTextSplitter(
+                        chunk_size=cfg.splitter_chunk_size,
+                        chunk_overlap=cfg.splitter_chunk_overlap,
+                    ).split_documents(raw_documents)
+
+                    # Add documents to the list
+                    all_docs.extend(documents)
+
+        # Load the extracted graph
+        extracted_graph = {dic["name"]: dic for dic in state["dic_extracted_graph"]}
+        # logger.log(logging.INFO, "Extracted graph: %s", extracted_graph)
+
+        # Set another prompt template
+        prompt_template = ChatPromptTemplate.from_messages(
+            [("system", cfg.prompt_graphrag_w_docs), ("human", "{input}")]
+        )
+
+        # Prepare chain with retrieved documents
+        qa_chain = create_stuff_documents_chain(state["llm_model"], prompt_template)
+        rag_chain = create_retrieval_chain(
+            InMemoryVectorStore.from_documents(
+                documents=all_docs, embedding=state["embedding_model"]
+            ).as_retriever(
+                search_type=cfg.retriever_search_type,
+                search_kwargs={
+                    "k": cfg.retriever_k,
+                    "fetch_k": cfg.retriever_fetch_k,
+                    "lambda_mult": cfg.retriever_lambda_mult,
+                },
+            ),
+            qa_chain,
+        )
+
+        # Invoke the chain
+        response = rag_chain.invoke(
+            {
+                "input": prompt,
+                "subgraph_summary": extracted_graph[extraction_name]["graph_summary"],
+            }
+        )
+
+        return Command(
+            update={
+                # update the message history
+                "messages": [ToolMessage(content=response, tool_call_id=tool_call_id)]
+            }
+        )

aiagents4pharma/talk2knowledgegraphs/tools/load_arguments.py

@@ -0,0 +1,22 @@
+"""
+A utility module for defining the dataclasses
+for the arguments to set up initial settings
+"""
+
+from dataclasses import dataclass
+from typing import Annotated
+
+
+@dataclass
+class ArgumentData:
+    """
+    Dataclass for storing the argument data.
+    """
+
+    extraction_name: Annotated[
+        str,
+        """An AI assigned _ separated name of the subgraph extraction
+                                    based on human query and the context of the graph reasoning
+                                    experiment.
+                                    This must be set before the subgraph extraction is invoked.""",
+    ]