aiagents4pharma 1.18.0__py3-none-any.whl → 1.19.0__py3-none-any.whl

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_ollama.py

@@ -6,20 +6,21 @@ import pytest
 import ollama
 from ..utils.enrichments.ollama import EnrichmentWithOllama
 
+
 @pytest.fixture(name="ollama_config")
 def fixture_ollama_config():
     """Return a dictionary with Ollama configuration."""
     return {
-        "model_name": "smollm2:360m",
+        "model_name": "llama3.2:1b",
         "prompt_enrichment": """
-            Given the input as a list of strings, please return the list of addditional information of
-            each input terms using your prior knowledge.
+            Given the input as a list of strings, please return the list of addditional information 
+            of each input terms using your prior knowledge.
 
             Example:
             Input: ['acetaminophen', 'aspirin']
-            Ouput: ['acetaminophen is a medication used to treat pain and fever', 
+            Ouput: ['acetaminophen is a medication used to treat pain and fever',
             'aspirin is a medication used to treat pain, fever, and inflammation']
-            
+
             Do not include any pretext as the output, only the list of strings enriched.
 
             Input: {input}
@@ -28,10 +29,11 @@ def fixture_ollama_config():
         "streaming": False,
     }
 
+
 def test_no_model_ollama(ollama_config):
     """Test the case when the Ollama model is not available."""
     cfg = ollama_config
-    cfg_model = "smollm2:135m" # Choose a small model
+    cfg_model = "smollm2:135m"  # Choose a small model
 
     # Delete the Ollama model
     try:
@@ -41,7 +43,8 @@ def test_no_model_ollama(ollama_config):
 
     # Check if the model is available
     with pytest.raises(
-        ValueError, match=f"Error: Pulled {cfg_model} model and restarted Ollama server."
+        ValueError,
+        match=f"Error: Pulled {cfg_model} model and restarted Ollama server.",
     ):
         EnrichmentWithOllama(
             model_name=cfg_model,
@@ -51,7 +54,8 @@ def test_no_model_ollama(ollama_config):
         )
     ollama.delete(cfg_model)
 
-def test_enrich_nodes_ollama(ollama_config):
+
+def test_enrich_ollama(ollama_config):
     """Test the Ollama textual enrichment class for node enrichment."""
     # Prepare enrichment model
     cfg = ollama_config
@@ -63,37 +67,11 @@ def test_enrich_nodes_ollama(ollama_config):
     )
 
     # Perform enrichment for nodes
-    nodes = ["Adalimumab", "Infliximab"]
+    nodes = ["acetaminophen"]
     enriched_nodes = enr_model.enrich_documents(nodes)
     # Check the enriched nodes
-    assert len(enriched_nodes) == 2
-    assert all(
-        enriched_nodes[i] != nodes[i] for i in range(len(nodes))
-    )
-
-
-def test_enrich_relations_ollama(ollama_config):
-    """Test the Ollama textual enrichment class for relation enrichment."""
-    # Prepare enrichment model
-    cfg = ollama_config
-    enr_model = EnrichmentWithOllama(
-        model_name=cfg["model_name"],
-        prompt_enrichment=cfg["prompt_enrichment"],
-        temperature=cfg["temperature"],
-        streaming=cfg["streaming"],
-    )
-    # Perform enrichment for relations
-    relations = [
-        "IL23R-gene causation disease-inflammatory bowel diseases",
-        "NOD2-gene causation disease-inflammatory bowel diseases",
-    ]
-    enriched_relations = enr_model.enrich_documents(relations)
-    # Check the enriched relations
-    assert len(enriched_relations) == 2
-    assert all(
-        enriched_relations[i] != relations[i]
-        for i in range(len(relations))
-    )
+    assert len(enriched_nodes) == 1
+    assert all(enriched_nodes[i] != nodes[i] for i in range(len(nodes)))
 
 
 def test_enrich_ollama_rag(ollama_config):
@@ -107,11 +85,9 @@ def test_enrich_ollama_rag(ollama_config):
         streaming=cfg["streaming"],
     )
     # Perform enrichment for nodes
-    nodes = ["Adalimumab", "Infliximab"]
+    nodes = ["acetaminophen"]
     docs = [r"\path\to\doc1", r"\path\to\doc2"]
     enriched_nodes = enr_model.enrich_documents_with_rag(nodes, docs)
     # Check the enriched nodes
-    assert len(enriched_nodes) == 2
-    assert all(
-        enriched_nodes[i] != nodes[i] for i in range(len(nodes))
-    )
+    assert len(enriched_nodes) == 1
+    assert all(enriched_nodes[i] != nodes[i] for i in range(len(nodes)))

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_kg_utils.py

@@ -0,0 +1,79 @@
+"""
+Test cases for utils/kg_utils.py
+"""
+
+import pytest
+import networkx as nx
+import pandas as pd
+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/tools/__init__.py

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

aiagents4pharma/talk2knowledgegraphs/tools/graphrag_reasoning.py

@@ -0,0 +1,143 @@
+"""
+Tool for performing Graph RAG reasoning.
+"""
+
+import logging
+from typing import Type, Annotated
+from pydantic import BaseModel, Field
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.messages import ToolMessage
+from langchain_core.tools.base import InjectedToolCallId
+from langchain_core.tools import BaseTool
+from langchain_core.vectorstores import InMemoryVectorStore
+from langchain.chains.retrieval import create_retrieval_chain
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import PyPDFLoader
+from langgraph.types import Command
+from langgraph.prebuilt import InjectedState
+import hydra
+
+# 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.""",
+    ]

