Use withStructuredOutput with multiple tools

[tristanretali]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

// @ts-nocheck
import { ChatMistralAI } from '@langchain/mistralai';
import * as output_parsers_1 from '@langchain/core/output_parsers';
import { zodToJsonSchema } from 'zod-to-json-schema';
import * as openai_tools_1 from '@langchain/core/output_parsers/openai_tools';
import * as runnables_1 from '@langchain/core/runnables';
import { secureHeapUsed } from 'node:crypto';
import { ZodSchema } from 'zod';

export default class CustomMistralAI extends ChatMistralAI {
	withStructuredOutput(outputSchema, config) {
		// eslint-disable-next-line @typescript-eslint/no-explicit-any

		if (outputSchema.length == undefined) {
			let schema;
			let name;
			let method;
			let includeRaw;
			if (isStructuredOutputMethodParams(outputSchema)) {
				schema = outputSchema.schema;
				name = outputSchema.name;
				method = outputSchema.method;
				includeRaw = outputSchema.includeRaw;
			} else {
				schema = outputSchema;
				name = config?.name;
				method = config?.method;
				includeRaw = config?.includeRaw;
			}
			let llm;
			let outputParser;
			if (method === 'jsonMode') {
				llm = this.bind({
					response_format: { type: 'json_object' }
				});
				if (isZodSchema(schema)) {
					outputParser = output_parsers_1.StructuredOutputParser.fromZodSchema(schema);
				} else {
					outputParser = new output_parsers_1.JsonOutputParser();
				}
			} else {
				const functionName = name ?? 'extract';
				if (isZodSchema(schema)) {
					const asJsonSchema = zodToJsonSchema(schema);
					llm = this.bind({
						tools: [
							{
								type: 'function',
								function: {
									name: functionName,
									description: asJsonSchema.description,
									parameters: asJsonSchema
								}
							}
						],
						tool_choice: 'any'
					});
					outputParser = new openai_tools_1.JsonOutputKeyToolsParser({
						returnSingle: false,
						keyName: functionName,
						zodSchema: schema
					});
				} else {
					llm = this.bind({
						tools: [
							{
								type: 'function',
								function: {
									name: functionName,
									description: schema.description,
									parameters: schema
								}
							}
						],
						tool_choice: 'any'
					});
					outputParser = new openai_tools_1.JsonOutputKeyToolsParser({
						returnSingle: true,
						keyName: functionName
					});
				}
			}
			if (!includeRaw) {
				return llm.pipe(outputParser);
			}
			const parserAssign = runnables_1.RunnablePassthrough.assign({
				// eslint-disable-next-line @typescript-eslint/no-explicit-any
				parsed: (input, config2) => outputParser.invoke(input.raw, config2)
			});
			const parserNone = runnables_1.RunnablePassthrough.assign({
				parsed: () => null
			});
			const parsedWithFallback = parserAssign.withFallbacks({
				fallbacks: [parserNone]
			});
			return runnables_1.RunnableSequence.from([
				{
					raw: llm
				},
				parsedWithFallback
			]);
		} else {
			let schemas = [];
			let names = [];
			let methods = [];
			let includeRaws = [];
			for (var output of outputSchema) {
				if (isStructuredOutputMethodParams(output)) {
					schemas.push(output.schema);
					names.push(output.name);
					methods.push(output.method);
					includeRaws.push(output.includeRaw);
				} else {
					schemas.push(outputSchema);
					names.push(config?.name);
					methods.push(config?.method);
					includeRaws.push(config?.includeRaw);
				}
			}
			let llm;
			let outputParser = [];
			if (methods[0] === 'jsonMode') {
				llm = this.bind({
					response_format: { type: 'json_object' }
				});
				for (var schema of schemas) {
					if (isZodSchema(schema)) {
						outputParser.push(output_parsers_1.StructuredOutputParser.fromZodSchema(schema));
					} else {
						outputParser.push(new output_parsers_1.JsonOutputParser());
					}
				}
			} else {
				if (isZodSchema(schemas[0])) {
					llm = this.bind({
						tools: createTools(schemas, names),
						tool_choice: 'any'
					});
				} else {
					llm = this.bind({
						tools: createTools(schemas, names),
						tool_choice: 'any'
					});
				}
			}
			if (!includeRaws[0]) {
				return llm;
			}
			const parserAssign = runnables_1.RunnablePassthrough.assign({
				// eslint-disable-next-line @typescript-eslint/no-explicit-any
				parsed: (input, config2) => outputParser[0].invoke(input.raw, config2)
			});
			const parserNone = runnables_1.RunnablePassthrough.assign({
				parsed: () => null
			});
			const parsedWithFallback = parserAssign.withFallbacks({
				fallbacks: [parserNone]
			});
			return runnables_1.RunnableSequence.from([
				{
					raw: llm
				},
				parsedWithFallback
			]);
		}
	}
}

