langchain 1.0.0-alpha.6 → 1.0.0-alpha.8

package/dist/agents/middlewareAgent/nodes/AgentNode.js

@@ -1,33 +1,59 @@
 import { initChatModel } from "../../../chat_models/universal.js";
 import { MultipleStructuredOutputsError } from "../../errors.js";
-import { bindTools, getPromptRunnable, hasSupportForJsonSchemaOutput, hasToolCalls, validateLLMHasNoBoundTools } from "../../utils.js";
+import { bindTools, getPromptRunnable, hasToolCalls, isClientTool, validateLLMHasNoBoundTools } from "../../utils.js";
 import { RunnableCallable } from "../../RunnableCallable.js";
 import { mergeAbortSignals } from "../../nodes/utils.js";
 import { withAgentName } from "../../withAgentName.js";
-import { ProviderStrategy, ToolStrategy, transformResponseFormat } from "../../responses.js";
-import { parseToolCalls, parseToolResults } from "./utils.js";
-import { AIMessage, BaseMessage, SystemMessage, ToolMessage } from "@langchain/core/messages";
+import { ProviderStrategy, ToolStrategy, hasSupportForJsonSchemaOutput, transformResponseFormat } from "../../responses.js";
+import { AIMessage, ToolMessage } from "@langchain/core/messages";
 import { Command } from "@langchain/langgraph";
-import { getSchemaDescription } from "@langchain/core/utils/types";
+import { getSchemaDescription, interopParse } from "@langchain/core/utils/types";
 import { z } from "zod/v3";
 
 //#region src/agents/middlewareAgent/nodes/AgentNode.ts
 var AgentNode = class extends RunnableCallable {
 	#options;
-	#structuredToolInfo = {};
+	#runState = { runModelCallCount: 0 };
 	constructor(options) {
 		super({
 			name: options.name ?? "model",
 			func: (input, config) => this.#run(input, config)
 		});
 		this.#options = options;
+	}
+	/**
+	* Returns response format primtivies based on given model and response format provided by the user.
+	*
+	* If the the user selects a tool output:
+	* - return a record of tools to extract structured output from the model's response
+	*
+	* if the the user selects a native schema output or if the model supports JSON schema output:
+	* - return a provider strategy to extract structured output from the model's response
+	*
+	* @param model - The model to get the response format for.
+	* @returns The response format.
+	*/
+	#getResponseFormat(model) {
+		if (!this.#options.responseFormat) return void 0;
+		const strategies = transformResponseFormat(this.#options.responseFormat, void 0, model);
+		/**
+		* we either define a list of provider strategies or a list of tool strategies
+		*/
+		const isProviderStrategy = strategies.every((format) => format instanceof ProviderStrategy);
 		/**
 		* Populate a list of structured tool info.
 		*/
-		this.#structuredToolInfo = transformResponseFormat(this.#options.responseFormat).filter((format) => format instanceof ToolStrategy).reduce((acc, format) => {
-			acc[format.name] = format;
-			return acc;
-		}, {});
+		if (!isProviderStrategy) return {
+			type: "tool",
+			tools: strategies.filter((format) => format instanceof ToolStrategy).reduce((acc, format) => {
+				acc[format.name] = format;
+				return acc;
+			}, {})
+		};
+		return {
+			type: "native",
+			strategy: strategies[0]
+		};
 	}
 	async #run(state, config) {
 		/**
@@ -40,13 +66,20 @@ var AgentNode = class extends RunnableCallable {
 		* return directly without invoking the model again
 		*/
 		return { messages: [] };
+		const privateState = this.getState()._privateState;
 		const response = await this.#invokeModel(state, config);
+		this.#runState.runModelCallCount++;
+		const _privateState = {
+			...privateState,
+			threadLevelCallCount: privateState.threadLevelCallCount + 1
+		};
 		/**
 		* if we were able to generate a structured response, return it
 		*/
 		if ("structuredResponse" in response) return {
 			messages: [...state.messages, ...response.messages || []],
-			structuredResponse: response.structuredResponse
+			structuredResponse: response.structuredResponse,
+			_privateState
 		};
 		/**
 		* if we need to direct the agent to the model, return the update
@@ -54,12 +87,18 @@ var AgentNode = class extends RunnableCallable {
 		if (response instanceof Command) return response;
 		response.name = this.name;
 		response.lc_kwargs.name = this.name;
-		if (this.#areMoreStepsNeeded(state, response)) return { messages: [new AIMessage({
-			content: "Sorry, need more steps to process this request.",
-			name: this.name,
-			id: response.id
-		})] };
-		return { messages: [response] };
+		if (this.#areMoreStepsNeeded(state, response)) return {
+			messages: [new AIMessage({
+				content: "Sorry, need more steps to process this request.",
+				name: this.name,
+				id: response.id
+			})],
+			_privateState
+		};
+		return {
+			messages: [response],
+			_privateState
+		};
 	}
 	/**
 	* Derive the model from the options.
@@ -67,71 +106,127 @@ var AgentNode = class extends RunnableCallable {
 	* @param config - The config of the agent.
 	* @returns The model.
 	*/