aiagents4pharma/talk2knowledgegraphs/tools/subgraph_extraction.py

@@ -0,0 +1,305 @@
+"""
+Tool for performing 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.chains.retrieval import create_retrieval_chain
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.vectorstores import InMemoryVectorStore
+from langchain_core.tools import BaseTool
+from langchain_core.messages import ToolMessage
+from langchain_core.tools.base import InjectedToolCallId
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langgraph.types import Command
+from langgraph.prebuilt import InjectedState
+import torch
+from torch_geometric.data import Data
+from ..utils.extractions.pcst import PCSTPruning
+from ..utils.embeddings.ollama import EmbeddingWithOllama
+from .load_arguments import ArgumentData
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class SubgraphExtractionInput(BaseModel):
+    """
+    SubgraphExtractionInput 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 SubgraphExtractionTool(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] = SubgraphExtractionInput
+
+    def perform_endotype_filtering(
+        self,
+        prompt: str,
+        state: Annotated[dict, InjectedState],
+        cfg: hydra.core.config_store.ConfigStore,
+    ) -> str:
+        """
+        Perform endotype filtering based on the uploaded files and prepare the prompt.
+
+        Args:
+            prompt: The prompt to interact with the backend.
+            state: Injected state for the tool.
+            cfg: Hydra configuration object.
+        """
+        # Loop through the uploaded files
+        all_genes = []
+        for uploaded_file in state["uploaded_files"]:
+            if uploaded_file["file_type"] == "endotype":
+                # Load the PDF file
+                docs = PyPDFLoader(file_path=uploaded_file["file_path"]).load()
+
+                # Split the text into chunks
+                splits = RecursiveCharacterTextSplitter(
+                    chunk_size=cfg.splitter_chunk_size,
+                    chunk_overlap=cfg.splitter_chunk_overlap,
+                ).split_documents(docs)
+
+                # Create a chat prompt template
+                prompt_template = ChatPromptTemplate.from_messages(
+                    [
+                        ("system", cfg.prompt_endotype_filtering),
+                        ("human", "{input}"),
+                    ]
+                )
+
+                qa_chain = create_stuff_documents_chain(
+                    state["llm_model"], prompt_template
+                )
+                rag_chain = create_retrieval_chain(
+                    InMemoryVectorStore.from_documents(
+                        documents=splits, 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,
+                )
+                results = rag_chain.invoke({"input": prompt})
+                all_genes.append(results["answer"])
+
+        # Prepare the prompt
+        if len(all_genes) > 0:
+            prompt = " ".join(
+                [prompt, cfg.prompt_endotype_addition, ", ".join(all_genes)]
+            )
+
+        return prompt
+
+    def prepare_final_subgraph(self,
+                               subgraph: dict,
+                               pyg_graph: Data,
+                               textualized_graph: pd.DataFrame) -> dict:
+        """
+        Prepare the subgraph based on the extracted subgraph.
+
+        Args:
+            subgraph: The extracted subgraph.
+            pyg_graph: The PyTorch Geometric graph.
+            textualized_graph: The textualized graph.
+
+        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"]],
+            node_id=np.array(pyg_graph.node_id)[subgraph["nodes"]].tolist(),
+            node_name=np.array(pyg_graph.node_id)[subgraph["nodes"]].tolist(),
+            enriched_node=np.array(pyg_graph.enriched_node)[subgraph["nodes"]].tolist(),
+            num_nodes=len(subgraph["nodes"]),
+            # Edge features
+            edge_index=torch.LongTensor(
+                [
+                    [
+                        mapping[i]
+                        for i in pyg_graph.edge_index[:, subgraph["edges"]][0].tolist()
+                    ],
+                    [
+                        mapping[i]
+                        for i in pyg_graph.edge_index[:, subgraph["edges"]][1].tolist()
+                    ],
+                ]
+            ),
+            edge_attr=pyg_graph.edge_attr[subgraph["edges"]],
+            edge_type=np.array(pyg_graph.edge_type)[subgraph["edges"]].tolist(),
+            relation=np.array(pyg_graph.edge_type)[subgraph["edges"]].tolist(),
+            label=np.array(pyg_graph.edge_type)[subgraph["edges"]].tolist(),
+            enriched_edge=np.array(pyg_graph.enriched_edge)[subgraph["edges"]].tolist(),
+        )
+
+        # Networkx DiGraph construction to be visualized in the frontend
+        nx_graph = nx.DiGraph()
+        for n in pyg_graph.node_name:
+            nx_graph.add_node(n)
+        for i, e in enumerate(
+            [
+                [pyg_graph.node_name[i], pyg_graph.node_name[j]]
+                for (i, j) in pyg_graph.edge_index.transpose(1, 0)
+            ]
+        ):
+            nx_graph.add_edge(
+                e[0],
+                e[1],
+                relation=pyg_graph.edge_type[i],
+                label=pyg_graph.edge_type[i],
+            )
+
+        # Prepare the textualized subgraph
+        textualized_graph = (
+            textualized_graph["nodes"].iloc[subgraph["nodes"]].to_csv(index=False)
+            + "\n"
+            + textualized_graph["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/subgraph_extraction=default"]
+            )
+            cfg = cfg.tools.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 prompt construction along with a list of endotypes
+        if len(state["uploaded_files"]) != 0 and "endotype" in [
+            f["file_type"] for f in state["uploaded_files"]
+        ]:
+            prompt = self.perform_endotype_filtering(prompt, state, cfg)
+
+        # Prepare embedding model and embed the user prompt as query
+        query_emb = torch.tensor(
+            EmbeddingWithOllama(model_name=cfg.ollama_embeddings[0]).embed_query(prompt)
+        ).float()
+
+        # Prepare the PCSTPruning object and extract the subgraph
+        # Parameters were set in the configuration file obtained from Hydra
+        subgraph = PCSTPruning(
+            state["topk_nodes"],
+            state["topk_edges"],
+            cfg.cost_e,
+            cfg.c_const,
+            cfg.root,
+            cfg.num_clusters,
+            cfg.pruning,
+            cfg.verbosity_level,
+        ).extract_subgraph(initial_graph["pyg"], query_emb)
+
+        # Prepare subgraph as a NetworkX graph and textualized graph
+        final_subgraph = self.prepare_final_subgraph(
+            subgraph, initial_graph["pyg"], initial_graph["text"]
+        )
+
+        # 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,
+                    )
+                ],
+            }
+        )