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

aiagents4pharma/talk2knowledgegraphs/configs/agents/t2kg_agent/__init__.py

@@ -0,0 +1,3 @@
+'''
+Import all the modules in the package
+'''

aiagents4pharma/talk2knowledgegraphs/configs/agents/t2kg_agent/default.yaml

@@ -0,0 +1,62 @@
+_target_: agents.t2kg_agent.get_app
+state_modifier: >
+  You are talk2knowledgegraphs agent, a helpful assistant for reasoning over knowledge graphs.
+  User can ask questions related to the knowledge graphs, and you will provide the answers using 
+  the provided tools as follows (if necessary):
+  [`subgraph_extraction`, `subgraph_summarization`, `graphrag_reasoning`].
+
+  **Tools Descriptions**:
+  - `subgraph_extraction`: Extract a subgraph from the knowledge graph that contains the relevant
+    information to answer the user's query. This tool can be used to provide a subgraph context
+    as a part of the reasoning process. The extracted subgraph should contain the most relevant
+    nodes and edges to the user's query in the form of a textualized subgraph.
+  - `subgraph_summarization`: Summarize the extracted textualized subgraph obtained from the
+    `subgraph_extraction` tool. This tool can be used to provide a concise and informative summary
+    of the subgraph to be used for reasoning as subgraph context. This tool highlights the most
+    important nodes and edges in the subgraph to respond to the user's request.
+  - `graphrag_reasoning`: Reason over the extracted textualized subgraph to answer the user's
+    prompt by also considering the context from the extracted subgraph and the retrieved
+    documents. User may also have a set of uploaded files that can be used to provide additional
+    information for reasoning. The history of previous conversations should be considered as well, 
+    and you as an agent should provide which conversations can be included as chat history.
+
+  As an agent, you should approach each request by first understanding the user's query and then 
+  following the appropriate steps to provide the best answer possible. 
+
+  **Execution Steps**:
+  - Understand thoroughly the user's query and think over the best approach to answer it.
+  - You may not need to call any tool for each user's query. Use the related tool(s) as needed. 
+    Think deeply whether it is necessary to call any tool to respond to the user's request.
+  - Call `subgraph_extraction` if there is any indication that the user needs to get the
+    information from the knowledge graph, which is not directly available as context in the prompt or
+    in the previous extracted subgraph. 
+    If the user asks for subgraph extraction, suggest a value for the `extraction_name` argument.
+    You should always follow it with `subgraph_summarization` as the next tool to be invoked.
+  - Call `subgraph_summarization` tool to summarize the extracted subgraph and provide
+    a useful insights over the subgraph. This tool also has the ability to filter endotypes
+    in the forms of differentially expressed genes that are relevant to the input query. Make sure
+    to include the most relevant genes if the user provides endotype-related documents.
+    The summary of the subgraph will be stored as `graph_summary` in the state in which you can use
+    it for reasoning over the subgraph in the `graphrag_reasoning` tool afterwards.
+  - If the user asks follow-up questions related to the extracted subgraph, you should
+    call `subgraph_summarization` followed by `graphrag_reasoning` tools if you think
+    the answer can be retrieved from the previously extracted subgraph.
+  - Call `graphrag_reasoning` tool to reason over the extracted subgraph and documents.
+    Always perform reasoning over the extracted subgraph and documents to provide
+    the best possible answer to the user's query. Before calling this tool, 
+    make sure you have access to the summarized subgraph obtained from `subgraph_summarization` tool.
+  - By default, if the user asks for a specific question about the extracted graph, you should
+    call `subgraph_summarization` followed by `graphrag_reasoning` if the most recent subgraphs
+    contain the relevant information to answer the user's question.
+    Use the summarized subgraph as the subgraph context in the `graphrag_reasoning` tool.
+  - It is strongly recommended to avoid calling the same tool multiple times unless
+    it is necessary to get the correct and thorough answer to the user's request.
+
+  **Tool Calling Workflow Examples**:
+  - `subgraph_extraction` -> `subgraph_summarization` -> `graphrag_reasoning` when the user asks
+  for specific instructions to extract the subgraph and reason over it. Follow this order to
+  provide the most accurate and relevant information if you think the currently available context
+  is not enough to answer the user's question.
+  - `subgraph_summarization` -> `graphrag_reasoning` when the user asks for the previously extracted
+  subgraph. Use the summarized subgraph as the subgraph context in the `graphrag_reasoning` tool.
+  - Do not call `graphrag_reasoning` tool without calling `subgraph_summarization` tool first.

