feat: upgrade ai package to version six and the corresponding @ai-sdk/* packages to compatible versions (#18172)
Used the migration guide to carry out this upgrade: https://ai-sdk.dev/docs/migration-guides/migration-guide-6-0 I have not been able to test locally due to credits. <img width="220" height="450" alt="image" src="https://github.com/user-attachments/assets/050b34b9-3239-4010-8c47-b43d44571994" /> --------- Co-authored-by: Félix Malfait <felix.malfait@gmail.com>
This commit is contained in:
Abdullah.
+16
-15
@@ -2,9 +2,9 @@ import { Injectable, Logger } from '@nestjs/common';
|
||||
import { InjectRepository } from '@nestjs/typeorm';
|
||||
|
||||
import {
|
||||
generateObject,
|
||||
generateText,
|
||||
jsonSchema,
|
||||
Output,
|
||||
stepCountIs,
|
||||
type ToolSet,
|
||||
} from 'ai';
|
||||
@@ -20,6 +20,7 @@ import { ToolCategory } from 'src/engine/core-modules/tool-provider/enums/tool-c
|
||||
import { ToolRegistryService } from 'src/engine/core-modules/tool-provider/services/tool-registry.service';
|
||||
import { type AgentExecutionResult } from 'src/engine/metadata-modules/ai/ai-agent-execution/types/agent-execution-result.type';
|
||||
import { extractCacheCreationTokensFromSteps } from 'src/engine/metadata-modules/ai/ai-billing/utils/extract-cache-creation-tokens.util';
|
||||
import { mergeLanguageModelUsage } from 'src/engine/metadata-modules/ai/ai-billing/utils/merge-language-model-usage.util';
|
||||
import {
|
||||
AgentException,
|
||||
AgentExceptionCode,
|
||||
@@ -219,7 +220,7 @@ export class AgentAsyncExecutorService {
|
||||
};
|
||||
}
|
||||
|
||||
const output = await generateObject({
|
||||
const structuredResult = await generateText({
|
||||
system: WORKFLOW_SYSTEM_PROMPTS.OUTPUT_GENERATOR,
|
||||
model: registeredModel.model,
|
||||
prompt: `Based on the following execution results, generate the structured output according to the schema:
|
||||
@@ -227,23 +228,23 @@ export class AgentAsyncExecutorService {
|
||||
Execution Results: ${textResponse.text}
|
||||
|
||||
Please generate the structured output based on the execution results and context above.`,
|
||||
schema: jsonSchema(agentSchema),
|
||||
output: Output.object({ schema: jsonSchema(agentSchema) }),
|
||||
experimental_telemetry: AI_TELEMETRY_CONFIG,
|
||||
});
|
||||
|
||||
if (structuredResult.output == null) {
|
||||
throw new AgentException(
|
||||
'Failed to generate structured output from execution results',
|
||||
AgentExceptionCode.AGENT_EXECUTION_FAILED,
|
||||
);
|
||||
}
|
||||
|
||||
return {
|
||||
result: output.object as object,
|
||||
usage: {
|
||||
inputTokens:
|
||||
(textResponse.usage?.inputTokens ?? 0) +
|
||||
(output.usage?.inputTokens ?? 0),
|
||||
outputTokens:
|
||||
(textResponse.usage?.outputTokens ?? 0) +
|
||||
(output.usage?.outputTokens ?? 0),
|
||||
totalTokens:
|
||||
(textResponse.usage?.totalTokens ?? 0) +
|
||||
(output.usage?.totalTokens ?? 0),
|
||||
},
|
||||
result: structuredResult.output as object,
|
||||
usage: mergeLanguageModelUsage(
|
||||
textResponse.usage,
|
||||
structuredResult.usage,
|
||||
),
|
||||
cacheCreationTokens,
|
||||
};
|
||||
} catch (error) {
|
||||
|
||||
+7
-3
@@ -1,4 +1,4 @@
|
||||
import { generateObject, type LanguageModel, NoSuchToolError } from 'ai';
|
||||
import { type LanguageModel, NoSuchToolError, Output, generateText } from 'ai';
|
||||
import { type z } from 'zod';
|
||||
|
||||
import { AI_TELEMETRY_CONFIG } from 'src/engine/metadata-modules/ai/ai-models/constants/ai-telemetry.const';
|
||||
@@ -41,9 +41,9 @@ export const repairToolCall = async ({
|
||||
}
|
||||
|
||||
try {
|
||||
const { object: repairedInput } = await generateObject({
|
||||
const { output: repairedInput } = await generateText({
|
||||
model,
|
||||
schema: schema as z.ZodTypeAny,
|
||||
output: Output.object({ schema: schema as z.ZodTypeAny }),
|
||||
prompt: [
|
||||
`The AI model attempted to call the tool "${toolCall.toolName}" with invalid input.`,
|
||||
``,
|
||||
@@ -62,6 +62,10 @@ export const repairToolCall = async ({
|
||||
experimental_telemetry: AI_TELEMETRY_CONFIG,
|
||||
});
|
||||
|
||||
if (repairedInput == null) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return {
|
||||
type: 'tool-call',
|
||||
toolCallId: toolCall.toolCallId,
|
||||
|
||||
+77
-4
@@ -38,10 +38,20 @@ describe('AIBillingService', () => {
|
||||
cacheCreationCostPerMillionTokens: 3.75,
|
||||
};
|
||||
|
||||
const defaultTokenDetails = {
|
||||
inputTokenDetails: {
|
||||
noCacheTokens: 0,
|
||||
cacheReadTokens: 0,
|
||||
cacheWriteTokens: 0,
|
||||
},
|
||||
outputTokenDetails: { textTokens: 0, reasoningTokens: 0 },
|
||||
};
|
||||
|
||||
const mockTokenUsage = {
|
||||
inputTokens: 1000,
|
||||
outputTokens: 500,
|
||||
totalTokens: 1500,
|
||||
...defaultTokenDetails,
|
||||
};
|
||||
|
||||
beforeEach(async () => {
|
||||
@@ -92,7 +102,12 @@ describe('AIBillingService', () => {
|
||||
inputTokens: 1000,
|
||||
outputTokens: 500,
|
||||
totalTokens: 1500,
|
||||
cachedInputTokens: 600,
|
||||
inputTokenDetails: {
|
||||
noCacheTokens: 400,
|
||||
cacheReadTokens: 600,
|
||||
cacheWriteTokens: 0,
|
||||
},
|
||||
outputTokenDetails: { textTokens: 500, reasoningTokens: 0 },
|
||||
},
|
||||
});
|
||||
|
||||
@@ -118,7 +133,12 @@ describe('AIBillingService', () => {
|
||||
inputTokens: 400,
|
||||
outputTokens: 500,
|
||||
totalTokens: 900,
|
||||
cachedInputTokens: 600,
|
||||
inputTokenDetails: {
|
||||
noCacheTokens: 400,
|
||||
cacheReadTokens: 600,
|
||||
cacheWriteTokens: 0,
|
||||
},
|
||||
outputTokenDetails: { textTokens: 500, reasoningTokens: 0 },
|
||||
},
|
||||
cacheCreationTokens: 200,
|
||||
},
|
||||
@@ -139,7 +159,12 @@ describe('AIBillingService', () => {
|
||||
inputTokens: 1000,
|
||||
outputTokens: 500,
|
||||
totalTokens: 2000,
|
||||
reasoningTokens: 500,
|
||||
inputTokenDetails: {
|
||||
noCacheTokens: 0,
|
||||
cacheReadTokens: 0,
|
||||
cacheWriteTokens: 0,
|
||||
},
|
||||
outputTokenDetails: { textTokens: 0, reasoningTokens: 500 },
|
||||
},
|
||||
});
|
||||
|
||||
@@ -152,6 +177,42 @@ describe('AIBillingService', () => {
|
||||
expect(costInDollars).toBeCloseTo(0.0075);
|
||||
});
|
||||
|
||||
it('should use outputTokenDetails.reasoningTokens when present (SDK-aligned shape)', () => {
|
||||
const costInDollars = service.calculateCost('gpt-4o', {
|
||||
usage: {
|
||||
inputTokens: 1000,
|
||||
outputTokens: 500,
|
||||
totalTokens: 2000,
|
||||
inputTokenDetails: {
|
||||
noCacheTokens: 0,
|
||||
cacheReadTokens: 0,
|
||||
cacheWriteTokens: 0,
|
||||
},
|
||||
outputTokenDetails: { textTokens: 0, reasoningTokens: 500 },
|
||||
},
|
||||
});
|
||||
|
||||
expect(costInDollars).toBeCloseTo(0.0075);
|
||||
});
|
||||
|
||||
it('should use inputTokenDetails.cacheReadTokens when present (SDK-aligned shape)', () => {
|
||||
const costInDollars = service.calculateCost('gpt-4o', {
|
||||
usage: {
|
||||
inputTokens: 1000,
|
||||
outputTokens: 500,
|
||||
totalTokens: 1500,
|
||||
inputTokenDetails: {
|
||||
noCacheTokens: 400,
|
||||
cacheReadTokens: 600,
|
||||
cacheWriteTokens: 0,
|
||||
},
|
||||
outputTokenDetails: { textTokens: 500, reasoningTokens: 0 },
|
||||
},
|
||||
});
|
||||
|
||||
expect(costInDollars).toBeCloseTo(0.00675);
|
||||
});
|
||||
|
||||
it('should fall back to input rate when cachedInputCostPerMillionTokens is undefined', () => {
|
||||
mockAiModelRegistryService.getEffectiveModelConfig.mockReturnValue({
|
||||
...openaiModelConfig,
|
||||
@@ -163,7 +224,12 @@ describe('AIBillingService', () => {
|
||||
inputTokens: 1000,
|
||||
outputTokens: 500,
|
||||
totalTokens: 1500,
|
||||
cachedInputTokens: 600,
|
||||
inputTokenDetails: {
|
||||
noCacheTokens: 400,
|
||||
cacheReadTokens: 600,
|
||||
cacheWriteTokens: 0,
|
||||
},
|
||||
outputTokenDetails: { textTokens: 500, reasoningTokens: 0 },
|
||||
},
|
||||
});
|
||||
|
||||
@@ -202,6 +268,12 @@ describe('AIBillingService', () => {
|
||||
outputTokens: 1000,
|
||||
totalTokens: 251_000,
|
||||
cachedInputTokens: 100_000,
|
||||
inputTokenDetails: {
|
||||
noCacheTokens: 0,
|
||||
cacheReadTokens: 100_000,
|
||||
cacheWriteTokens: 0,
|
||||
},
|
||||
outputTokenDetails: { textTokens: 1000, reasoningTokens: 0 },
|
||||
},
|
||||
},
|
||||
);
|
||||
@@ -240,6 +312,7 @@ describe('AIBillingService', () => {
|
||||
inputTokens: 50_000,
|
||||
outputTokens: 1000,
|
||||
totalTokens: 51_000,
|
||||
...defaultTokenDetails,
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
+2
-2
@@ -37,8 +37,8 @@ export class AIBillingService {
|
||||
const breakdown = computeCostBreakdown(model, {
|
||||
inputTokens: usage.inputTokens,
|
||||
outputTokens: usage.outputTokens,
|
||||
reasoningTokens: usage.reasoningTokens,
|
||||
cachedInputTokens: usage.cachedInputTokens,
|
||||
reasoningTokens: usage.outputTokenDetails?.reasoningTokens,
|
||||
cachedInputTokens: usage.inputTokenDetails?.cacheReadTokens,
|
||||
cacheCreationTokens,
|
||||
});
|
||||
|
||||
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
import { type LanguageModelUsage } from 'ai';
|
||||
|
||||
const sum = (a: number | undefined, b: number | undefined): number =>
|
||||
(a ?? 0) + (b ?? 0);
|
||||
|
||||
export const mergeLanguageModelUsage = (
|
||||
a: LanguageModelUsage,
|
||||
b: LanguageModelUsage,
|
||||
): LanguageModelUsage => {
|
||||
const inA = a.inputTokenDetails;
|
||||
const inB = b.inputTokenDetails;
|
||||
const outA = a.outputTokenDetails;
|
||||
const outB = b.outputTokenDetails;
|
||||
|
||||
return {
|
||||
inputTokens: sum(a.inputTokens, b.inputTokens),
|
||||
outputTokens: sum(a.outputTokens, b.outputTokens),
|
||||
totalTokens: sum(a.totalTokens, b.totalTokens),
|
||||
inputTokenDetails: {
|
||||
noCacheTokens: sum(inA?.noCacheTokens, inB?.noCacheTokens),
|
||||
cacheReadTokens: sum(inA?.cacheReadTokens, inB?.cacheReadTokens),
|
||||
cacheWriteTokens: sum(inA?.cacheWriteTokens, inB?.cacheWriteTokens),
|
||||
},
|
||||
outputTokenDetails: {
|
||||
textTokens: sum(outA?.textTokens, outB?.textTokens),
|
||||
reasoningTokens: sum(outA?.reasoningTokens, outB?.reasoningTokens),
|
||||
},
|
||||
};
|
||||
};
|
||||
+6
-5
@@ -131,7 +131,8 @@ export class AgentChatStreamingService {
|
||||
messageMetadata: ({ part }) => {
|
||||
if (part.type === 'finish-step') {
|
||||
const stepInput = part.usage?.inputTokens ?? 0;
|
||||
const stepCached = part.usage?.cachedInputTokens ?? 0;
|
||||
const stepCached =
|
||||
part.usage?.inputTokenDetails?.cacheReadTokens ?? 0;
|
||||
const stepCacheCreation = extractCacheCreationTokens(
|
||||
(
|
||||
part as {
|
||||
@@ -245,8 +246,8 @@ function computeStreamCosts(
|
||||
| {
|
||||
inputTokens?: number;
|
||||
outputTokens?: number;
|
||||
cachedInputTokens?: number;
|
||||
reasoningTokens?: number;
|
||||
inputTokenDetails?: { cacheReadTokens?: number };
|
||||
outputTokenDetails?: { reasoningTokens?: number };
|
||||
}
|
||||
| undefined,
|
||||
cacheCreationTokens: number,
|
||||
@@ -254,8 +255,8 @@ function computeStreamCosts(
|
||||
const breakdown = computeCostBreakdown(modelConfig, {
|
||||
inputTokens: totalUsage?.inputTokens,
|
||||
outputTokens: totalUsage?.outputTokens,
|
||||
cachedInputTokens: totalUsage?.cachedInputTokens,
|
||||
reasoningTokens: totalUsage?.reasoningTokens,
|
||||
cachedInputTokens: totalUsage?.inputTokenDetails?.cacheReadTokens,
|
||||
reasoningTokens: totalUsage?.outputTokenDetails?.reasoningTokens,
|
||||
cacheCreationTokens,
|
||||
});
|
||||
|
||||
|
||||
+35
-14
@@ -124,8 +124,6 @@ export class ChatExecutionService {
|
||||
toolContext,
|
||||
);
|
||||
|
||||
const preloadedToolNames = Object.keys(preloadedTools);
|
||||
|
||||
const modelId = workspace.smartModel;
|
||||
|
||||
this.aiModelRegistryService.validateModelAvailability(modelId, workspace);
|
||||
@@ -139,13 +137,21 @@ export class ChatExecutionService {
|
||||
registeredModel.modelId,
|
||||
);
|
||||
|
||||
const { tools: nativeSearchTools, callableToolNames: searchToolNames } =
|
||||
this.getNativeWebSearchTools(registeredModel.inferenceProvider);
|
||||
|
||||
// Direct tools: native provider tools + preloaded tools.
|
||||
// These are callable directly AND as fallback through execute_tool.
|
||||
const directTools: ToolSet = {
|
||||
...wrapToolsWithOutputSerialization(preloadedTools),
|
||||
...this.getNativeWebSearchTool(registeredModel.inferenceProvider),
|
||||
...nativeSearchTools,
|
||||
};
|
||||
|
||||
const preloadedToolNames = [
|
||||
...Object.keys(preloadedTools),
|
||||
...searchToolNames,
|
||||
];
|
||||
|
||||
// ToolSet is constant for the entire conversation — no mutation.
|
||||
// learn_tools returns schemas as text; execute_tool dispatches to cached tools.
|
||||
const activeTools: ToolSet = {
|
||||
@@ -205,9 +211,11 @@ export class ChatExecutionService {
|
||||
: undefined,
|
||||
};
|
||||
|
||||
const modelMessages = await convertToModelMessages(processedMessages);
|
||||
|
||||
const stream = streamText({
|
||||
model: registeredModel.model,
|
||||
messages: [systemMessage, ...convertToModelMessages(processedMessages)],
|
||||
messages: [systemMessage, ...modelMessages],
|
||||
tools: activeTools,
|
||||
stopWhen: stepCountIs(AGENT_CONFIG.MAX_STEPS),
|
||||
experimental_telemetry: AI_TELEMETRY_CONFIG,
|
||||
@@ -318,33 +326,46 @@ export class ChatExecutionService {
|
||||
return context;
|
||||
}
|
||||
|
||||
private getNativeWebSearchTool(
|
||||
inferenceProvider: InferenceProvider,
|
||||
): ToolSet {
|
||||
private getNativeWebSearchTools(inferenceProvider: InferenceProvider): {
|
||||
tools: ToolSet;
|
||||
callableToolNames: string[];
|
||||
} {
|
||||
switch (inferenceProvider) {
|
||||
case InferenceProvider.ANTHROPIC:
|
||||
return { web_search: anthropic.tools.webSearch_20250305() };
|
||||
return {
|
||||
tools: { web_search: anthropic.tools.webSearch_20250305() },
|
||||
callableToolNames: ['web_search'],
|
||||
};
|
||||
case InferenceProvider.BEDROCK: {
|
||||
const bedrockProvider =
|
||||
this.aiModelRegistryService.getBedrockProvider();
|
||||
|
||||
if (bedrockProvider) {
|
||||
return {
|
||||
web_search:
|
||||
bedrockProvider.tools.webSearch_20250305() as ToolSet[string],
|
||||
tools: {
|
||||
web_search:
|
||||
bedrockProvider.tools.webSearch_20250305() as ToolSet[string],
|
||||
},
|
||||
callableToolNames: ['web_search'],
|
||||
};
|
||||
}
|
||||
|
||||
return {};
|
||||
return { tools: {}, callableToolNames: [] };
|
||||
}
|
||||
case InferenceProvider.OPENAI:
|
||||
return { web_search: openai.tools.webSearch() };
|
||||
return {
|
||||
tools: { web_search: openai.tools.webSearch() },
|
||||
callableToolNames: ['web_search'],
|
||||
};
|
||||
case InferenceProvider.GROQ:
|
||||
return {
|
||||
web_search: groq.tools.browserSearch({}) as ToolSet[string],
|
||||
tools: {
|
||||
web_search: groq.tools.browserSearch({}) as ToolSet[string],
|
||||
},
|
||||
callableToolNames: [],
|
||||
};
|
||||
default:
|
||||
return {};
|
||||
return { tools: {}, callableToolNames: [] };
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+17
-5
@@ -246,6 +246,7 @@ ${skillsList}`;
|
||||
preloadedTools: string[],
|
||||
): string {
|
||||
const preloadedSet = new Set(preloadedTools);
|
||||
const hasWebSearch = preloadedSet.has('web_search');
|
||||
|
||||
const toolsByCategory = new Map<string, ToolIndexEntry[]>();
|
||||
|
||||
@@ -259,6 +260,14 @@ ${skillsList}`;
|
||||
|
||||
const sections: string[] = [];
|
||||
|
||||
const webSearchLine = hasWebSearch
|
||||
? `- \`web_search\` ✓: Search the web for real-time information (ALWAYS use this for current data, news, research)`
|
||||
: `- Web search is automatically available — the model will search the web when needed. Do NOT call \`web_search\` as a tool.`;
|
||||
|
||||
const otherPreloadedTools = preloadedTools.filter(
|
||||
(name) => name !== 'web_search',
|
||||
);
|
||||
|
||||
sections.push(`
|
||||
## Available Tools
|
||||
|
||||
@@ -266,8 +275,8 @@ You have access to ${toolCatalog.length} tools plus native web search. Some are
|
||||
To use any other tool, first call \`${LEARN_TOOLS_TOOL_NAME}\` to learn its schema, then call \`${EXECUTE_TOOL_TOOL_NAME}\` to run it.
|
||||
|
||||
### Pre-loaded Tools (ready to use now)
|
||||
- \`web_search\` ✓: Search the web for real-time information (ALWAYS use this for current data, news, research)
|
||||
${preloadedTools.length > 0 ? preloadedTools.map((toolName) => `- \`${toolName}\` ✓`).join('\n') : ''}
|
||||
${webSearchLine}
|
||||
${otherPreloadedTools.length > 0 ? otherPreloadedTools.map((toolName) => `- \`${toolName}\` ✓`).join('\n') : ''}
|
||||
|
||||
### Tool Catalog by Category`);
|
||||
|
||||
@@ -301,11 +310,14 @@ ${tools
|
||||
.join('\n')}`);
|
||||
}
|
||||
|
||||
const webSearchInstruction = hasWebSearch
|
||||
? `1. **Web search** (\`web_search\`): Use for ANY request requiring current/real-time information from the internet\n`
|
||||
: '';
|
||||
|
||||
sections.push(`
|
||||
### How to Use Tools
|
||||
1. **Web search** (\`web_search\`): Use for ANY request requiring current/real-time information from the internet
|
||||
2. **Pre-loaded tools** (marked with ✓): Use directly
|
||||
3. **Other tools**: First call \`${LEARN_TOOLS_TOOL_NAME}({toolNames: ["tool_name"]})\` to learn the schema, then call \`${EXECUTE_TOOL_TOOL_NAME}({toolName: "tool_name", arguments: {...}})\` to run it`);
|
||||
${webSearchInstruction}${hasWebSearch ? '2' : '1'}. **Pre-loaded tools** (marked with ✓): Use directly
|
||||
${hasWebSearch ? '3' : '2'}. **Other tools**: First call \`${LEARN_TOOLS_TOOL_NAME}({toolNames: ["tool_name"]})\` to learn the schema, then call \`${EXECUTE_TOOL_TOOL_NAME}({toolName: "tool_name", arguments: {...}})\` to run it`);
|
||||
|
||||
return sections.join('\n');
|
||||
}
|
||||
|
||||
+1
-1
@@ -35,7 +35,7 @@ export class AiService {
|
||||
maxOutputTokens?: number;
|
||||
model: LanguageModel;
|
||||
};
|
||||
}) {
|
||||
}): ReturnType<typeof streamText> {
|
||||
return streamText({
|
||||
model: options.model,
|
||||
messages,
|
||||
|
||||
Reference in New Issue
Block a user