Skip to content

Commit bb85051

Browse files
dqbddqbd
authored
fix(langgraph): apply Zod reducer when attempting to validate input (#1241)
2 parents a512380 + 4def333 commit bb85051

File tree

4 files changed

+275
-101
lines changed

4 files changed

+275
-101
lines changed

‎libs/langgraph/src/graph/state.ts

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ import { isPregelLike } from "../pregel/utils/subgraph.js";
5959
import {
6060
AnyZodObject,
6161
getChannelsFromZod,
62+
applyZodPlugin,
6263
isAnyZodObject,
6364
ZodToStateDefinition,
6465
} from "./zod/state.js";
@@ -982,18 +983,27 @@ export class CompiledStateGraph<
982983
protected async _validateInput(
983984
input: UpdateType<ToStateDefinition<I>>
984985
): Promise<UpdateType<ToStateDefinition<I>>> {
985-
let inputSchema = this.builder._inputRuntimeDefinition;
986-
if (inputSchema === PartialStateSchema) {
987-
inputSchema = this.builder._schemaRuntimeDefinition?.partial();
988-
}
986+
const schema = (() => {
987+
const input = this.builder._inputRuntimeDefinition;
988+
const schema = this.builder._schemaRuntimeDefinition;
989+
990+
const apply = (schema: AnyZodObject | undefined) => {
991+
if (schema == null) return undefined;
992+
return applyZodPlugin(schema, { reducer: true });
993+
};
994+
995+
if (isAnyZodObject(input)) return apply(input);
996+
if (input === PartialStateSchema) return apply(schema)?.partial();
997+
return undefined;
998+
})();
989999

9901000
if (isCommand(input)) {
9911001
const parsedInput = input;
992-
if (input.update && isAnyZodObject(inputSchema))
993-
parsedInput.update = inputSchema.parse(input.update);
1002+
if (input.update && schema != null)
1003+
parsedInput.update = schema.parse(input.update);
9941004
return parsedInput;
9951005
}
996-
if (isAnyZodObject(inputSchema)) return inputSchema.parse(input);
1006+
if (schema != null) return schema.parse(input);
9971007
return input;
9981008
}
9991009

‎libs/langgraph/src/graph/zod/schema.ts

Lines changed: 8 additions & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -1,64 +1,10 @@
1-
import { z } from "zod";
1+
import type { z } from "zod";
22
import { zodToJsonSchema as _zodToJsonSchema } from "zod-to-json-schema";
3-
import { getMeta } from "./state.js";
4-
5-
const TYPE_CACHE: Record<string, WeakMap<z.AnyZodObject, z.AnyZodObject>> = {};
6-
const DESCRIPTION_PREFIX = "lg:";
3+
import { applyZodPlugin, applyExtraFromDescription } from "./state.js";
74

85
const PartialStateSchema = Symbol.for("langgraph.state.partial");
96
type PartialStateSchema = typeof PartialStateSchema;
10-
11-
function applyPlugin(
12-
schema: z.AnyZodObject,
13-
actions: {
14-
/** Apply .langgraph.reducer calls */
15-
reducer?: boolean;
16-
17-
/** Apply .langgraph.metadata() calls */
18-
jsonSchemaExtra?: boolean;
19-
20-
/** Apply .partial() */
21-
partial?: boolean;
22-
}
23-
) {
24-
const cacheKey = [
25-
`reducer:${actions.reducer ?? false}`,
26-
`jsonSchemaExtra:${actions.jsonSchemaExtra ?? false}`,
27-
`partial:${actions.partial ?? false}`,
28-
].join("|");
29-
30-
TYPE_CACHE[cacheKey] ??= new WeakMap();
31-
const cache = TYPE_CACHE[cacheKey];
32-
33-
if (cache.has(schema)) return cache.get(schema)!;
34-
let shape = z.object({
35-
...Object.fromEntries(
36-
Object.entries(schema.shape as Record<string, z.ZodTypeAny>).map(
37-
([key, input]): [string, z.ZodTypeAny] => {
38-
const meta = getMeta(input);
39-
let output = actions.reducer ? meta?.reducer?.schema ?? input : input;
40-
41-
if (actions.jsonSchemaExtra) {
42-
const strMeta = JSON.stringify({
43-
...meta?.jsonSchemaExtra,
44-
description: output.description ?? input.description,
45-
});
46-
47-
if (strMeta !== "{}") {
48-
output = output.describe(`${DESCRIPTION_PREFIX}${strMeta}`);
49-
}
50-
}
51-
52-
return [key, output];
53-
}
54-
)
55-
),
56-
});
57-
58-
if (actions.partial) shape = shape.partial();
59-
cache.set(schema, shape);
60-
return shape;
61-
}
7+
type JsonSchema = ReturnType<typeof _zodToJsonSchema>;
628

639
// Using a subset of types to avoid circular type import
6410
interface GraphWithZodLike {
@@ -83,37 +29,6 @@ function isGraphWithZodLike(graph: unknown): graph is GraphWithZodLike {
8329
return true;
8430
}
8531

86-
type JsonSchema = ReturnType<typeof _zodToJsonSchema>;
87-
88-
function applyExtraFromDescription(schema: unknown): unknown {
89-
if (Array.isArray(schema)) {
90-
return schema.map(applyExtraFromDescription);
91-
}
92-
93-
if (typeof schema === "object" && schema != null) {
94-
const output = Object.fromEntries(
95-
Object.entries(schema).map(([key, value]) => [
96-
key,
97-
applyExtraFromDescription(value),
98-
])
99-
);
100-
101-
if (
102-
"description" in output &&
103-
typeof output.description === "string" &&
104-
output.description.startsWith(DESCRIPTION_PREFIX)
105-
) {
106-
const strMeta = output.description.slice(DESCRIPTION_PREFIX.length);
107-
delete output.description;
108-
Object.assign(output, JSON.parse(strMeta));
109-
}
110-
111-
return output;
112-
}
113-
114-
return schema as JsonSchema;
115-
}
116-
11732
function toJsonSchema(schema: z.ZodType): JsonSchema {
11833
return applyExtraFromDescription(_zodToJsonSchema(schema)) as JsonSchema;
11934
}
@@ -127,7 +42,7 @@ export function getStateTypeSchema(graph: unknown): JsonSchema | undefined {
12742
if (!isGraphWithZodLike(graph)) return undefined;
12843
const schemaDef = graph.builder._schemaRuntimeDefinition;
12944
if (!schemaDef) return undefined;
130-
return toJsonSchema(applyPlugin(schemaDef, { jsonSchemaExtra: true }));
45+
return toJsonSchema(applyZodPlugin(schemaDef, { jsonSchemaExtra: true }));
13146
}
13247

13348
/**
@@ -141,7 +56,7 @@ export function getUpdateTypeSchema(graph: unknown): JsonSchema | undefined {
14156
if (!schemaDef) return undefined;
14257

14358
return toJsonSchema(
144-
applyPlugin(schemaDef, {
59+
applyZodPlugin(schemaDef, {
14560
reducer: true,
14661
jsonSchemaExtra: true,
14762
partial: true,
@@ -164,7 +79,7 @@ export function getInputTypeSchema(graph: unknown): JsonSchema | undefined {
16479

16580
if (!schemaDef) return undefined;
16681
return toJsonSchema(
167-
applyPlugin(schemaDef, {
82+
applyZodPlugin(schemaDef, {
16883
reducer: true,
16984
jsonSchemaExtra: true,
17085
partial: true,
@@ -181,7 +96,7 @@ export function getOutputTypeSchema(graph: unknown): JsonSchema | undefined {
18196
if (!isGraphWithZodLike(graph)) return undefined;
18297
const schemaDef = graph.builder._outputRuntimeDefinition;
18398
if (!schemaDef) return undefined;
184-
return toJsonSchema(applyPlugin(schemaDef, { jsonSchemaExtra: true }));
99+
return toJsonSchema(applyZodPlugin(schemaDef, { jsonSchemaExtra: true }));
185100
}
186101

187102
/**
@@ -193,5 +108,5 @@ export function getConfigTypeSchema(graph: unknown): JsonSchema | undefined {
193108
if (!isGraphWithZodLike(graph)) return undefined;
194109
const configDef = graph.builder._configRuntimeSchema;
195110
if (!configDef) return undefined;
196-
return toJsonSchema(applyPlugin(configDef, { jsonSchemaExtra: true }));
111+
return toJsonSchema(applyZodPlugin(configDef, { jsonSchemaExtra: true }));
197112
}

‎libs/langgraph/src/graph/zod/state.ts

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,3 +122,103 @@ export function getChannelsFromZod<T extends z.ZodRawShape>(
122122
}
123123
return channels as ZodToStateDefinition<z.ZodObject<T>>;
124124
}
125+
126+
const ZOD_TYPE_CACHE: Record<
127+
string,
128+
WeakMap<z.AnyZodObject, z.AnyZodObject>
129+
> = {};
130+
131+
const ZOD_DESCRIPTION_PREFIX = "lg:";
132+
133+
export function applyZodPlugin(
134+
schema: z.AnyZodObject,
135+
actions: {
136+
/** Apply .langgraph.reducer calls */
137+
reducer?: boolean;
138+
139+
/** Apply .langgraph.metadata() calls */
140+
jsonSchemaExtra?: boolean;
141+
142+
/** Apply .partial() */
143+
partial?: boolean;
144+
}
145+
) {
146+
const cacheKey = [
147+
`reducer:${actions.reducer ?? false}`,
148+
`jsonSchemaExtra:${actions.jsonSchemaExtra ?? false}`,
149+
`partial:${actions.partial ?? false}`,
150+
].join("|");
151+
152+
ZOD_TYPE_CACHE[cacheKey] ??= new WeakMap();
153+
const cache = ZOD_TYPE_CACHE[cacheKey];
154+
155+
if (cache.has(schema)) return cache.get(schema)!;
156+
157+
let shape = schema.extend({
158+
...Object.fromEntries(
159+
Object.entries(schema.shape as Record<string, z.ZodTypeAny>).map(
160+
([key, input]): [string, z.ZodTypeAny] => {
161+
const meta = getMeta(input);
162+
let output = actions.reducer ? meta?.reducer?.schema ?? input : input;
163+
164+
if (actions.jsonSchemaExtra) {
165+
const strMeta = JSON.stringify({
166+
...meta?.jsonSchemaExtra,
167+
description: output.description ?? input.description,
168+
});
169+
170+
if (strMeta !== "{}") {
171+
output = output.describe(`${ZOD_DESCRIPTION_PREFIX}${strMeta}`);
172+
}
173+
}
174+
175+
return [key, output];
176+
}
177+
)
178+
),
179+
});
180+
181+
// using zObject.extend() will set `unknownKeys` to `passthrough`
182+
// which trips up `zod-to-json-schema`
183+
if (
184+
"_def" in shape &&
185+
shape._def != null &&
186+
typeof shape._def === "object" &&
187+
"unknownKeys" in shape._def
188+
) {
189+
shape._def.unknownKeys = "strip";
190+
}
191+
192+
if (actions.partial) shape = shape.partial();
193+
cache.set(schema, shape);
194+
return shape;
195+
}
196+
197+
export function applyExtraFromDescription(schema: unknown): unknown {
198+
if (Array.isArray(schema)) {
199+
return schema.map(applyExtraFromDescription);
200+
}
201+
202+
if (typeof schema === "object" && schema != null) {
203+
const output = Object.fromEntries(
204+
Object.entries(schema).map(([key, value]) => [
205+
key,
206+
applyExtraFromDescription(value),
207+
])
208+
);
209+
210+
if (
211+
"description" in output &&
212+
typeof output.description === "string" &&
213+
output.description.startsWith(ZOD_DESCRIPTION_PREFIX)
214+
) {
215+
const strMeta = output.description.slice(ZOD_DESCRIPTION_PREFIX.length);
216+
delete output.description;
217+
Object.assign(output, JSON.parse(strMeta));
218+
}
219+
220+
return output;
221+
}
222+
223+
return schema;
224+
}

0 commit comments

Comments
 (0)