Skip to content

Commit d42d299

Browse files
fix(langchain/createAgent): improve interop between Zod v3 and v4
1 parent d823370 commit d42d299

File tree

11 files changed

+197
-100
lines changed

11 files changed

+197
-100
lines changed

libs/langchain-core/src/utils/types/zod.ts

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,29 @@ export type ZodObjectV3 = z3.ZodObject<any, any, any, any>;
2222

2323
export type ZodObjectV4 = z4.$ZodObject;
2424

25+
export type ZodDefaultV3<T extends z3.ZodTypeAny> = z3.ZodDefault<T>;
26+
export type ZodDefaultV4<T extends z4.SomeType> = z4.$ZodDefault<T>;
27+
export type ZodOptionalV3<T extends z3.ZodTypeAny> = z3.ZodOptional<T>;
28+
export type ZodOptionalV4<T extends z4.SomeType> = z4.$ZodOptional<T>;
29+
2530
// eslint-disable-next-line @typescript-eslint/no-explicit-any
2631
export type InteropZodType<Output = any, Input = Output> =
2732
| z3.ZodType<Output, z3.ZodTypeDef, Input>
2833
| z4.$ZodType<Output, Input>;
2934

3035
export type InteropZodObject = ZodObjectV3 | ZodObjectV4;
36+
export type InteropZodDefault<T = InteropZodObjectShape> =
37+
T extends z3.ZodTypeAny
38+
? ZodDefaultV3<T>
39+
: T extends z4.SomeType
40+
? ZodDefaultV4<T>
41+
: never;
42+
export type InteropZodOptional<T = InteropZodObjectShape> =
43+
T extends z3.ZodTypeAny
44+
? ZodOptionalV3<T>
45+
: T extends z4.SomeType
46+
? ZodOptionalV4<T>
47+
: never;
3148

3249
export type InteropZodObjectShape<
3350
T extends InteropZodObject = InteropZodObject
@@ -39,6 +56,18 @@ export type InteropZodObjectShape<
3956

4057
export type InteropZodIssue = z3.ZodIssue | z4.$ZodIssue;
4158

59+
export type InteropZodInput<T> = T extends z3.ZodType<
60+
unknown,
61+
z3.ZodTypeDef,
62+
infer Input
63+
>
64+
? Input
65+
: T extends z4.$ZodType<unknown, infer Input>
66+
? Input
67+
: T extends { _zod: { input: infer Input } }
68+
? Input
69+
: never;
70+
4271
// Simplified type inference to avoid circular dependencies
4372
export type InferInteropZodInput<T> = T extends z3.ZodType<
4473
unknown,

