Skip to content

Commit 3a167a5

Browse files
monadoidmonadoid
committed
fix: keep stg-1555 on ai sdk v2 surface
1 parent f8c727d commit 3a167a5

18 files changed

Lines changed: 165 additions & 416 deletions

File tree

‎packages/core/lib/v3/agent/tools/ariaTree.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ export const ariaTreeTool = (v3: V3, toolTimeout?: number) =>
4747
};
4848
}
4949
},
50-
toModelOutput: ({ output: result }) => {
50+
toModelOutput: (result) => {
5151
if (result.success === false || result.error !== undefined) {
5252
return {
5353
type: "content",

‎packages/core/lib/v3/agent/tools/click.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ export const clickTool = (v3: V3, provider?: string) =>
8686
};
8787
}
8888
},
89-
toModelOutput: ({ output: result }) => {
89+
toModelOutput: (result) => {
9090
if (result.success === false || result.error !== undefined) {
9191
return {
9292
type: "content",

‎packages/core/lib/v3/agent/tools/dragAndDrop.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ export const dragAndDropTool = (v3: V3, provider?: string) =>
101101
};
102102
}
103103
},
104-
toModelOutput: ({ output: result }) => {
104+
toModelOutput: (result) => {
105105
if (result.success === false || result.error !== undefined) {
106106
return {
107107
type: "content",

‎packages/core/lib/v3/agent/tools/fillFormVision.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,7 @@ MANDATORY USE CASES (always use fillFormVision for these):
150150
};
151151
}
152152
},
153-
toModelOutput: ({ output: result }) => {
153+
toModelOutput: (result) => {
154154
if (result.success === false || result.error !== undefined) {
155155
return {
156156
type: "content",

‎packages/core/lib/v3/agent/tools/screenshot.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ export const screenshotTool = (v3: V3) =>
3030
};
3131
}
3232
},
33-
toModelOutput: ({ output: result }) => {
33+
toModelOutput: (result) => {
3434
if (result.success === false || result.error !== undefined) {
3535
return {
3636
type: "content",

‎packages/core/lib/v3/agent/tools/scroll.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ export const scrollTool = (v3: V3) =>
6464
scrolledPixels: scrollDistance,
6565
};
6666
},
67-
toModelOutput: ({ output: result }) => {
67+
toModelOutput: (result) => {
6868
if (result.success === false || result.error !== undefined) {
6969
return {
7070
type: "content",
@@ -169,7 +169,7 @@ export const scrollVisionTool = (v3: V3, provider?: string) =>
169169
screenshotBase64,
170170
};
171171
},
172-
toModelOutput: ({ output: result }) => {
172+
toModelOutput: (result) => {
173173
if (result.success === false || result.error !== undefined) {
174174
return {
175175
type: "content",

‎packages/core/lib/v3/agent/tools/type.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ export const typeTool = (v3: V3, provider?: string, variables?: Variables) => {
103103
};
104104
}
105105
},
106-
toModelOutput: ({ output: result }) => {
106+
toModelOutput: (result) => {
107107
if (result.success === false || result.error !== undefined) {
108108
return {
109109
type: "content",

‎packages/core/lib/v3/agent/tools/wait.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ export const waitTool = (v3: V3, mode?: AgentToolMode) =>
4040

4141
return { success: true, waited: timeMs };
4242
},
43-
toModelOutput: ({ output: result }) => {
43+
toModelOutput: (result) => {
4444
if (result.success === false || result.error !== undefined) {
4545
return {
4646
type: "content",

‎packages/core/lib/v3/external_clients/aisdk.ts‎

Lines changed: 47 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -1,64 +1,24 @@
11
import {
2-
AssistantModelMessage,
2+
CoreAssistantMessage,
3+
ModelMessage,
4+
CoreSystemMessage,
5+
Tool,
6+
CoreUserMessage,
7+
generateObject,
38
generateText,
49
ImagePart,
5-
ModelMessage,
6-
Output,
7-
SystemModelMessage,
810
TextPart,
9-
UserModelMessage,
10-
type Tool,
1111
} from "ai";
12-
import type { LanguageModelV2, LanguageModelV3 } from "@ai-sdk/provider";
12+
import type { LanguageModelV2 } from "@ai-sdk/provider";
1313
import { CreateChatCompletionOptions, LLMClient } from "../llm/LLMClient.js";
1414
import { AvailableModel } from "../types/public/index.js";
1515
import { ChatCompletion } from "openai/resources";
1616

17-
function getReasoningTokens(
18-
usage?: {
19-
outputTokenDetails?: { reasoningTokens?: number };
20-
reasoningTokens?: number;
21-
} | null,
22-
): number {
23-
return (
24-
usage?.outputTokenDetails?.reasoningTokens ?? usage?.reasoningTokens ?? 0
25-
);
26-
}
27-
28-
function getCachedInputTokens(
29-
usage?: {
30-
inputTokenDetails?: { cacheReadTokens?: number };
31-
cachedInputTokens?: number;
32-
} | null,
33-
): number {
34-
return (
35-
usage?.inputTokenDetails?.cacheReadTokens ?? usage?.cachedInputTokens ?? 0
36-
);
37-
}
38-
39-
function toLLMUsage(usage?: {
40-
inputTokens?: number;
41-
outputTokens?: number;
42-
totalTokens?: number;
43-
outputTokenDetails?: { reasoningTokens?: number };
44-
reasoningTokens?: number;
45-
inputTokenDetails?: { cacheReadTokens?: number };
46-
cachedInputTokens?: number;
47-
}) {
48-
return {
49-
prompt_tokens: usage?.inputTokens ?? 0,
50-
completion_tokens: usage?.outputTokens ?? 0,
51-
reasoning_tokens: getReasoningTokens(usage),
52-
cached_input_tokens: getCachedInputTokens(usage),
53-
total_tokens: usage?.totalTokens ?? 0,
54-
};
55-
}
56-
5717
export class AISdkClient extends LLMClient {
5818
public type = "aisdk" as const;
59-
private model: LanguageModelV2 | LanguageModelV3;
19+
private model: LanguageModelV2;
6020

61-
constructor({ model }: { model: LanguageModelV2 | LanguageModelV3 }) {
21+
constructor({ model }: { model: LanguageModelV2 }) {
6222
super(model.modelId as AvailableModel);
6323
this.model = model;
6424
}
@@ -70,7 +30,7 @@ export class AISdkClient extends LLMClient {
7030
(message) => {
7131
if (Array.isArray(message.content)) {
7232
if (message.role === "system") {
73-
const systemMessage: SystemModelMessage = {
33+
const systemMessage: CoreSystemMessage = {
7434
role: "system",
7535
content: message.content
7636
.map((c) => ("text" in c ? c.text : ""))
@@ -86,32 +46,32 @@ export class AISdkClient extends LLMClient {
8646
image: content.image_url.url,
8747
};
8848
return imageContent;
49+
} else {
50+
const textContent: TextPart = {
51+
type: "text",
52+
text: content.text,
53+
};
54+
return textContent;
8955
}
90-
91-
const textContent: TextPart = {
92-
type: "text",
93-
text: content.text,
94-
};
95-
return textContent;
9656
});
9757

9858
if (message.role === "user") {
99-
const userMessage: UserModelMessage = {
59+
const userMessage: CoreUserMessage = {
10060
role: "user",
10161
content: contentParts,
10262
};
10363
return userMessage;
64+
} else {
65+
const textOnlyParts = contentParts.map((part) => ({
66+
type: "text" as const,
67+
text: part.type === "image" ? "[Image]" : part.text,
68+
}));
69+
const assistantMessage: CoreAssistantMessage = {
70+
role: "assistant",
71+
content: textOnlyParts,
72+
};
73+
return assistantMessage;
10474
}
105-
106-
const textOnlyParts = contentParts.map((part) => ({
107-
type: "text" as const,
108-
text: part.type === "image" ? "[Image]" : part.text,
109-
}));
110-
const assistantMessage: AssistantModelMessage = {
111-
role: "assistant",
112-
content: textOnlyParts,
113-
};
114-
return assistantMessage;
11575
}
11676

11777
return {
@@ -122,28 +82,27 @@ export class AISdkClient extends LLMClient {
12282
);
12383

12484
if (options.response_model) {
125-
const response = await generateText({
85+
const response = await generateObject({
12686
model: this.model,
12787
messages: formattedMessages,
128-
output: Output.object({
129-
schema: options.response_model.schema,
130-
name: options.response_model.name,
131-
}),
132-
maxOutputTokens: options.maxOutputTokens,
133-
temperature: options.temperature,
134-
topP: options.top_p,
135-
frequencyPenalty: options.frequency_penalty,
136-
presencePenalty: options.presence_penalty,
88+
schema: options.response_model.schema,
13789
});
13890

13991
return {
140-
data: response.output,
141-
usage: toLLMUsage(response.usage),
92+
data: response.object,
93+
usage: {
94+
prompt_tokens: response.usage.inputTokens ?? 0,
95+
completion_tokens: response.usage.outputTokens ?? 0,
96+
reasoning_tokens: response.usage.reasoningTokens ?? 0,
97+
cached_input_tokens: response.usage.cachedInputTokens ?? 0,
98+
total_tokens: response.usage.totalTokens ?? 0,
99+
},
142100
} as T;
143101
}
144102

145103
const tools: Record<string, Tool> = {};
146-
for (const rawTool of options.tools ?? []) {
104+
105+
for (const rawTool of options.tools) {
147106
tools[rawTool.name] = {
148107
description: rawTool.description,
149108
inputSchema: rawTool.parameters,
@@ -153,25 +112,18 @@ export class AISdkClient extends LLMClient {
153112
const response = await generateText({
154113
model: this.model,
155114
messages: formattedMessages,
156-
tools: Object.keys(tools).length > 0 ? tools : undefined,
157-
toolChoice:
158-
Object.keys(tools).length > 0
159-
? options.tool_choice === "required"
160-
? "required"
161-
: options.tool_choice === "none"
162-
? "none"
163-
: "auto"
164-
: undefined,
165-
maxOutputTokens: options.maxOutputTokens,
166-
temperature: options.temperature,
167-
topP: options.top_p,
168-
frequencyPenalty: options.frequency_penalty,
169-
presencePenalty: options.presence_penalty,
115+
tools,
170116
});
171117

172118
return {
173119
data: response.text,
174-
usage: toLLMUsage(response.usage),
120+
usage: {
121+
prompt_tokens: response.usage.inputTokens ?? 0,
122+
completion_tokens: response.usage.outputTokens ?? 0,
123+
reasoning_tokens: response.usage.reasoningTokens ?? 0,
124+
cached_input_tokens: response.usage.cachedInputTokens ?? 0,
125+
total_tokens: response.usage.totalTokens ?? 0,
126+
},
175127
} as T;
176128
}
177129
}

‎packages/core/lib/v3/flowlogger/FlowLogger.ts‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -674,9 +674,8 @@ export class FlowLogger {
674674
// while leaving model execution behavior unchanged.
675675
static createLlmLoggingMiddleware(
676676
modelId: string,
677-
): Pick<LanguageModelMiddleware, "specificationVersion" | "wrapGenerate"> {
677+
): Pick<LanguageModelMiddleware, "wrapGenerate"> {
678678
return {
679-
specificationVersion: "v3",
680679
wrapGenerate: async ({ doGenerate, params }) => {
681680
const llmRequestId = uuidv7();
682681
FlowLogger.logLlmRequest({
@@ -697,8 +696,8 @@ export class FlowLogger {
697696
requestId: llmRequestId,
698697
model: modelId,
699698
output: FlowLogger.buildMiddlewareOutputSummary(res),
700-
inputTokens: result.usage?.inputTokens?.total,
701-
outputTokens: result.usage?.outputTokens?.total,
699+
inputTokens: result.usage?.inputTokens,
700+
outputTokens: result.usage?.outputTokens,
702701
});
703702

704703
return result;

0 commit comments

Comments
 (0)