diff --git a/libs/langchain-core/src/utils/types/tests/zod.test.ts b/libs/langchain-core/src/utils/types/tests/zod.test.ts index 7d7cabad9b25..0e61ae58454a 100644 --- a/libs/langchain-core/src/utils/types/tests/zod.test.ts +++ b/libs/langchain-core/src/utils/types/tests/zod.test.ts @@ -1525,6 +1525,209 @@ describe("Zod utility functions", () => { expect(elementShape.name).toBeInstanceOf(z4.ZodString); expect(elementShape.age).toBeInstanceOf(z4.ZodNumber); }); + + it("should cache and reuse sanitized sub-schemas when same schema is used in multiple properties", () => { + // Create a shared sub-schema that will be used in multiple places + const addressSchema = z4.object({ + street: z4.string().transform((s) => s.toUpperCase()), + city: z4.string(), + }); + + const inputSchema = z4.object({ + homeAddress: addressSchema, + workAddress: addressSchema, + billingAddress: addressSchema, + }); + + const result = interopZodTransformInputSchema(inputSchema, true); + + expect(result).toBeInstanceOf(z4.ZodObject); + const resultShape = getInteropZodObjectShape(result as any); + + // All three address properties should reference the exact same sanitized instance + expect(resultShape.homeAddress).toBe(resultShape.workAddress); + expect(resultShape.workAddress).toBe(resultShape.billingAddress); + expect(resultShape.homeAddress).toBe(resultShape.billingAddress); + }); + + it("should cache and reuse sanitized sub-schemas when same schema is used in array and as property", () => { + // Create a shared user schema + const userSchema = z4.object({ + name: z4.string().transform((s) => s.toUpperCase()), + age: z4.number(), + }); + + const inputSchema = z4.object({ + primaryUser: userSchema, + secondaryUser: userSchema, + allUsers: z4.array(userSchema), + }); + + const result = interopZodTransformInputSchema(inputSchema, true); + + expect(result).toBeInstanceOf(z4.ZodObject); + const resultShape = getInteropZodObjectShape(result as any); + + // Extract the array element schema + const arrayElement = (resultShape.allUsers as any)._zod.def.element; + + // All user schema references should be the exact same instance + expect(resultShape.primaryUser).toBe(resultShape.secondaryUser); + expect(resultShape.primaryUser).toBe(arrayElement); + expect(resultShape.secondaryUser).toBe(arrayElement); + }); + + it("should cache and reuse deeply nested shared schemas", () => { + // Create a deeply nested shared schema + const tagSchema = z4.object({ + name: z4.string().transform((s) => s.toLowerCase()), + color: z4.string(), + }); + + const postSchema = z4.object({ + title: z4.string(), + tags: z4.array(tagSchema), + primaryTag: tagSchema, + }); + + const inputSchema = z4.object({ + posts: z4.array(postSchema), + featuredPost: postSchema, + }); + + const result = interopZodTransformInputSchema(inputSchema, true); + + expect(result).toBeInstanceOf(z4.ZodObject); + const resultShape = getInteropZodObjectShape(result as any); + + // Extract the post schema from the array + const postsArrayElement = (resultShape.posts as any)._zod.def.element; + + // The featuredPost and the array element should be the same instance + expect(resultShape.featuredPost).toBe(postsArrayElement); + + // Extract tag schemas from both places + const featuredPostShape = getInteropZodObjectShape( + resultShape.featuredPost as any + ); + const postsArrayElementShape = getInteropZodObjectShape( + postsArrayElement as any + ); + + // Tag schemas should also be reused + const featuredPostTagsArray = (featuredPostShape.tags as any)._zod.def + .element; + const postsArrayTagsArray = (postsArrayElementShape.tags as any)._zod + .def.element; + + expect(featuredPostShape.primaryTag).toBe(featuredPostTagsArray); + expect(postsArrayElementShape.primaryTag).toBe(postsArrayTagsArray); + expect(featuredPostTagsArray).toBe(postsArrayTagsArray); + }); + + it("should handle caching when same schema appears at different nesting levels", () => { + // Create a schema that appears at different nesting levels + const metadataSchema = z4.object({ + key: z4.string(), + value: z4.string().transform((s) => s.trim()), + }); + + const inputSchema = z4.object({ + topLevelMetadata: metadataSchema, + nested: z4.object({ + nestedMetadata: metadataSchema, + deepNested: z4.object({ + deepMetadata: metadataSchema, + }), + }), + metadataList: z4.array(metadataSchema), + }); + + const result = interopZodTransformInputSchema(inputSchema, true); + + expect(result).toBeInstanceOf(z4.ZodObject); + const resultShape = getInteropZodObjectShape(result as any); + const nestedShape = getInteropZodObjectShape(resultShape.nested as any); + const deepNestedShape = getInteropZodObjectShape( + nestedShape.deepNested as any + ); + const metadataListElement = (resultShape.metadataList as any)._zod.def + .element; + + // All metadata schemas should reference the same sanitized instance + expect(resultShape.topLevelMetadata).toBe(nestedShape.nestedMetadata); + expect(nestedShape.nestedMetadata).toBe(deepNestedShape.deepMetadata); + expect(deepNestedShape.deepMetadata).toBe(metadataListElement); + expect(resultShape.topLevelMetadata).toBe(metadataListElement); + }); + + it("should generate JSON schema with $ref for reused sub-schemas", () => { + // Create a shared address schema with transforms + const addressSchema = z4.object({ + street: z4.string().transform((s) => s.toUpperCase()), + city: z4.string(), + zipCode: z4.string(), + }); + + const inputSchema = z4.object({ + homeAddress: addressSchema, + workAddress: addressSchema, + billingAddress: addressSchema, + shippingAddresses: z4.array(addressSchema), + }); + + // Get the sanitized input schema with transforms removed + const sanitizedSchema = interopZodTransformInputSchema( + inputSchema, + true + ); + + // Generate JSON schema with refs for reused schemas + const jsonSchema = z4.toJSONSchema(sanitizedSchema as any, { + cycles: "ref", + reused: "ref", + override(ctx) { + ctx.jsonSchema.title = "TestSchema"; + }, + }); + + // Verify exact JSON schema structure with $ref for reused schemas + expect(jsonSchema).toEqual({ + $schema: "https://json-schema.org/draft/2020-12/schema", + title: "TestSchema", + type: "object", + properties: { + homeAddress: { $ref: "#/$defs/__schema0" }, + workAddress: { $ref: "#/$defs/__schema0" }, + billingAddress: { $ref: "#/$defs/__schema0" }, + shippingAddresses: { + title: "TestSchema", + type: "array", + items: { $ref: "#/$defs/__schema0" }, + }, + }, + required: [ + "homeAddress", + "workAddress", + "billingAddress", + "shippingAddresses", + ], + additionalProperties: false, + $defs: { + __schema0: { + title: "TestSchema", + type: "object", + properties: { + street: { title: "TestSchema", type: "string" }, + city: { title: "TestSchema", type: "string" }, + zipCode: { title: "TestSchema", type: "string" }, + }, + required: ["street", "city", "zipCode"], + additionalProperties: false, + }, + }, + }); + }); }); it("should throw error for non-schema values", () => { diff --git a/libs/langchain-core/src/utils/types/zod.ts b/libs/langchain-core/src/utils/types/zod.ts index b2e902ac2c8b..b13e48b25715 100644 --- a/libs/langchain-core/src/utils/types/zod.ts +++ b/libs/langchain-core/src/utils/types/zod.ts @@ -779,23 +779,24 @@ function isZodTransformV4(schema: InteropZodType): schema is z4.$ZodPipe { return isZodSchemaV4(schema) && schema._zod.def.type === "pipe"; } -/** - * Returns the input type of a Zod transform schema, for both v3 and v4. - * If the schema is not a transform, returns undefined. If `recursive` is true, - * recursively processes nested object schemas and arrays of object schemas. - * - * @param schema - The Zod schema instance (v3 or v4) - * @param {boolean} [recursive=false] - Whether to recursively process nested objects/arrays. - * @returns The input Zod schema of the transform, or undefined if not a transform - */ -export function interopZodTransformInputSchema( +function interopZodTransformInputSchemaImpl( schema: InteropZodType, - recursive = false + recursive: boolean, + cache: WeakMap ): InteropZodType { + const cached = cache.get(schema); + if (cached !== undefined) { + return cached; + } + // Zod v3: ._def.schema is the input schema for ZodEffects (transform) if (isZodSchemaV3(schema)) { if (isZodTransformV3(schema)) { - return interopZodTransformInputSchema(schema._def.schema, recursive); + return interopZodTransformInputSchemaImpl( + schema._def.schema, + recursive, + cache + ); } // TODO: v3 schemas aren't recursively handled here // (currently not necessary since zodToJsonSchema handles this) @@ -806,9 +807,10 @@ export function interopZodTransformInputSchema( if (isZodSchemaV4(schema)) { let outputSchema: InteropZodType = schema; if (isZodTransformV4(schema)) { - outputSchema = interopZodTransformInputSchema( + outputSchema = interopZodTransformInputSchemaImpl( schema._zod.def.in, - recursive + recursive, + cache ); } if (recursive) { @@ -818,9 +820,10 @@ export function interopZodTransformInputSchema( for (const [key, keySchema] of Object.entries( outputSchema._zod.def.shape )) { - outputShape[key] = interopZodTransformInputSchema( + outputShape[key] = interopZodTransformInputSchemaImpl( keySchema, - recursive + recursive, + cache ) as z4.$ZodType; } outputSchema = clone(outputSchema, { @@ -830,9 +833,10 @@ export function interopZodTransformInputSchema( } // Handle nested array schemas else if (isZodArrayV4(outputSchema)) { - const elementSchema = interopZodTransformInputSchema( + const elementSchema = interopZodTransformInputSchemaImpl( outputSchema._zod.def.element, - recursive + recursive, + cache ); outputSchema = clone(outputSchema, { ...outputSchema._zod.def, @@ -842,12 +846,30 @@ export function interopZodTransformInputSchema( } const meta = globalRegistry.get(schema); if (meta) globalRegistry.add(outputSchema as z4.$ZodType, meta); + cache.set(schema, outputSchema); return outputSchema; } throw new Error("Schema must be an instance of z3.ZodType or z4.$ZodType"); } +/** + * Returns the input type of a Zod transform schema, for both v3 and v4. + * If the schema is not a transform, returns undefined. If `recursive` is true, + * recursively processes nested object schemas and arrays of object schemas. + * + * @param schema - The Zod schema instance (v3 or v4) + * @param {boolean} [recursive=false] - Whether to recursively process nested objects/arrays. + * @returns The input Zod schema of the transform, or undefined if not a transform + */ +export function interopZodTransformInputSchema( + schema: InteropZodType, + recursive = false +): InteropZodType { + const cache = new WeakMap(); + return interopZodTransformInputSchemaImpl(schema, recursive, cache); +} + /** * Creates a modified version of a Zod object schema where fields matching a predicate are made optional. * Supports both Zod v3 and v4 schemas and preserves the original schema version.