Injecting Tool Args into MCP Tools

[Blake-Martin-code]

I may have an unique use case.

As an overview, I am trying to inject tool arguments into my models mcp tool calls. Basically I have a vector Search tool that accepts 'source_filters': list[str]. This argument narrows down the vector search to chunks that have this meta data. It allows the user to specify which sources they want to look at.

This source_filters is a tool argument which the model can fill out but I override to the users selection and ensure the user has permissions to the sources incase of prompt injection. The way I override this is by passing the user selected sources and allowed sources to the RunnableConfig inside my custom agent graph. I then created a custom Tool Node which inherits from tool node which replaces the tool args before the tool is invoked.

class ChatAgent:

    def __init__(self, model: ChatOpenAI, tools: list, prompt: str = None, name: str = "chat_agent"):

        self.llm = model.bind_tools(tools)
        self.tools = tools
        self.prompt = prompt
        self.name = name

    def update_system_prompt(self, state: MessagesState, config: RunnableConfig):

        # update prompt with current date
        prompt = self.prompt.format(
            current_date=datetime.now().strftime("%Y-%m-%d")
        )
        # if the first message is a system message, update the content
        if isinstance(state["messages"][0], SystemMessage):
            state["messages"][0].content = prompt
        else:
            # otherwise, insert a new system message at the beginning
            state["messages"].insert(0, SystemMessage(content=prompt))

    def chat(self, state: MessagesState, config: RunnableConfig):

        # update prompt in messages. 
        # This handles replacing the placeholder system prompt and updating potential variables
        self.update_system_prompt(state, config)

        # determine if streaming is enabled
        # this is to handle the case where the model does not support streaming tool calls
        if int(os.getenv('STREAMING', '1')):
            return {"messages": [self.llm.invoke(state["messages"])]}
        else:
            return {"messages": [self.llm.invoke(state["messages"], stream=False)]}
    
    def build_graph(self):

        self.graph = StateGraph(MessagesState)
        self.graph.add_node("chatbot", self.chat)

        self.graph.add_node("tools", CustomToolNode(tools=self.tools))

        self.graph.add_conditional_edges(
            "chatbot",
            tools_condition,
            "tools",
        )
        self.graph.add_edge("tools", "chatbot")
        self.graph.add_edge(START, "chatbot")

        self.graph = self.graph.compile(name=self.name)
    
    def get_graph(self):
        return self.graph 

class CustomToolNode(ToolNode):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
    
    def handle_data_sources(self, tool_call: dict[str, Any], config: RunnableConfig) -> list[str]:
            # handle data sources filtering. Current implementation allows model to choose sources, but we need to filter them to only the ones that are available.
            # future implementation could be to only allow user sources to be chosen. Similar to the rag algorithm.
            config = config.get("configurable", {})
            model_chosen_sources = tool_call['args'].get("sources", [])
            if model_chosen_sources:
                sources = [source for source in model_chosen_sources if source in config.get("sources", [])]
                tool_call['args']['sources'] = sources
                return tool_call
            else:
                tool_call['args']['sources'] = config.get("sources", [])
                return tool_call

    def handle_tool_call(self, tool_call: dict[str, Any], config: RunnableConfig) -> dict[str, Any]:
        # handle tool call - custom tool arg injection
        tool_name = tool_call['name']
        match tool_name:
            case "searchVectorDatabase":
                tool_call = self.handle_data_sources(tool_call, config)
            case _:
                return tool_call
        return tool_call

    def _run_one(
        self,
        call: ToolCall,
        input_type: Literal["list", "dict", "tool_calls"],
        config: RunnableConfig,
    ) -> ToolMessage:
        """Run a single tool call synchronously."""
        if invalid_tool_message := self._validate_tool_call(call):
            return invalid_tool_message
        try:
            # handle tool call - custom tool arg injection handles individual cases for mcp tools
            call_args = self.handle_tool_call(
                {**call, **{"type": "tool_call"}}, 
                config
            )
            response = self.tools_by_name[call["name"]].invoke(call_args, config)

I did it this way because my mcp servers are using langchain_mcp_adapters and I cannot use the InjectedArgs through mcp.

With langchain v1 I was hoping I could just use the default create_agents with no custom agent or custom tool node. Then I could wrap my tool calls in a wrap_tool_call middle ware, grab the config and update the tool args. However It appears that I cannot get the config in the below code:

class ToolMonitoringMiddleware(AgentMiddleware):
    def wrap_tool_call(
        self,
        request: ToolCallRequest,
        handler: Callable[[ToolCallRequest], ToolMessage | Command],
    ) -> ToolMessage | Command:
        print(f"Executing tool: {request.tool_call['name']}")
        print(f"Arguments: {request.tool_call['args']}")
        print(f"Request: {request}")

        try:
            result = handler(request)
            print(f"Tool completed successfully")
            return result
        except Exception as e:
            print(f"Tool failed: {e}")
            raise

Is anyone aware of how to implement this or can guide me?