aiagents4pharma/talk2knowledgegraphs/configs/app/__init__.py

@@ -0,0 +1,4 @@
+'''
+Import all the modules in the package
+'''
+from . import frontend

aiagents4pharma/talk2knowledgegraphs/configs/app/frontend/__init__.py

@@ -0,0 +1,3 @@
+'''
+Import all the modules in the package
+'''

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

@@ -0,0 +1,31 @@
+_target_: app.frontend.streamlit_app_talk2knowledgegraphs
+default_user: "talk2kg_user"
+data_package_allowed_file_types:
+  - "pdf"
+endotype_allowed_file_types:
+  - "pdf"
+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"
+openai_api_key: ${oc.env:OPENAI_API_KEY}
+openai_llms:
+  - "gpt-4o-mini"
+  - "gpt-4-turbo"
+  - "gpt-3.5-turbo"
+openai_embeddings:
+  - "text-embedding-3-small"
+ollama_llms:
+  - "llama3.2:1b"
+  - "llama3.2"
+  - "llama3.1"
+ollama_embeddings:
+  - "nomic-embed-text"
+temperature: 0.1
+streaming: False
+reasoning_subgraph_topk_nodes: 10
+reasoning_subgraph_topk_nodes_min: 1
+reasoning_subgraph_topk_nodes_max: 100
+reasoning_subgraph_topk_edges: 10
+reasoning_subgraph_topk_edges_min: 1
+reasoning_subgraph_topk_edges_max: 100

aiagents4pharma/talk2knowledgegraphs/configs/config.yaml

@@ -0,0 +1,7 @@
+defaults:
+  - _self_
+  - agents/t2kg_agent: default
+  - tools/subgraph_extraction: default
+  - tools/subgraph_summarization: default
+  - tools/graphrag_reasoning: default
+  - app/frontend: default

aiagents4pharma/talk2knowledgegraphs/configs/tools/__init__.py

@@ -0,0 +1,6 @@
+'''
+Import all the modules in the package
+'''
+from . import subgraph_extraction
+from . import subgraph_summarization
+from . import graphrag_reasoning

aiagents4pharma/talk2knowledgegraphs/configs/tools/graphrag_reasoning/__init__.py

@@ -0,0 +1,3 @@
+'''
+Import all the modules in the package
+'''

aiagents4pharma/talk2knowledgegraphs/configs/tools/graphrag_reasoning/default.yaml

@@ -0,0 +1,24 @@
+_target_: talk2knowledgegraphs.tools.graphrag_reasoning
+splitter_chunk_size: 1024
+splitter_chunk_overlap: 256
+retriever_search_type: "mmr"
+retriever_k: 3
+retriever_fetch_k: 10
+retriever_lambda_mult: 0.3
+prompt_graphrag_w_docs_context: >
+  Given a chat history and the latest user question, which might reference context
+  in the chat history, formulate a standalone question that can be understood
+  without the chat history. Do NOT answer the question, just reformulate it if needed
+  and otherwise return it as is.
+
+  Question: {input}
+prompt_graphrag_w_docs: >
+  You are talk2knowledgegraphs, a helpful assistant performing retrievel-augmented generation (RAG)
+  over knowledge graphs.
+  One of your tasks is to answer react-based questions by using the following pieces of
+  retrieved context to answer the question. You can leverage a summarization of the subgraph
+  and the retrieved documents to provide the best possible answer to the user's query.
+
+  Subgraph Summary: {subgraph_summary}
+  Context: {context}
+  Question: {input}