libs/langchain/src/agents/middlewareAgent/ReactAgent.ts

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -662,9 +662,9 @@ export class ReactAgent<
662662
/**
663663
* Initialize middleware states if not already present in the input state.
664664
*/
665-
#initializeMiddlewareStates(
665+
async #initializeMiddlewareStates(
666666
state: InvokeStateParameter<TMiddleware>
667-
): InvokeStateParameter<TMiddleware> {
667+
): Promise<InvokeStateParameter<TMiddleware>> {
668668
if (
669669
!this.options.middleware ||
670670
this.options.middleware.length === 0 ||
@@ -674,7 +674,7 @@ export class ReactAgent<
674674
return state;
675675
}
676676

677-
const defaultStates = initializeMiddlewareStates(
677+
const defaultStates = await initializeMiddlewareStates(
678678
this.options.middleware,
679679
state
680680
);
@@ -736,15 +736,15 @@ export class ReactAgent<
736736
* console.log(result.structuredResponse.weather); // outputs: "It's sunny and 75°F."
737737
* ```
738738
*/
739-
invoke(
739+
async invoke(
740740
state: InvokeStateParameter<TMiddleware>,
741741
config?: InvokeConfiguration<
742742
InferContextInput<ContextSchema> &
743743
InferMiddlewareContextInputs<TMiddleware>
744744
>
745745
) {
746746
type FullState = MergedAgentState<StructuredResponseFormat, TMiddleware>;
747-
const initializedState = this.#initializeMiddlewareStates(state);
747+
const initializedState = await this.#initializeMiddlewareStates(state);
748748
return this.#graph.invoke(
749749
initializedState,
750750
config as unknown as InferContextInput<ContextSchema> &
@@ -808,7 +808,7 @@ export class ReactAgent<
808808
InferMiddlewareContextInputs<TMiddleware>
809809
>
810810
): Promise<IterableReadableStream<any>> {
811-
const initializedState = this.#initializeMiddlewareStates(state);
811+
const initializedState = await this.#initializeMiddlewareStates(state);
812812
return this.#graph.streamEvents(initializedState, {
813813
...config,
814814
version: "v2",

libs/langchain/src/agents/middlewareAgent/middleware.ts

Lines changed: 84 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,10 @@
11
/* eslint-disable @typescript-eslint/no-explicit-any */
2-
import { z } from "zod/v3";
2+
import type {
3+
InteropZodObject,
4+
InteropZodDefault,
5+
InteropZodOptional,
6+
InferInteropZodInput,
7+
} from "@langchain/core/utils/types";
38
import type {
49
AgentMiddleware,
510
Runtime,
@@ -40,11 +45,11 @@ import type {
4045
* ```
4146
*/
4247
export function createMiddleware<
43-
TSchema extends z.ZodObject<any> | undefined = undefined,
48+
TSchema extends InteropZodObject | undefined = undefined,
4449
TContextSchema extends
45-
| z.ZodObject<any>
46-
| z.ZodOptional<z.ZodObject<any>>
47-
| z.ZodDefault<z.ZodObject<any>>
50+
| InteropZodObject
51+
| InteropZodOptional
52+
| InteropZodDefault
4853
| undefined = undefined
4954
>(config: {
5055
/**
@@ -90,18 +95,20 @@ export function createMiddleware<
9095
* @returns The modified model request or undefined to pass through
9196
*/
9297
modifyModelRequest?: (
93-
request: ModelRequest,
94-
state: (TSchema extends z.ZodObject<any> ? z.infer<TSchema> : {}) &
98+
options: ModelRequest,
99+
state: (TSchema extends InteropZodObject
100+
? InferInteropZodInput<TSchema>
101+
: {}) &
95102
AgentBuiltInState,
96103
runtime: Runtime<
97-
(TSchema extends z.ZodObject<any> ? z.infer<TSchema> : {}) &
104+
(TSchema extends InteropZodObject ? InferInteropZodInput<TSchema> : {}) &
98105
AgentBuiltInState,
99-
TContextSchema extends z.ZodObject<any>
100-
? z.infer<TContextSchema>
101-
: TContextSchema extends z.ZodDefault<z.ZodObject<any>>
102-
? z.infer<TContextSchema>
103-
: TContextSchema extends z.ZodOptional<z.ZodObject<any>>
104-
? Partial<z.infer<TContextSchema>>
106+
TContextSchema extends InteropZodObject
107+
? InferInteropZodInput<TContextSchema>
108+
: TContextSchema extends InteropZodDefault
109+
? InferInteropZodInput<TContextSchema>
110+
: TContextSchema extends InteropZodOptional
111+
? Partial<InferInteropZodInput<TContextSchema>>
105112
: never
106113
>
107114
) => Promise<ModelRequest | void> | ModelRequest | void;
@@ -114,27 +121,35 @@ export function createMiddleware<
114121
* @returns The modified middleware state or undefined to pass through
115122
*/
116123
beforeModel?: (
117-
state: (TSchema extends z.ZodObject<any> ? z.infer<TSchema> : {}) &
124+
state: (TSchema extends InteropZodObject
125+
? InferInteropZodInput<TSchema>
126+
: {}) &
118127
AgentBuiltInState,
119128
runtime: Runtime<
120-
(TSchema extends z.ZodObject<any> ? z.infer<TSchema> : {}) &
129+
(TSchema extends InteropZodObject ? InferInteropZodInput<TSchema> : {}) &
121130
AgentBuiltInState,
122-
TContextSchema extends z.ZodObject<any>
123-
? z.infer<TContextSchema>
124-
: TContextSchema extends z.ZodDefault<z.ZodObject<any>>
125-
? z.infer<TContextSchema>
126-
: TContextSchema extends z.ZodOptional<z.ZodObject<any>>
127-
? Partial<z.infer<TContextSchema>>
131+
TContextSchema extends InteropZodObject
132+
? InferInteropZodInput<TContextSchema>
133+
: TContextSchema extends InteropZodDefault
134+
? InferInteropZodInput<TContextSchema>
135+
: TContextSchema extends InteropZodOptional
136+
? Partial<InferInteropZodInput<TContextSchema>>
128137
: never
129138
>
130139
) =>
131140
| Promise<
132141
MiddlewareResult<
133-
Partial<TSchema extends z.ZodObject<any> ? z.infer<TSchema> : {}>
142+
Partial<
143+
TSchema extends InteropZodObject
144+
? InferInteropZodInput<TSchema>
145+
: {}
146+
>
134147
>
135148
>
136149
| MiddlewareResult<
137-
Partial<TSchema extends z.ZodObject<any> ? z.infer<TSchema> : {}>
150+
Partial<
151+
TSchema extends InteropZodObject ? InferInteropZodInput<TSchema> : {}
152+
>
138153
>;
139154
/**
140155
* The function to run after the model call. This function is called after the model is invoked and before any tools are called.
@@ -145,27 +160,35 @@ export function createMiddleware<
145160
* @returns The modified middleware state or undefined to pass through
146161
*/
147162
afterModel?: (
148-
state: (TSchema extends z.ZodObject<any> ? z.infer<TSchema> : {}) &
163+
state: (TSchema extends InteropZodObject
164+
? InferInteropZodInput<TSchema>
165+
: {}) &
149166
AgentBuiltInState,
150167
runtime: Runtime<
151-
(TSchema extends z.ZodObject<any> ? z.infer<TSchema> : {}) &
168+
(TSchema extends InteropZodObject ? InferInteropZodInput<TSchema> : {}) &
152169
AgentBuiltInState,
153-
TContextSchema extends z.ZodObject<any>
154-
? z.infer<TContextSchema>
155-
: TContextSchema extends z.ZodDefault<z.ZodObject<any>>
156-
? z.infer<TContextSchema>
157-
: TContextSchema extends z.ZodOptional<z.ZodObject<any>>
158-
? Partial<z.infer<TContextSchema>>
170+
TContextSchema extends InteropZodObject
171+
? InferInteropZodInput<TContextSchema>
172+
: TContextSchema extends InteropZodDefault
173+
? InferInteropZodInput<TContextSchema>
174+
: TContextSchema extends InteropZodOptional
175+
? Partial<InferInteropZodInput<TContextSchema>>
159176
: never
160177
>
161178
) =>
162179
| Promise<
163180
MiddlewareResult<
164-
Partial<TSchema extends z.ZodObject<any> ? z.infer<TSchema> : {}>
181+
Partial<
182+
TSchema extends InteropZodObject
183+
? InferInteropZodInput<TSchema>
184+
: {}
185+
>
165186
>
166187
>
167188
| MiddlewareResult<
168-
Partial<TSchema extends z.ZodObject<any> ? z.infer<TSchema> : {}>
189+
Partial<
190+
TSchema extends InteropZodObject ? InferInteropZodInput<TSchema> : {}
191+
>
169192
>;
170193
}): AgentMiddleware<TSchema, TContextSchema, any> {
171194
const middleware: AgentMiddleware<TSchema, TContextSchema, any> = {
@@ -183,14 +206,16 @@ export function createMiddleware<
183206
options,
184207
state,
185208
runtime as Runtime<
186-
(TSchema extends z.ZodObject<any> ? z.infer<TSchema> : {}) &
209+
(TSchema extends InteropZodObject
210+
? InferInteropZodInput<TSchema>
211+
: {}) &
187212
AgentBuiltInState,
188-
TContextSchema extends z.ZodObject<any>
189-
? z.infer<TContextSchema>
190-
: TContextSchema extends z.ZodDefault<z.ZodObject<any>>
191-
? z.infer<TContextSchema>
192-
: TContextSchema extends z.ZodOptional<z.ZodObject<any>>
193-
? Partial<z.infer<TContextSchema>>
213+
TContextSchema extends InteropZodObject
214+
? InferInteropZodInput<TContextSchema>
215+
: TContextSchema extends InteropZodDefault
216+
? InferInteropZodInput<TContextSchema>
217+
: TContextSchema extends InteropZodOptional
218+
? Partial<InferInteropZodInput<TContextSchema>>
194219
: never
195220
>
196221
)
@@ -203,14 +228,16 @@ export function createMiddleware<
203228
config.beforeModel!(
204229
state,
205230
runtime as Runtime<
206-
(TSchema extends z.ZodObject<any> ? z.infer<TSchema> : {}) &
231+
(TSchema extends InteropZodObject
232+
? InferInteropZodInput<TSchema>
233+
: {}) &
207234
AgentBuiltInState,
208-
TContextSchema extends z.ZodObject<any>
209-
? z.infer<TContextSchema>
210-
: TContextSchema extends z.ZodDefault<z.ZodObject<any>>
211-
? z.infer<TContextSchema>
212-
: TContextSchema extends z.ZodOptional<z.ZodObject<any>>
213-
? Partial<z.infer<TContextSchema>>
235+
TContextSchema extends InteropZodObject
236+
? InferInteropZodInput<TContextSchema>
237+
: TContextSchema extends InteropZodDefault
238+
? InferInteropZodInput<TContextSchema>
239+
: TContextSchema extends InteropZodOptional
240+
? Partial<InferInteropZodInput<TContextSchema>>
214241
: never
215242
>
216243
)
@@ -223,14 +250,16 @@ export function createMiddleware<
223250
config.afterModel!(
224251
state,
225252
runtime as Runtime<
226-
(TSchema extends z.ZodObject<any> ? z.infer<TSchema> : {}) &
253+
(TSchema extends InteropZodObject
254+
? InferInteropZodInput<TSchema>
255+
: {}) &
227256
AgentBuiltInState,
228-
TContextSchema extends z.ZodObject<any>
229-
? z.infer<TContextSchema>
230-
: TContextSchema extends z.ZodDefault<z.ZodObject<any>>
231-
? z.infer<TContextSchema>
232-
: TContextSchema extends z.ZodOptional<z.ZodObject<any>>
233-
? Partial<z.infer<TContextSchema>>
257+
TContextSchema extends InteropZodObject
258+
? InferInteropZodInput<TContextSchema>
259+
: TContextSchema extends InteropZodDefault
260+
? InferInteropZodInput<TContextSchema>
261+
: TContextSchema extends InteropZodOptional
262+
? Partial<InferInteropZodInput<TContextSchema>>
234263
: never
235264
>
236265
)

libs/langchain/src/agents/middlewareAgent/middleware/dynamicSystemPrompt.ts

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,11 @@ import { SystemMessage } from "@langchain/core/messages";
22
import { createMiddleware } from "../middleware.js";
33
import type { Runtime, AgentBuiltInState } from "../types.js";
44

5+
export type DynamicSystemPromptMiddlewareConfig<TContextSchema> = (
6+
state: AgentBuiltInState,
7+
runtime: Runtime<AgentBuiltInState, TContextSchema>
8+
) => string | Promise<string>;
9+
510
/**
611
* Dynamic System Prompt Middleware
712
*
@@ -41,10 +46,7 @@ import type { Runtime, AgentBuiltInState } from "../types.js";
4146
* @public
4247
*/
4348
export function dynamicSystemPromptMiddleware<TContextSchema = unknown>(
44-
fn: (
45-
state: AgentBuiltInState,
46-
runtime: Runtime<AgentBuiltInState, TContextSchema>
47-
) => string | Promise<string>
49+
fn: DynamicSystemPromptMiddlewareConfig<TContextSchema>
4850
) {
4951
return createMiddleware({
5052
name: "DynamicSystemPromptMiddleware",

libs/langchain/src/agents/middlewareAgent/middleware/hitl.ts

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,14 @@
11
/* eslint-disable @typescript-eslint/no-explicit-any */
22
import { z } from "zod/v3";
33
import { AIMessage, ToolMessage } from "@langchain/core/messages";
4+
import {
5+
InferInteropZodInput,
6+
InteropZodObject,
7+
} from "@langchain/core/utils/types";
48
import { interrupt } from "@langchain/langgraph";
59

610
import { createMiddleware } from "../middleware.js";
11+
import type { AgentMiddleware } from "../types.js";
712

813
const ToolConfigSchema = z.object({
914
/**
@@ -154,6 +159,9 @@ const contextSchema = z
154159
descriptionPrefix: z.string().default("Tool execution requires approval"),
155160
})
156161
.optional();
162+
export type HumanInTheLoopMiddlewareConfig = InferInteropZodInput<
163+
typeof contextSchema
164+
>;
157165

158166
/**
159167
* Creates a Human-in-the-Loop (HITL) middleware for tool approval and oversight.
@@ -333,8 +341,12 @@ const contextSchema = z
333341
* @public
334342
*/
335343
export function humanInTheLoopMiddleware(
336-
options: z.input<typeof contextSchema> = {}
337-
) {
344+
options: HumanInTheLoopMiddlewareConfig = {}
345+
): AgentMiddleware<
346+
undefined,
347+
InteropZodObject,
348+
HumanInTheLoopMiddlewareConfig
349+
> {
338350
return createMiddleware({
339351
name: "HumanInTheLoopMiddleware",
340352
contextSchema,

0 commit comments

Comments
 (0)