aiagents4pharma 0.0.0__py3-none-any.whl

aiagents4pharma/talk2knowledgegraphs/tools/multimodal_subgraph_extraction.py

@@ -0,0 +1,374 @@
+"""
+Tool for performing multimodal subgraph extraction.
+"""
+
+import logging
+from typing import Annotated
+
+import hydra
+import joblib
+import networkx as nx
+import numpy as np
+import pandas as pd
+import torch
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import BaseTool
+from langchain_core.tools.base import InjectedToolCallId
+from langgraph.prebuilt import InjectedState
+from langgraph.types import Command
+from pydantic import BaseModel, Field
+from torch_geometric.data import Data
+
+from ..utils.embeddings.ollama import EmbeddingWithOllama
+from ..utils.extractions.multimodal_pcst import MultimodalPCSTPruning
+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,
+            strict=False,
+        )
+        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 using secure joblib
+        initial_graph["pyg"] = joblib.load(initial_graph["source"]["kg_pyg_path"])
+        initial_graph["text"] = joblib.load(initial_graph["source"]["kg_text_path"])
+
+        # 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/tools/subgraph_extraction.py

@@ -0,0 +1,291 @@
+"""
+Tool for performing subgraph extraction.
+"""
+
+import logging
+from typing import Annotated
+
+import hydra
+import joblib
+import networkx as nx
+import numpy as np
+import pandas as pd
+import torch
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain.chains.retrieval import create_retrieval_chain
+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 langchain_text_splitters import RecursiveCharacterTextSplitter
+from langgraph.prebuilt import InjectedState
+from langgraph.types import Command
+from pydantic import BaseModel, Field
+from torch_geometric.data import Data
+
+from ..utils.embeddings.ollama import EmbeddingWithOllama
+from ..utils.extractions.pcst import PCSTPruning
+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 using secure joblib
+        initial_graph["pyg"] = joblib.load(initial_graph["source"]["kg_pyg_path"])
+        initial_graph["text"] = joblib.load(initial_graph["source"]["kg_text_path"])
+
+        # 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,
+                    )
+                ],
+            }
+        )