Ai Chat - Caching optim (#20126)

EDIT : - solving auto-caching from Anthropic by updating ai-sdk/anthropic + adding providerOption at stream level - concerning Bedrock, it needs breakpoint **1. Breakpoints were only on the system prompt** The code already placed a cache marker on the system prompt (~10K tokens). But the conversation history — which can grow to hundreds of thousands of tokens — had no marker, so Anthropic re-read it at full price on every turn. The fix adds a prepareStep hook inside streamText that stamps the last message with a cache breakpoint before every LLM call. Anthropic then caches the entire conversation prefix, and subsequent turns read it at $0.30/M instead of $3/M. prepareStep is used rather than a one-shot pre-processing step because an agentic turn makes multiple internal LLM calls as tool results accumulate — the hook refreshes the breakpoint before each one. **2. Bedrock was using the wrong field** The system prompt marker for Bedrock was set as cacheControl: { type: 'ephemeral' } — which is the Anthropic wire format. The Bedrock Converse API expects cachePoint: { type: 'default' }. The system prompt was silently not being cached on Bedrock at all. Both the system prompt and the new prepareStep now go through a shared getCacheProviderOptions helper that returns the correct field per provider. **3. Persisted cached token usage to monitor cache strat. efficiency**
2026-04-29 16:42:12 +02:00
parent 11628d19a3
commit fd6d5f895d
9 changed files with 195 additions and 151 deletions
@@ -60,6 +60,12 @@ export class AgentChatThreadEntity {
  @Column({ type: 'bigint', default: 0 })
  totalOutputCredits: number;

+  @Column({ type: 'bigint', default: 0 })
+  totalCacheReadTokens: number;
+
+  @Column({ type: 'bigint', default: 0 })
+  totalCacheCreationTokens: number;
+
  @Column({ type: 'varchar', nullable: true })
  activeStreamId: string | null;

@@ -188,6 +188,7 @@ export class StreamAgentChatJob {
        outputTokens: 0,
        inputCredits: 0,
        outputCredits: 0,
+        cacheReadTokens: 0,
      };
      let lastStepConversationSize = 0;
      let totalCacheCreationTokens = 0;
@@ -279,6 +280,7 @@ export class StreamAgentChatJob {
                    workspaceId: data.workspaceId,
                    streamUsage,
                    lastStepConversationSize,
+                    totalCacheCreationTokens,
                    modelConfig,
                    userMessagePromise,
                  });
@@ -359,6 +361,7 @@ export class StreamAgentChatJob {
      outputTokens: number;
      inputCredits: number;
      outputCredits: number;
+      cacheReadTokens: number;
    }) => void;
    onUpdateConversationSize: (size: number) => void;
    onUpdateCacheCreationTokens: (tokens: number) => void;
@@ -395,6 +398,7 @@ export class StreamAgentChatJob {
        outputTokens: part.totalUsage?.outputTokens ?? 0,
        inputCredits,
        outputCredits,
+        cacheReadTokens: breakdown.tokenCounts.cachedInputTokens,
      });

      return {
@@ -422,6 +426,7 @@ export class StreamAgentChatJob {
    workspaceId,
    streamUsage,
    lastStepConversationSize,
+    totalCacheCreationTokens,
    modelConfig,
    userMessagePromise,
  }: {
@@ -433,8 +438,10 @@ export class StreamAgentChatJob {
      outputTokens: number;
      inputCredits: number;
      outputCredits: number;
+      cacheReadTokens: number;
    };
    lastStepConversationSize: number;
+    totalCacheCreationTokens: number;
    modelConfig: AiModelConfig;
    userMessagePromise: Promise<{ turnId: string | null }>;
  }): Promise<void> {
@@ -459,6 +466,10 @@ export class StreamAgentChatJob {
        `"totalInputCredits" + ${streamUsage.inputCredits}`,
      totalOutputCredits: () =>
        `"totalOutputCredits" + ${streamUsage.outputCredits}`,
+      totalCacheReadTokens: () =>
+        `"totalCacheReadTokens" + ${streamUsage.cacheReadTokens}`,
+      totalCacheCreationTokens: () =>
+        `"totalCacheCreationTokens" + ${totalCacheCreationTokens}`,
      contextWindowTokens: modelConfig.contextWindowTokens,
      conversationSize: lastStepConversationSize,
    });