-	#deriveModel(state, config) {
-		if (this.#options.model) {
-			if (typeof this.#options.model === "string") return initChatModel(this.#options.model);
-			throw new Error("`model` option must be a string.");
-		}
-		const model = this.#options.llm;
-		/**
-		* If the model is a function, call it to get the model.
-		*/
-		if (typeof model === "function") return model(state, config);
-		if (model) return model;
-		throw new Error("No model option was provided, either via `model` or via `llm` option.");
+	#deriveModel() {
+		if (typeof this.#options.model === "string") return initChatModel(this.#options.model);
+		if (this.#options.model) return this.#options.model;
+		throw new Error("No model option was provided, either via `model` option.");
 	}
 	async #invokeModel(state, config, options = {}) {
-		const model = await this.#deriveModel(state, config);
+		const model = await this.#deriveModel();
 		/**
-		* Execute prepareModelRequest hooks from beforeModelNodes
+		* Execute modifyModelRequest hooks from beforeModelNodes
 		*/
-		const preparedOptions = await this.#executePrepareModelRequestHooks(model, state, config);
+		let preparedOptions = await this.#executePrepareModelRequestHooks(model, state, config);
 		/**
-		* If user provides a model in the preparedOptions, use it,
-		* otherwise use the model from the options
+		* Retry loop for model invocation with error handling
+		* Hard limit of 100 attempts to prevent infinite loops from buggy middleware
 		*/
-		const finalModel = preparedOptions?.model ?? model;
-		/**
-		* Check if the LLM already has bound tools and throw if it does.
-		*/
-		validateLLMHasNoBoundTools(finalModel);
-		const modelWithTools = await this.#bindTools(finalModel, preparedOptions);
-		let modelInput = this.#getModelInputState(state);
-		if (preparedOptions?.messages) modelInput = {
-			...modelInput,
-			messages: preparedOptions.messages
-		};
-		const signal = mergeAbortSignals(this.#options.signal, config.signal);
-		const invokeConfig = {
-			...config,
-			signal
-		};
-		const response = await modelWithTools.invoke(modelInput, invokeConfig);
-		/**
-		* if the user requests a native schema output, try to parse the response
-		* and return the structured response if it is valid
-		*/
-		if (this.#options.responseFormat instanceof ProviderStrategy) {
-			const structuredResponse = this.#options.responseFormat.parse(response);
-			if (structuredResponse) return {
-				structuredResponse,
-				messages: [response]
+		const maxAttempts = 100;
+		for (let attempt = 1; attempt <= maxAttempts; attempt++) try {
+			/**
+			* If user provides a model in the preparedOptions, use it,
+			* otherwise use the model from the options
+			*/
+			const finalModel = preparedOptions?.model ?? model;
+			/**
+			* Check if the LLM already has bound tools and throw if it does.
+			*/
+			validateLLMHasNoBoundTools(finalModel);
+			const structuredResponseFormat = this.#getResponseFormat(finalModel);
+			const modelWithTools = await this.#bindTools(finalModel, preparedOptions, structuredResponseFormat);
+			let modelInput = this.#getModelInputState(state);
+			/**
+			* Use messages from preparedOptions if provided
+			*/
+			if (preparedOptions?.messages) modelInput = {
+				...modelInput,
+				messages: preparedOptions.messages
+			};
+			const signal = mergeAbortSignals(this.#options.signal, config.signal);
+			const invokeConfig = {
+				...config,
+				signal
 			};
+			const response = await modelWithTools.invoke(modelInput, invokeConfig);
+			/**
+			* if the user requests a native schema output, try to parse the response
+			* and return the structured response if it is valid
+			*/
+			if (structuredResponseFormat?.type === "native") {
+				const structuredResponse = structuredResponseFormat.strategy.parse(response);
+				if (structuredResponse) return {
+					structuredResponse,
+					messages: [response]
+				};
+				return response;
+			}
+			if (!structuredResponseFormat || !response.tool_calls) return response;
+			const toolCalls = response.tool_calls.filter((call) => call.name in structuredResponseFormat.tools);
+			/**
+			* if there were not structured tool calls, we can return the response
+			*/
+			if (toolCalls.length === 0) return response;
+			/**
+			* if there were multiple structured tool calls, we should throw an error as this
+			* scenario is not defined/supported.
+			*/
+			if (toolCalls.length > 1) return this.#handleMultipleStructuredOutputs(response, toolCalls, structuredResponseFormat);
+			const toolStrategy = structuredResponseFormat.tools[toolCalls[0].name];
+			const toolMessageContent = toolStrategy?.options?.toolMessageContent;
+			return this.#handleSingleStructuredOutput(response, toolCalls[0], structuredResponseFormat, toolMessageContent ?? options.lastMessage);
+		} catch (error) {
+			const retryMiddleware = this.#options.retryModelRequestHookMiddleware ?? [];
+			let shouldRetry = false;
+			for (const [middleware, getMiddlewareState] of retryMiddleware) if (middleware.retryModelRequest) {
+				/**
+				* Cast config to LangGraphRunnableConfig to access LangGraph-specific properties
+				*/
+				const lgConfig = config;
+				/**
+				* Merge context with default context of middleware
+				*/
+				const context = middleware.contextSchema ? interopParse(middleware.contextSchema, lgConfig?.context || {}) : lgConfig?.context;
+				/**
+				* Create runtime
+				*/
+				const privateState = this.getState()._privateState;
+				const runtime = {
+					...privateState,
+					context,
+					writer: lgConfig.writer,
+					interrupt: lgConfig.interrupt,
+					signal: lgConfig.signal
+				};
+				const retryRequest = await middleware.retryModelRequest(
+					error,
+					{
+						model: preparedOptions?.model ?? model,
+						systemPrompt: preparedOptions?.systemPrompt,
+						messages: preparedOptions?.messages ?? state.messages,
+						tools: this.#options.toolClasses
+					},
+					{
+						...getMiddlewareState(),
+						messages: state.messages
+					},
+					/**
+					* ensure runtime is frozen to prevent modifications
+					*/
+					Object.freeze({
+						...runtime,
+						context
+					}),
+					attempt
+				);
+				if (retryRequest) {
+					preparedOptions = retryRequest;
+					shouldRetry = true;
+					break;
+				}
+			}
+			if (!shouldRetry) throw error;
 		}
-		if (!response.tool_calls) return response;
-		const toolCalls = response.tool_calls.filter((call) => call.name in this.#structuredToolInfo);
-		/**
-		* if there were not structured tool calls, we can return the response
-		*/
-		if (toolCalls.length === 0) return response;
-		/**
-		* if there were multiple structured tool calls, we should throw an error as this
-		* scenario is not defined/supported.
-		*/
-		if (toolCalls.length > 1) return this.#handleMultipleStructuredOutputs(response, toolCalls);
-		const toolStrategy = this.#structuredToolInfo[toolCalls[0].name];
-		const toolMessageContent = toolStrategy?.options?.toolMessageContent;
-		return this.#handleSingleStructuredOutput(response, toolCalls[0], toolMessageContent ?? options.lastMessage);
+		throw new Error(`Maximum retry attempts (${maxAttempts}) exceeded`);
 	}
 	/**
 	* If the model returns multiple structured outputs, we need to handle it.
@@ -139,21 +234,17 @@ var AgentNode = class extends RunnableCallable {
 	* @param toolCalls - The tool calls that were made
 	* @returns The response from the model
 	*/
-	#handleMultipleStructuredOutputs(response, toolCalls) {
-		/**
-		* the following should never happen, let's throw an error if it does
-		*/
-		if (this.#options.responseFormat instanceof ProviderStrategy) throw new Error("Multiple structured outputs should not apply to native structured output responses");
+	#handleMultipleStructuredOutputs(response, toolCalls, responseFormat) {
 		const multipleStructuredOutputsError = new MultipleStructuredOutputsError(toolCalls.map((call) => call.name));
-		return this.#handleToolStrategyError(multipleStructuredOutputsError, response, toolCalls[0]);
+		return this.#handleToolStrategyError(multipleStructuredOutputsError, response, toolCalls[0], responseFormat);
 	}
 	/**
 	* If the model returns a single structured output, we need to handle it.
 	* @param toolCall - The tool call that was made
 	* @returns The structured response and a message to the LLM if needed
 	*/
-	#handleSingleStructuredOutput(response, toolCall, lastMessage) {
-		const tool = this.#structuredToolInfo[toolCall.name];
+	#handleSingleStructuredOutput(response, toolCall, responseFormat, lastMessage) {
+		const tool = responseFormat.tools[toolCall.name];
 		try {
 			const structuredResponse = tool.parse(toolCall.args);
 			return {
@@ -161,10 +252,10 @@ var AgentNode = class extends RunnableCallable {
 				messages: [response, new AIMessage(lastMessage ?? `Returning structured response: ${JSON.stringify(structuredResponse)}`)]
 			};
 		} catch (error) {
-			return this.#handleToolStrategyError(error, response, toolCall);
+			return this.#handleToolStrategyError(error, response, toolCall, responseFormat);
 		}
 	}
-	async #handleToolStrategyError(error, response, toolCall) {
+	async #handleToolStrategyError(error, response, toolCall, responseFormat) {
 		/**
 		* Using the `errorHandler` option of the first `ToolStrategy` entry is sufficient here.
 		* There is technically only one `ToolStrategy` entry in `structuredToolInfo` if the user
@@ -172,7 +263,7 @@ var AgentNode = class extends RunnableCallable {
 		* schema objects, these will be transformed into multiple `ToolStrategy` entries but all
 		* with the same `handleError` option.
 		*/
-		const errorHandler = Object.values(this.#structuredToolInfo).at(0)?.options?.handleError;
+		const errorHandler = Object.values(responseFormat.tools).at(0)?.options?.handleError;
 		const toolCallId = toolCall.id;
 		if (!toolCallId) throw new Error("Tool call ID is required to handle tool output errors. Please provide a tool call ID.");
 		/**
@@ -231,45 +322,80 @@ var AgentNode = class extends RunnableCallable {
 		};
 	}
 	async #executePrepareModelRequestHooks(model, state, config) {
-		if (!this.#options.prepareModelRequestHookMiddleware || this.#options.prepareModelRequestHookMiddleware.length === 0) return void 0;
-		let systemMessage;
-		if (typeof this.#options.prompt === "string") systemMessage = new SystemMessage(this.#options.prompt);
-		else if (BaseMessage.isInstance(this.#options.prompt)) systemMessage = this.#options.prompt;
+		if (!this.#options.modifyModelRequestHookMiddleware || this.#options.modifyModelRequestHookMiddleware.length === 0) return void 0;
+		/**
+		* Get the prompt for system message
+		*/
+		const systemPrompt = this.#options.systemPrompt;
+		/**
+		* Prepare the initial call options
+		*/
 		let currentOptions = {
 			model,
-			systemMessage,
+			systemPrompt,
 			messages: state.messages,
-			tools: []
+			tools: this.#options.toolClasses
 		};
-		const middlewareList = this.#options.prepareModelRequestHookMiddleware;
+		/**
+		* Execute modifyModelRequest hooks from all middleware
+		*/
+		const middlewareList = this.#options.modifyModelRequestHookMiddleware;
 		for (const [middleware, getMiddlewareState] of middlewareList) {
-			const context = {
-				...middleware.contextSchema?.parse({}) || {},
-				...config?.context || {}
-			};
+			/**
+			* Merge context with default context of middleware
+			*/
+			const context = middleware.contextSchema ? interopParse(middleware.contextSchema, config?.context || {}) : config?.context;
+			/**
+			* Create runtime
+			*/
+			const privateState = this.getState()._privateState;
 			const runtime = {
-				toolCalls: parseToolCalls(state.messages),
-				toolResults: parseToolResults(state.messages),
-				context
-			};
-			const result = await middleware.prepareModelRequest(currentOptions, {
-				messages: state.messages,
-				...getMiddlewareState()
-			}, {
-				...runtime,
-				context
-			});
-			if (result) currentOptions = {
-				...currentOptions,
-				...result
+				...privateState,
+				context,
+				writer: config.writer,
+				interrupt: config.interrupt,
+				signal: config.signal
 			};
+			const result = await middleware.modifyModelRequest(
+				currentOptions,
+				{
+					...getMiddlewareState(),
+					messages: state.messages
+				},
+				/**
+				* ensure runtime is frozen to prevent modifications
+				*/
+				Object.freeze({
+					...runtime,
+					context
+				})
+			);
+			if (result) {
+				const modifiedTools = result.tools ?? [];
+				/**
+				* Verify that the user didn't add any new tools.
+				* We can't allow this as the ToolNode is already initiated with given tools.
+				*/
+				const newTools = modifiedTools.filter((tool) => isClientTool(tool) && !this.#options.toolClasses.some((t) => t.name === tool.name));
+				if (newTools.length > 0) throw new Error(`You have added a new tool in "modifyModelRequest" hook of middleware "${middleware.name}": ${newTools.map((tool) => tool.name).join(", ")}. This is not supported.`);
+				/**
+				* Verify that user has not added or modified a tool with the same name.
+				* We can't allow this as the ToolNode is already initiated with given tools.
+				*/
+				const invalidTools = modifiedTools.filter((tool) => isClientTool(tool) && this.#options.toolClasses.every((t) => t !== tool));
+				if (invalidTools.length > 0) throw new Error(`You have modified a tool in "modifyModelRequest" hook of middleware "${middleware.name}": ${invalidTools.map((tool) => tool.name).join(", ")}. This is not supported.`);
+				currentOptions = {
+					...currentOptions,
+					...result
+				};
+			}
 		}
 		return currentOptions;
 	}
-	async #bindTools(model, preparedOptions) {
+	async #bindTools(model, preparedOptions, structuredResponseFormat) {
 		const options = {};
-		const structuredTools = Object.values(this.#structuredToolInfo);
-		const allTools = this.#options.toolClasses.concat(...structuredTools.map((toolStrategy) => toolStrategy.tool), ...preparedOptions?.tools || []);
+		const structuredTools = Object.values(structuredResponseFormat && "tools" in structuredResponseFormat ? structuredResponseFormat.tools : {});
+		const allTools = [...preparedOptions?.tools ?? [], ...structuredTools.map((toolStrategy) => toolStrategy.tool)];
 		/**
 		* If there are structured tools, we need to set the tool choice to "any"
 		* so that the model can choose to use a structured tool or not.
@@ -278,15 +404,15 @@ var AgentNode = class extends RunnableCallable {
 		/**
 		* check if the user requests a native schema output
 		*/
-		if (this.#options.responseFormat instanceof ProviderStrategy) {
+		if (structuredResponseFormat?.type === "native") {
 			/**
 			* if the model does not support JSON schema output, throw an error
 			*/
 			if (!hasSupportForJsonSchemaOutput(model)) throw new Error("Model does not support native structured output responses. Please use a model that supports native structured output responses or use a tool output.");
 			const jsonSchemaParams = {
-				name: this.#options.responseFormat.schema?.name ?? "extract",
-				description: getSchemaDescription(this.#options.responseFormat.schema),
-				schema: this.#options.responseFormat.schema,
+				name: structuredResponseFormat.strategy.schema?.name ?? "extract",
+				description: getSchemaDescription(structuredResponseFormat.strategy.schema),
+				schema: structuredResponseFormat.strategy.schema,
 				strict: true
 			};
 			Object.assign(options, {
@@ -296,7 +422,7 @@ var AgentNode = class extends RunnableCallable {
 				},
 				ls_structured_output_format: {
 					kwargs: { method: "json_schema" },
-					schema: this.#options.responseFormat.schema
+					schema: structuredResponseFormat.strategy.schema
 				},
 				strict: true
 			});
@@ -311,11 +437,28 @@ var AgentNode = class extends RunnableCallable {
 		/**
 		* Create a model runnable with the prompt and agent name
 		*/
-		const modelRunnable = getPromptRunnable(preparedOptions?.systemMessage ?? this.#options.prompt).pipe(this.#options.includeAgentName === "inline" ? withAgentName(modelWithTools, this.#options.includeAgentName) : modelWithTools);
+		const modelRunnable = getPromptRunnable(preparedOptions?.systemPrompt ?? this.#options.systemPrompt).pipe(this.#options.includeAgentName === "inline" ? withAgentName(modelWithTools, this.#options.includeAgentName) : modelWithTools);
 		return modelRunnable;
 	}
 	static get nodeOptions() {
-		return { input: z.object({ messages: z.array(z.custom()) }) };
+		return { input: z.object({
+			messages: z.array(z.custom()),
+			_privateState: z.object({ threadLevelCallCount: z.number() })
+		}) };
+	}
+	getState() {
+		const origState = super.getState() ?? { _privateState: {
+			threadLevelCallCount: 0,
+			runModelCallCount: 0
+		} };
+		return {
+			messages: [],
+			...origState,
+			_privateState: {
+				...origState._privateState ?? {},
+				...this.#runState
+			}
+		};
 	}
 };