Support returning state in wrapModelCall hook

OSS Product Help LangChain
,
3

Let me describe my current issue in more detail:

Currently I am trying to create a custom middleware called preHITLValidationMiddleware, which is designed to use in sequence with humanInTheLoopMiddleware, since I wish to have the model generated args to be validated first, before raising an interrupt to the client.

Sample code:

import { AIMessage, createMiddleware } from "langchain";
import { validateToolCalls } from "./validate-tool-calls";

const maxRetries = 3;

export const preHITLValidationMiddleware = createMiddleware({
    name: "PreHITLValidationMiddleware",
    wrapModelCall: async (request, handler) => {

        /* Initial attempt + retries */
        for (let attempt = 0; attempt <= maxRetries; attempt++) {
            /* 1. Call the model */
            const response = await handler(request);

            /* 2. If no tool calls, return immediately */
            if (!response.tool_calls?.length) return response;

            /* 3. Assign a tool id to tool calls that don't have one */
            const toolCallsWithId = response.tool_calls.map((call) => ({
                ...call,
                id: call.id ?? crypto.randomUUID(),
            }));

            /* 4. Validate tool calls */
            const result = await validateToolCalls(
                toolCallsWithId,
                request.tools,
            );

            if (result.success) {
                /* 5. Success, replace the tool messages args with transformed args */
                return new AIMessage({
                    ...response,
                    tool_calls: toolCallsWithId.map((call) => ({
                        ...call,
                        /* ! here should be correct */
                        args: result.argsByToolCallId[call.id]!,
                    })),
                });
            } else {
                /* 4. Failed, run handler again */
                if (attempt < maxRetries) {
                    request = {
                        ...request,
                        messages: [
                            ...request.messages,
                            ...result.invalidToolCallMessages,
                        ],
                    };
                } else {
                    return new AIMessage({
                        content: "Failed to validate tool calls ",
                    });
                    // finalMessage = result.invalidToolCallMessages
                    //     .map((m) => m.name)
                    //     .join("\n\n");
                }
            }
        }
        return new AIMessage({
            invalid_tool_calls: result.invalidToolCallMessages,
        });
    },
});

As you can see it is kind of like the toolRetryMIddleware, where provided a maxRetry, it will loop for certain of times, and if failed it will append a ToolMessage with status error to the request before passing to the model. Now, since wrap model call only allow returning AIMessage, the tool message containing the error status will be lost in the final output. If i use beforeModel or afterModel hooks instead, i will lose track of the retry state, and the graph might run in a n infinite loop potentially.