function createTools(schemas: ZodSchema[], names: string[]) {
	let tools = [];
	for (let i = 0; i < schemas.length; i++) {
		let functionName = names[i] ?? 'extract';
		let asJsonSchema = zodToJsonSchema(schemas[i]);
		tools.push({
			type: 'function',
			function: {
				name: functionName,
				description: asJsonSchema.description,
				parameters: asJsonSchema
			}
		});
	}
	return tools;
}

function isZodSchema(
	// eslint-disable-next-line @typescript-eslint/no-explicit-any
	input
) {
	// Check for a characteristic method of Zod schemas
	return typeof input?.parse === 'function';
}

function isStructuredOutputMethodParams(
	x
	// eslint-disable-next-line @typescript-eslint/no-explicit-any
) {
	return (
		x !== undefined &&
		// eslint-disable-next-line @typescript-eslint/no-explicit-any
		typeof x.schema === 'object'
	);
}

Description

I tried to use the withStructuredOutput() method with a Mistral model to implement multiple tools. But I found nothing on the Langchain documentation or any forum. So, I needed to make my own. I would like to know if this is a normal thing or you plan to implement it?

System Info

platform: linux
Node.js v20.7.0
yarn version v1.22.19

[dosubot[bot]]

Hey there, @tristanretali! I'm here to help you with any bugs, questions, or contributions while we wait for a human maintainer. Let's tackle this together!

Yes, there are plans to officially support the withStructuredOutput() method with a Mistral model in LangChain. The documentation explicitly mentions that LangChain implements a common interface called .withStructuredOutput for several popular model providers, including Mistral ^[1].

Regarding your custom implementation for handling multiple tools, it is indeed a valid approach. However, there are some performance considerations and optimizations you can apply when using the withStructuredOutput() method with multiple tools:

1. Force Tool Selection: You can use the force option in the configuration to enforce the use of a specific tool, which can help optimize the tool selection process and reduce overhead.

   const modelWithStructuredOutput = model.withStructuredOutput(
     outputSchema,
     {
       name: "toolName",
       force: true,
     }
   );

2. Efficient Schema Handling: Ensure that the schemas are efficiently handled and validated. Using Zod schemas can help in maintaining a consistent and validated structure for the output.

3. Error Handling: Properly handle errors and invalid input schemas to ensure robustness. The FunctionCallStructuredOutputParser class provides functionality for parsing and validating the structured output based on a JSON schema or Zod schema ^[2].

4. Include Raw Output for Debugging: If necessary, include the raw output for debugging purposes by setting includeRaw to true.

   const modelWithStructuredOutput = model.withStructuredOutput(
     outputSchema,
     {
       name: "toolName",
       includeRaw: true,
     }
   );

By following these best practices, you can optimize the performance and reliability of using the withStructuredOutput() method with multiple tools in LangChain.

To continue talking to Dosu, mention @dosu.