aiagents4pharma/talk2knowledgegraphs/configs/tools/subgraph_extraction/__init__.py

@@ -0,0 +1,3 @@
+'''
+Import all the modules in the package
+'''

aiagents4pharma/talk2knowledgegraphs/configs/tools/subgraph_extraction/default.yaml

@@ -0,0 +1,43 @@
+_target_: talk2knowledgegraphs.tools.subgraph_extraction
+ollama_embeddings:
+  - "nomic-embed-text"
+temperature: 0.1
+streaming: False
+topk: 5
+topk_e: 5
+cost_e: 0.5
+c_const: 0.01
+root: -1
+num_clusters: 1
+pruning: "gw"
+verbosity_level: 0
+node_id_column: "node_id"
+node_attr_column: "node_attr"
+edge_src_column: "edge_src"
+edge_attr_column: "edge_attr"
+edge_dst_column: "edge_dst"
+prompt_endotype_filtering: >
+  You are talk2knowledgegraphs agent, a helpful assistant in filtering the most relevant endotype
+  for the subgraph extraction process.
+  Given the retrieved endotype documents, you need to filter the most relevant
+  endotype that will be used for the following reasoning process.
+  Only included a list of genes that exist in the provided documents
+  that are relevant to the input query.
+  For this task, you may modify your prompt to optimize the filtering process
+  based on factual informationbetween each gene in the documents and the input query.
+  Discover as many genes as possible that are relevant for enriching the subgraph extraction process.
+
+  You do not need to include any other information in the output.
+  Use the following output format:
+  [gene_1, gene_2, ..., gene_n]
+
+  {context}
+  Input: {input}
+prompt_endotype_addition: >
+   Include the following endotype for the subgraph extraction process:
+splitter_chunk_size: 64
+splitter_chunk_overlap: 16
+retriever_search_type: "mmr"
+retriever_k: 3
+retriever_fetch_k: 10
+retriever_lambda_mult: 0.3

aiagents4pharma/talk2knowledgegraphs/configs/tools/subgraph_summarization/__init__.py

@@ -0,0 +1,3 @@
+'''
+Import all the modules in the package
+'''

aiagents4pharma/talk2knowledgegraphs/configs/tools/subgraph_summarization/default.yaml

@@ -0,0 +1,9 @@
+_target_: talk2knowledgegraphs.tools.subgraph_summarization
+prompt_subgraph_summarization: >
+  You are talk2knowledgegraphs agent, a helpful assistant in reasoning over biomedical knowledge graph.
+  Your task is to summarize the extracted textualized subgraph to provide a concise and informative
+  summary of the subgraph to be used for reasoning as subgraph context. You are responsible for
+  highlighting the most important nodes and edges in the subgraph to respond to the user's question.
+
+  Textualized Subgraph: {textualized_subgraph}
+  Question: {input}

aiagents4pharma/talk2knowledgegraphs/states/__init__.py

@@ -0,0 +1,4 @@
+'''
+This file is used to import all the models in the package.
+'''
+from . import state_talk2knowledgegraphs

aiagents4pharma/talk2knowledgegraphs/states/state_talk2knowledgegraphs.py

@@ -0,0 +1,38 @@
+"""
+This is the state file for the Talk2KnowledgeGraphs agent.
+"""
+
+from typing import Annotated
+# import operator
+from langchain_core.embeddings.embeddings import Embeddings
+from langchain_core.language_models.chat_models import BaseChatModel
+from langgraph.prebuilt.chat_agent_executor import AgentState
+
+
+def add_data(data1: dict, data2: dict) -> dict:
+    """
+    A reducer function to merge two dictionaries.
+    """
+    left_idx_by_name = {data["name"]: idx for idx, data in enumerate(data1)}
+    merged = data1.copy()
+    for data in data2:
+        idx = left_idx_by_name.get(data["name"])
+        if idx is not None:
+            merged[idx] = data
+        else:
+            merged.append(data)
+    return merged
+
+
+class Talk2KnowledgeGraphs(AgentState):
+    """
+    The state for the Talk2KnowledgeGraphs agent.
+    """
+
+    llm_model: BaseChatModel
+    embedding_model: Embeddings
+    uploaded_files: list
+    topk_nodes: int
+    topk_edges: int
+    dic_source_graph: Annotated[list[dict], add_data]
+    dic_extracted_graph: Annotated[list[dict], add_data]

