diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent/utils/repair-tool-call.util.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent/utils/repair-tool-call.util.ts index ddf71c0b63..5a4af61549 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent/utils/repair-tool-call.util.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent/utils/repair-tool-call.util.ts @@ -1,6 +1,17 @@ -import { type LanguageModel, NoSuchToolError, Output, generateText } from 'ai'; +import { + type LanguageModel, + type LanguageModelUsage, + NoSuchToolError, + Output, + type StepResult, + type ToolSet, + generateText, +} from 'ai'; import { type z } from 'zod'; +import { UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-operation-type.enum'; +import { AiBillingService } from 'src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service'; +import { extractCacheCreationTokensFromSteps } from 'src/engine/metadata-modules/ai/ai-billing/utils/extract-cache-creation-tokens.util'; import { AI_TELEMETRY_CONFIG } from 'src/engine/metadata-modules/ai/ai-models/constants/ai-telemetry.const'; type ToolCall = { @@ -10,18 +21,28 @@ type ToolCall = { input: string; }; +type RepairToolCallBillingContext = { + aiBillingService: AiBillingService; + modelId: string; + workspaceId: string; + userWorkspaceId: string | null; + operationType: UsageOperationType; +}; + export const repairToolCall = async ({ toolCall, tools, inputSchema, error, model, + billingContext, }: { toolCall: ToolCall; tools: Record; inputSchema: (toolCall: { toolName: string }) => unknown; error: Error; model: LanguageModel; + billingContext?: RepairToolCallBillingContext; }): Promise => { // Don't attempt to fix invalid tool names if (NoSuchToolError.isInstance(error)) { @@ -40,8 +61,11 @@ export const repairToolCall = async ({ return null; } + let usage: LanguageModelUsage | undefined; + let steps: StepResult[] | undefined; + try { - const { output: repairedInput } = await generateText({ + const result = await generateText({ model, output: Output.object({ schema: schema as z.ZodTypeAny }), prompt: [ @@ -62,6 +86,11 @@ export const repairToolCall = async ({ experimental_telemetry: AI_TELEMETRY_CONFIG, }); + usage = result.usage; + steps = result.steps; + + const repairedInput = result.output; + if (repairedInput == null) { return null; } @@ -75,5 +104,20 @@ export const repairToolCall = async ({ } catch { // If repair fails, return null to let the error propagate return null; + } finally { + if (billingContext && usage) { + const cacheCreationTokens = steps + ? extractCacheCreationTokensFromSteps(steps) + : 0; + + billingContext.aiBillingService.calculateAndBillUsage( + billingContext.modelId, + { usage, cacheCreationTokens }, + billingContext.workspaceId, + billingContext.operationType, + null, + billingContext.userWorkspaceId, + ); + } } }; diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service.ts index 621970a50a..4f5fbfa7bb 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service.ts @@ -328,8 +328,11 @@ export class AgentChatService { return null; } - const title = - await this.titleGenerationService.generateThreadTitle(messageContent); + const title = await this.titleGenerationService.generateThreadTitle( + messageContent, + workspaceId, + thread.userWorkspaceId, + ); await this.threadRepository.update(threadId, { title }); diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-title-generation.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-title-generation.service.ts index 8fed08f177..736a0c82f8 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-title-generation.service.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-title-generation.service.ts @@ -1,7 +1,15 @@ import { Injectable, Logger } from '@nestjs/common'; -import { generateText } from 'ai'; +import { + type LanguageModelUsage, + type StepResult, + type ToolSet, + generateText, +} from 'ai'; +import { UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-operation-type.enum'; +import { AiBillingService } from 'src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service'; +import { extractCacheCreationTokensFromSteps } from 'src/engine/metadata-modules/ai/ai-billing/utils/extract-cache-creation-tokens.util'; import { AI_TELEMETRY_CONFIG } from 'src/engine/metadata-modules/ai/ai-models/constants/ai-telemetry.const'; import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models/services/ai-model-registry.service'; @@ -11,29 +19,55 @@ export class AgentTitleGenerationService { constructor( private readonly aiModelRegistryService: AiModelRegistryService, + private readonly aiBillingService: AiBillingService, ) {} - async generateThreadTitle(messageContent: string): Promise { + async generateThreadTitle( + messageContent: string, + workspaceId: string, + userWorkspaceId: string | null, + ): Promise { + const defaultModel = this.aiModelRegistryService.getDefaultSpeedModel(); + + if (!defaultModel) { + this.logger.warn('No default AI model available for title generation'); + + return this.generateFallbackTitle(messageContent); + } + + let usage: LanguageModelUsage | undefined; + let steps: StepResult[] | undefined; + try { - const defaultModel = this.aiModelRegistryService.getDefaultSpeedModel(); - - if (!defaultModel) { - this.logger.warn('No default AI model available for title generation'); - - return this.generateFallbackTitle(messageContent); - } - const result = await generateText({ model: defaultModel.model, prompt: `Generate a concise, descriptive title (maximum 60 characters) for a chat thread based on the following message. The title should capture the main topic or purpose of the conversation. Return only the title, nothing else. Message: "${messageContent}"`, experimental_telemetry: AI_TELEMETRY_CONFIG, }); + usage = result.usage; + steps = result.steps; + return this.cleanTitle(result.text); } catch (error) { this.logger.error('Failed to generate title with AI:', error); return this.generateFallbackTitle(messageContent); + } finally { + if (usage) { + const cacheCreationTokens = steps + ? extractCacheCreationTokensFromSteps(steps) + : 0; + + this.aiBillingService.calculateAndBillUsage( + defaultModel.modelId, + { usage, cacheCreationTokens }, + workspaceId, + UsageOperationType.AI_CHAT_TOKEN, + null, + userWorkspaceId, + ); + } } } diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts index 8905a72f0e..5cabaac7f5 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts @@ -348,6 +348,13 @@ export class ChatExecutionService { inputSchema, error, model: registeredModel.model, + billingContext: { + aiBillingService: this.aiBillingService, + modelId: registeredModel.modelId, + workspaceId: workspace.id, + userWorkspaceId, + operationType: UsageOperationType.AI_CHAT_TOKEN, + }, }); }, });