@@ -48,9 +48,10 @@ import {
  type ExtractedFile,
 } from 'src/engine/metadata-modules/ai/ai-chat/utils/extract-code-interpreter-files.util';
 import {
-  AI_SDK_ANTHROPIC,
-  AI_SDK_BEDROCK,
-} from 'src/engine/metadata-modules/ai/ai-models/constants/ai-sdk-package.const';
+  injectCacheBreakpoint,
+  getCacheProviderOptions,
+  getCallLevelCacheProviderOptions,
+} from 'src/engine/metadata-modules/ai/ai-chat/utils/inject-cache-breakpoint.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';
 import { type AiModelConfig } from 'src/engine/metadata-modules/ai/ai-models/types/ai-model-config.type';
@@ -237,12 +238,7 @@ export class ChatExecutionService {
    const systemMessage: SystemModelMessage = {
      role: 'system',
      content: systemPrompt,
-      providerOptions:
-        registeredModel.sdkPackage === AI_SDK_ANTHROPIC
-          ? { anthropic: { cacheControl: { type: 'ephemeral' } } }
-          : registeredModel.sdkPackage === AI_SDK_BEDROCK
-            ? { bedrock: { cacheControl: { type: 'ephemeral' } } }
-            : undefined,
+      providerOptions: getCacheProviderOptions(registeredModel.sdkPackage),
    };

    const rawModelMessages = await convertToModelMessages(processedMessages);
@@ -333,6 +329,12 @@ export class ChatExecutionService {
      abortSignal,
      stopWhen: stepCountIs(AGENT_CONFIG.MAX_STEPS),
      experimental_telemetry: AI_TELEMETRY_CONFIG,
+      providerOptions: getCallLevelCacheProviderOptions(
+        registeredModel.sdkPackage,
+      ),
+      prepareStep: ({ messages }) => ({
+        messages: injectCacheBreakpoint(messages, registeredModel.sdkPackage),
+      }),
      onAbort: async ({ steps }) => {
        await billUsageFromSteps(steps);
      },
@@ -0,0 +1,52 @@
+import { type ModelMessage } from 'ai';
+import { type ProviderOptions } from '@ai-sdk/provider-utils';
+
+import {
+  AI_SDK_ANTHROPIC,
+  AI_SDK_BEDROCK,
+} from 'src/engine/metadata-modules/ai/ai-models/constants/ai-sdk-package.const';
+
+export const getCallLevelCacheProviderOptions = (
+  sdkPackage: string,
+): ProviderOptions | undefined => {
+  if (sdkPackage === AI_SDK_ANTHROPIC) {
+    return { anthropic: { cacheControl: { type: 'ephemeral' } } };
+  }
+
+  return undefined;
+};
+
+export const getCacheProviderOptions = (
+  sdkPackage: string,
+): ProviderOptions | undefined => {
+  if (sdkPackage === AI_SDK_BEDROCK) {
+    return { bedrock: { cachePoint: { type: 'default' } } };
+  }
+
+  return undefined;
+};
+
+export const injectCacheBreakpoint = (
+  messages: ModelMessage[],
+  sdkPackage: string,
+): ModelMessage[] => {
+  if (messages.length === 0) return messages;
+
+  const cacheOptions = getCacheProviderOptions(sdkPackage);
+
+  if (!cacheOptions) return messages;
+
+  const lastIdx = messages.length - 1;
+
+  return messages.map((message, index) => {
+    if (index !== lastIdx) return message;
+
+    return {
+      ...message,
+      providerOptions: {
+        ...(message.providerOptions ?? {}),
+        ...cacheOptions,
+      },
+    };
+  });
+};
@@ -45,10 +45,10 @@ export class AiModelConfigService {
    model: RegisteredAiModel,
    options: NativeModelToolOptions,
  ): ToolSet {
-    const tools: ToolSet = {};
+    const tools: Record<string, unknown> = {};

    if (!options.webSearchEnabled) {
-      return tools;
+      return tools as ToolSet;
    }

    switch (model.sdkPackage) {
@@ -76,7 +76,7 @@ export class AiModelConfigService {
      }
    }

-    return tools;
+    return tools as ToolSet;
  }

  private getXaiProviderOptions(agent: FlatAgentWithRoleId): ProviderOptions {