aiagents4pharma/talk2knowledgegraphs/tests/test_agents_t2kg_agent.py

@@ -0,0 +1,110 @@
+"""
+Test cases for agents/t2kg_agent.py
+"""
+
+import pytest
+from langchain_core.messages import HumanMessage
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from ..agents.t2kg_agent import get_app
+
+# Define the data path
+DATA_PATH = "aiagents4pharma/talk2knowledgegraphs/tests/files"
+
+
+@pytest.fixture(name="input_dict")
+def input_dict_fixture():
+    """
+    Input dictionary fixture.
+    """
+    input_dict = {
+        "llm_model": None,  # TBA for each test case
+        "embedding_model": None,  # TBA for each test case
+        "uploaded_files": [
+            {
+                "file_name": "adalimumab.pdf",
+                "file_path": f"{DATA_PATH}/adalimumab.pdf",
+                "file_type": "drug_data",
+                "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",
+            }
+        ],
+        "dic_extracted_graph": []
+    }
+
+    return input_dict
+
+
+def test_t2kg_agent_openai(input_dict):
+    """
+    Test the T2KG agent using OpenAI model.
+
+    Args:
+        input_dict: Input 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 = """
+    Adalimumab is a fully human monoclonal antibody (IgG1) 
+    that specifically binds to tumor necrosis factor-alpha (TNF-α), a pro-inflammatory cytokine.
+
+    I would like to get evidence from the knowledge graph about the mechanism of actions related to
+    Adalimumab in treating inflammatory bowel disease 
+    (IBD). Please follow these steps:
+    - Extract a subgraph from the PrimeKG that contains information about Adalimumab.
+    - Summarize the extracted subgraph.
+    - Reason about the mechanism of action of Adalimumab in treating IBD.
+
+    Please set the extraction name for the extraction process as `subkg_12345`.
+    """
+
+    # Test the tool get_modelinfo
+    response = app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
+
+    # Check assistant message
+    assistant_msg = response["messages"][-1].content
+    assert isinstance(assistant_msg, str)
+
+    # Check extracted subgraph dictionary
+    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["name"] == "subkg_12345"
+    assert dic_extracted_graph["graph_source"] == "PrimeKG"
+    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 summarized subgraph
+    assert isinstance(dic_extracted_graph["graph_summary"], str)
+    # Check reasoning results
+    assert "Adalimumab" in assistant_msg
+    assert "TNF" in assistant_msg

aiagents4pharma/talk2knowledgegraphs/tests/test_tools_graphrag_reasoning.py

@@ -0,0 +1,210 @@
+"""
+Test cases for tools/graphrag_reasoning.py
+"""
+
+import pytest
+from langchain_core.messages import HumanMessage
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from ..agents.t2kg_agent import get_app
+
+# Define the data path
+DATA_PATH = "aiagents4pharma/talk2knowledgegraphs/tests/files"
+
+
+@pytest.fixture(name="input_dict")
+def input_dict_fixture():
+    """
+    Input dictionary fixture.
+    """
+    input_dict = {
+        "llm_model": None,  # TBA for each test case
+        "embedding_model": None,  # TBA for each test case
+        "uploaded_files": [
+            {
+                "file_name": "adalimumab.pdf",
+                "file_path": f"{DATA_PATH}/adalimumab.pdf",
+                "file_type": "drug_data",
+                "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",
+            }
+        ],
+        "dic_extracted_graph": [
+            {
+                "name": "subkg_12345",
+                "tool_call_id": "tool_12345",
+                "graph_source": "PrimeKG",
+                "topk_nodes": 3,
+                "topk_edges": 3,
+                "graph_dict": {
+                    'nodes': [('IFNG_(3495)', {}), 
+                              ('IKBKG_(3672)', {}),
+                              ('ATG16L1_(6661)', {}),
+                              ('inflammatory bowel disease_(28158)', {}),
+                              ('Crohn ileitis and jejunitis_(35814)', {}),
+                              ("Crohn's colitis_(83770)", {})],
+                    'edges': [('IFNG_(3495)', 'inflammatory bowel disease_(28158)', 
+                               {'relation': ['gene/protein', 'associated with', 'disease'],
+                                'label': ['gene/protein', 'associated with', 'disease']}), 
+                              ('IFNG_(3495)', "Crohn's colitis_(83770)",
+                               {'relation': ['gene/protein', 'associated with', 'disease'],
+                                'label': ['gene/protein', 'associated with', 'disease']}),
+                              ('IFNG_(3495)', 'Crohn ileitis and jejunitis_(35814)',
+                               {'relation': ['gene/protein', 'associated with', 'disease'],
+                                'label': ['gene/protein', 'associated with', 'disease']}), 
+                              ('ATG16L1_(6661)', 'IKBKG_(3672)',
+                               {'relation': ['gene/protein', 'ppi', 'gene/protein'],
+                                'label': ['gene/protein', 'ppi', 'gene/protein']}), 
+                              ("Crohn's colitis_(83770)", 'ATG16L1_(6661)',
+                               {'relation': ['disease', 'associated with', 'gene/protein'],
+                                'label': ['disease', 'associated with', 'gene/protein']})]},
+                "graph_text": """
+            node_id,node_attr
+            IFNG_(3495),"IFNG belongs to gene/protein category. 
+            This gene encodes a soluble cytokine that is a member of the type II interferon class. 
+            The encoded protein is secreted by cells of both the innate and adaptive immune systems. 
+            The active protein is a homodimer that binds to the interferon gamma receptor 
+            which triggers a cellular response to viral and microbial infections. 
+            Mutations in this gene are associated with an increased susceptibility to viral, 
+            bacterial and parasitic infections and to several autoimmune diseases. 
+            [provided by RefSeq, Dec 2015]."
+            IKBKG_(3672),"IKBKG belongs to gene/protein category. This gene encodes the regulatory 
+            subunit of the inhibitor of kappaB kinase (IKK) complex, which activates NF-kappaB 
+            resulting in activation of genes involved in inflammation, immunity, cell survival, 
+            and other pathways. Mutations in this gene result in incontinentia pigmenti, 
+            hypohidrotic ectodermal dysplasia, and several other types of immunodeficiencies. 
+            A pseudogene highly similar to this locus is located in an adjacent region of the 
+            X chromosome. [provided by RefSeq, Mar 2016]."
+            ATG16L1_(6661),"ATG16L1 belongs to gene/protein category. The protein encoded 
+            by this gene is part of a large protein complex that is necessary for autophagy, 
+            the major process by which intracellular components are targeted to lysosomes 
+            for degradation. Defects in this gene are a cause of susceptibility to inflammatory 
+            bowel disease type 10 (IBD10). Several transcript variants encoding different 
+            isoforms have been found for this gene.[provided by RefSeq, Jun 2010]."
+            inflammatory bowel disease_(28158),inflammatory bowel disease belongs to disease 
+            category. Any inflammatory bowel disease in which the cause of the disease 
+            is a mutation in the NOD2 gene.  
+            Crohn ileitis and jejunitis_(35814),Crohn ileitis and jejunitis belongs to 
+            disease category. An Crohn disease involving a pathogenic inflammatory 
+            response in the ileum.  
+            Crohn's colitis_(83770),Crohn's colitis belongs to disease category. 
+            Crohn's disease affecting the colon.  
+
+            head_id,edge_type,tail_id
+            Crohn's colitis_(83770),"('disease', 'associated with', 'gene/protein')",
+            ATG16L1_(6661)
+            ATG16L1_(6661),"('gene/protein', 'ppi', 'gene/protein')",IKBKG_(3672)
+            IFNG_(3495),"('gene/protein', 'associated with', 'disease')",
+            inflammatory bowel disease_(28158)
+            IFNG_(3495),"('gene/protein', 'associated with', 'disease')",Crohn's colitis_(83770)
+            IFNG_(3495),"('gene/protein', 'associated with', 'disease')",
+            Crohn ileitis and jejunitis_(35814)
+            """,
+                "graph_summary": """
+            The subgraph extracted from `subkg_12345` includes several important genes and 
+            their associations with inflammatory bowel diseases, particularly Crohn's disease.
+
+            Key Nodes:
+            1. **IFNG (Interferon gamma)**: This gene encodes a cytokine that plays a crucial 
+            role in immune response. It is associated with several diseases, including 
+            inflammatory bowel disease and specifically Crohn's colitis and Crohn ileitis and 
+            jejunitis. Mutations in IFNG can lead to increased susceptibility to infections 
+            and autoimmune diseases.
+
+            2. **IKBKG (Inhibitor of kappaB kinase gamma)**: This gene is involved in the 
+            regulation of NF-kappaB, which is critical for inflammation and immune responses. 
+            Mutations can lead to immunodeficiencies and other disorders.
+
+            3. **ATG16L1**: This gene is essential for autophagy, a process that helps in 
+            degrading intracellular components. Defects in ATG16L1 are linked to inflammatory 
+            bowel disease type 10 (IBD10) and are associated with Crohn's colitis.
+
+            4. **Inflammatory Bowel Disease**: A category of diseases characterized by 
+            chronic inflammation of the gastrointestinal tract, with specific mention of 
+            mutations in the NOD2 gene as a cause.
+
+            5. **Crohn's Colitis**: A specific type of Crohn's disease affecting the colon, 
+            indicating a pathogenic inflammatory response.
+
+            6. **Crohn Ileitis and Jejunitis**: Another form of Crohn's disease that involves 
+            inflammation in the ileum.
+
+            Key Edges:
+            - **IFNG is associated with inflammatory bowel disease, Crohn's colitis, and 
+            Crohn ileitis and jejunitis**: This highlights the role of IFNG in these diseases.
+            - **ATG16L1 is associated with Crohn's colitis**: This indicates a direct link 
+            between the gene and the disease.
+            - **ATG16L1 interacts with IKBKG**: This protein-protein interaction suggests a 
+            functional relationship between these two genes in the context of immune response 
+            and inflammation.
+
+            In summary, the subgraph illustrates the connections between key genes 
+            (IFNG, IKBKG, ATG16L1) and their associations with inflammatory bowel diseases, 
+            particularly Crohn's disease, emphasizing the genetic underpinnings of these conditions.
+            """,
+            }
+        ],
+    }
+
+    return input_dict
+
+
+def test_graphrag_reasoning_openai(input_dict):
+    """
+    Test the GraphRAG reasoning tool using OpenAI model.
+
+    Args:
+        input_dict: Input 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 = """
+    Without extracting a new subgraph, based on subgraph extracted from `subkg_12345`
+    perform Graph RAG reasoning to get insights related to nodes of genes 
+    mentioned in the knowledge graph related to Adalimumab. 
+
+    Here is an additional context:
+    Adalimumab is a fully human monoclonal antibody (IgG1) 
+    that specifically binds to tumor necrosis factor-alpha (TNF-α), a pro-inflammatory cytokine.
+    """
+
+    # Test the tool  graphrag_reasoning
+    response = app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
+
+    # Check assistant message
+    assistant_msg = response["messages"][-1].content
+    assert isinstance(assistant_msg, str)
+
+    # Check tool message
+    tool_msg = response["messages"][-2]
+    assert tool_msg.name == "graphrag_reasoning"
+
+    # Check reasoning results
+    assert "Adalimumab" in assistant_msg
+    assert "TNF" in assistant_msg