Skip to content

Commit e5e5fc8

Browse files
yaonyanyaonyan
committed
feat(ai): update step handling in stream processing to support dynamic maxSteps configuration
1 parent 61d3754 commit e5e5fc8

File tree

5 files changed

+44
-43
lines changed

5 files changed

+44
-43
lines changed

‎packages/ai/core/generate-text/__snapshots__/stream-text.test.ts.snap

Lines changed: 27 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result > callbac
127127
"type": "tool-result",
128128
},
129129
],
130-
"id": "msg-3",
130+
"id": "msg-1",
131131
"role": "tool",
132132
},
133133
{
@@ -137,7 +137,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result > callbac
137137
"type": "text",
138138
},
139139
],
140-
"id": "msg-2",
140+
"id": "msg-3",
141141
"role": "assistant",
142142
},
143143
],
@@ -194,7 +194,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result > callbac
194194
"type": "tool-result",
195195
},
196196
],
197-
"id": "msg-3",
197+
"id": "msg-1",
198198
"role": "tool",
199199
},
200200
],
@@ -275,7 +275,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result > callbac
275275
"type": "tool-result",
276276
},
277277
],
278-
"id": "msg-3",
278+
"id": "msg-1",
279279
"role": "tool",
280280
},
281281
{
@@ -285,7 +285,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result > callbac
285285
"type": "text",
286286
},
287287
],
288-
"id": "msg-2",
288+
"id": "msg-3",
289289
"role": "assistant",
290290
},
291291
],
@@ -367,7 +367,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result > callbac
367367
"type": "tool-result",
368368
},
369369
],
370-
"id": "msg-3",
370+
"id": "msg-1",
371371
"role": "tool",
372372
},
373373
],
@@ -448,7 +448,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result > callbac
448448
"type": "tool-result",
449449
},
450450
],
451-
"id": "msg-3",
451+
"id": "msg-1",
452452
"role": "tool",
453453
},
454454
{
@@ -458,7 +458,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result > callbac
458458
"type": "text",
459459
},
460460
],
461-
"id": "msg-2",
461+
"id": "msg-3",
462462
"role": "assistant",
463463
},
464464
],
@@ -534,7 +534,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result > should
534534
"warnings": undefined,
535535
},
536536
{
537-
"messageId": "msg-2",
537+
"messageId": "msg-3",
538538
"request": {},
539539
"type": "step-start",
540540
"warnings": [],
@@ -552,7 +552,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result > should
552552
"finishReason": "stop",
553553
"isContinued": false,
554554
"logprobs": undefined,
555-
"messageId": "msg-2",
555+
"messageId": "msg-3",
556556
"providerMetadata": undefined,
557557
"request": {},
558558
"response": {
@@ -752,7 +752,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result > value p
752752
"type": "tool-result",
753753
},
754754
],
755-
"id": "msg-3",
755+
"id": "msg-1",
756756
"role": "tool",
757757
},
758758
{
@@ -762,7 +762,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result > value p
762762
"type": "text",
763763
},
764764
],
765-
"id": "msg-2",
765+
"id": "msg-3",
766766
"role": "assistant",
767767
},
768768
]
@@ -818,7 +818,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result > value p
818818
"type": "tool-result",
819819
},
820820
],
821-
"id": "msg-3",
821+
"id": "msg-1",
822822
"role": "tool",
823823
},
824824
],
@@ -899,7 +899,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result > value p
899899
"type": "tool-result",
900900
},
901901
],
902-
"id": "msg-3",
902+
"id": "msg-1",
903903
"role": "tool",
904904
},
905905
{
@@ -909,7 +909,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result > value p
909909
"type": "text",
910910
},
911911
],
912-
"id": "msg-2",
912+
"id": "msg-3",
913913
"role": "assistant",
914914
},
915915
],
@@ -968,7 +968,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result with expe
968968
"type": "tool-result",
969969
},
970970
],
971-
"id": "msg-3",
971+
"id": "msg-1",
972972
"role": "tool",
973973
},
974974
{
@@ -978,7 +978,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result with expe
978978
"type": "text",
979979
},
980980
],
981-
"id": "msg-2",
981+
"id": "msg-3",
982982
"role": "assistant",
983983
},
984984
],
@@ -1024,7 +1024,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result with expe
10241024
"type": "tool-result",
10251025
},
10261026
],
1027-
"id": "msg-3",
1027+
"id": "msg-1",
10281028
"role": "tool",
10291029
},
10301030
],
@@ -1099,7 +1099,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result with expe
10991099
"type": "tool-result",
11001100
},
11011101
],
1102-
"id": "msg-3",
1102+
"id": "msg-1",
11031103
"role": "tool",
11041104
},
11051105
{
@@ -1109,7 +1109,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result with expe
11091109
"type": "text",
11101110
},
11111111
],
1112-
"id": "msg-2",
1112+
"id": "msg-3",
11131113
"role": "assistant",
11141114
},
11151115
],
@@ -1191,7 +1191,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result with expe
11911191
"warnings": undefined,
11921192
},
11931193
{
1194-
"messageId": "msg-2",
1194+
"messageId": "msg-3",
11951195
"request": {},
11961196
"type": "step-start",
11971197
"warnings": [],
@@ -1213,7 +1213,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result with expe
12131213
"finishReason": "stop",
12141214
"isContinued": false,
12151215
"logprobs": undefined,
1216-
"messageId": "msg-2",
1216+
"messageId": "msg-3",
12171217
"providerMetadata": undefined,
12181218
"request": {},
12191219
"response": {
@@ -1276,7 +1276,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result with expe
12761276
"type": "tool-result",
12771277
},
12781278
],
1279-
"id": "msg-3",
1279+
"id": "msg-1",
12801280
"role": "tool",
12811281
},
12821282
{
@@ -1286,7 +1286,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result with expe
12861286
"type": "text",
12871287
},
12881288
],
1289-
"id": "msg-2",
1289+
"id": "msg-3",
12901290
"role": "assistant",
12911291
},
12921292
]
@@ -1331,7 +1331,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result with expe
13311331
"type": "tool-result",
13321332
},
13331333
],
1334-
"id": "msg-3",
1334+
"id": "msg-1",
13351335
"role": "tool",
13361336
},
13371337
],
@@ -1406,7 +1406,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result with expe
14061406
"type": "tool-result",
14071407
},
14081408
],
1409-
"id": "msg-3",
1409+
"id": "msg-1",
14101410
"role": "tool",
14111411
},
14121412
{
@@ -1416,7 +1416,7 @@ exports[`streamText > options.maxSteps > 2 steps: initial, tool-result with expe
14161416
"type": "text",
14171417
},
14181418
],
1419-
"id": "msg-2",
1419+
"id": "msg-3",
14201420
"role": "assistant",
14211421
},
14221422
],

‎packages/ai/core/generate-text/generate-text.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -331,6 +331,7 @@ A function that attempts to repair a tool call that failed to parse.
331331
const prepareStepResult = await prepareStep?.({
332332
model,
333333
steps,
334+
maxSteps,
334335
stepNumber: stepCount,
335336
messages: promptMessages,
336337
});

‎packages/ai/core/generate-text/prepare-step.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ export type PrepareStepFunction<
2121
> = (options: {
2222
steps: Array<StepResult<NoInfer<TOOLS>>>;
2323
stepNumber: number;
24+
maxSteps: number;
2425
model: LanguageModel;
2526
messages: Array<LanguageModelV1Message>;
2627
}) => PromiseLike<PrepareStepResult<TOOLS>> | PrepareStepResult<TOOLS>;

‎packages/ai/core/generate-text/stream-text.test.ts

Lines changed: 4 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -2982,21 +2982,8 @@ describe('streamText', () => {
29822982
case 1: {
29832983
expect(mode).toStrictEqual({
29842984
type: 'regular',
2985-
toolChoice: { type: 'tool', toolName: 'tool1' },
2986-
tools: [
2987-
{
2988-
type: 'function',
2989-
name: 'tool1',
2990-
description: undefined,
2991-
parameters: {
2992-
$schema: 'http://json-schema.org/draft-07/schema#',
2993-
additionalProperties: false,
2994-
properties: { value: { type: 'string' } },
2995-
required: ['value'],
2996-
type: 'object',
2997-
},
2998-
},
2999-
],
2985+
toolChoice: { type: 'auto' },
2986+
tools: [],
30002987
});
30012988

30022989
expect(prompt).toStrictEqual([
@@ -3092,6 +3079,7 @@ describe('streamText', () => {
30923079
onStepFinishResults.push(event);
30933080
},
30943081
experimental_prepareStep: async ({ model, stepNumber, steps }) => {
3082+
console.log('Preparing step:', { model, stepNumber, steps });
30953083
expect(model).toStrictEqual(originalModel);
30963084

30973085
if (stepNumber === 0) {
@@ -3106,7 +3094,7 @@ describe('streamText', () => {
31063094
}
31073095

31083096
if (stepNumber === 1) {
3109-
expect(steps.length).toStrictEqual(0); // step 0 not yet recorded in streaming
3097+
expect(steps.length).toStrictEqual(1); // step 0 is now properly recorded due to race condition fix
31103098
return {
31113099
model: trueModel,
31123100
toolChoice: 'auto',

‎packages/ai/core/generate-text/stream-text.ts

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -647,6 +647,7 @@ class DefaultStreamTextResult<TOOLS extends ToolSet, OUTPUT, PARTIAL_OUTPUT>
647647
let stepType: 'initial' | 'continue' | 'tool-result' = 'initial';
648648
const recordedSteps: StepResult<TOOLS>[] = [];
649649
let rootSpan!: Span;
650+
let stepFinish!: DelayedPromise<void>;
650651

651652
const eventProcessor = new TransformStream<
652653
EnrichedStreamPart<TOOLS, PARTIAL_OUTPUT>,
@@ -798,6 +799,9 @@ class DefaultStreamTextResult<TOOLS extends ToolSet, OUTPUT, PARTIAL_OUTPUT>
798799
recordedResponse.messages.push(...stepMessages);
799800
recordedContinuationText = '';
800801
}
802+
803+
// Signal that the step is fully processed
804+
stepFinish.resolve();
801805
}
802806

803807
if (part.type === 'finish') {
@@ -983,6 +987,9 @@ class DefaultStreamTextResult<TOOLS extends ToolSet, OUTPUT, PARTIAL_OUTPUT>
983987
hasLeadingWhitespace: boolean;
984988
messageId: string;
985989
}) {
990+
// Create a new promise for this step
991+
stepFinish = new DelayedPromise<void>();
992+
986993
// after the 1st step, we need to switch to messages format:
987994
const promptFormat =
988995
responseMessages.length === 0 ? initialPrompt.type : 'messages';
@@ -996,6 +1003,7 @@ class DefaultStreamTextResult<TOOLS extends ToolSet, OUTPUT, PARTIAL_OUTPUT>
9961003
model,
9971004
steps: recordedSteps,
9981005
stepNumber: recordedSteps.length,
1006+
maxSteps,
9991007
messages: stepInputMessages
10001008
});
10011009

@@ -1453,6 +1461,9 @@ class DefaultStreamTextResult<TOOLS extends ToolSet, OUTPUT, PARTIAL_OUTPUT>
14531461

14541462
self.closeStream(); // close the stitchable stream
14551463
} else {
1464+
// wait for the step to be fully processed by the event processor
1465+
await stepFinish.value;
1466+
14561467
// append to messages for the next step:
14571468
if (stepType === 'continue') {
14581469
// continue step: update the last assistant message

0 commit comments

Comments
 (0)