Initial open-source release

packages/llm/src/message-builder.ts

@@ -0,0 +1,197 @@
+import type { LanguageModel, ModelMessage, ToolSet } from 'ai';
+import { isAnthropicProtocolModel } from './model-provider.js';
+import type { KloLlmProvider, KloPromptCacheTtl, KloPromptParts } from './types.js';
+
+type ToolMap = ToolSet | Record<string, Record<string, unknown>>;
+
+interface KloMessageBuilderOptions {
+  cacheSystem?: boolean;
+  cacheTools?: boolean;
+  cacheLastHistory?: boolean;
+}
+
+interface KloBuildInput {
+  parts: KloPromptParts;
+  history: ModelMessage[];
+  currentMessage: ModelMessage;
+  tools: ToolMap;
+  model: LanguageModel | string;
+}
+
+interface KloWrapSimpleInput {
+  system?: string;
+  messages?: ModelMessage[];
+  tools?: ToolMap;
+  model: LanguageModel | string;
+}
+
+interface KloBuildOutput {
+  messages: ModelMessage[];
+  tools: ToolMap;
+}
+
+export class KloMessageBuilder {
+  constructor(
+    private readonly provider: KloLlmProvider,
+    private readonly options: KloMessageBuilderOptions = {},
+  ) {}
+
+  build(input: KloBuildInput): KloBuildOutput {
+    const cfg = this.provider.promptCachingConfig();
+    const cachingActive = cfg.enabled && isAnthropicProtocolModel(input.model);
+    const ttls = this.resolveTtls(input.model);
+    const messages: ModelMessage[] = [];
+
+    const systemMessage: ModelMessage & { providerOptions?: unknown } = {
+      role: 'system',
+      content: input.parts.staticSystem,
+    };
+    if (cachingActive && this.cacheSystemEnabled()) {
+      systemMessage.providerOptions = this.provider.cacheMarker(ttls.systemTtl, input.model);
+    }
+    messages.push(systemMessage);
+
+    if (input.parts.dynamicSystem) {
+      messages.push({ role: 'system', content: input.parts.dynamicSystem });
+    }
+
+    const historyToEmit =
+      cachingActive && this.cacheHistoryEnabled()
+        ? this.markLastHistoryMessage(input.history, ttls.historyTtl, input.model)
+        : input.history;
+    messages.push(...historyToEmit);
+    messages.push(this.wrapLeading(input.currentMessage, input.parts.leadingUserContext));
+
+    return {
+      messages,
+      tools: this.sortAndMarkTools(input.tools, cachingActive, this.cacheToolsEnabled(), ttls.toolsTtl, input.model),
+    };
+  }
+
+  wrapSimple(input: KloWrapSimpleInput): KloBuildOutput {
+    const cfg = this.provider.promptCachingConfig();
+    const cachingActive = cfg.enabled && isAnthropicProtocolModel(input.model);
+    const ttls = this.resolveTtls(input.model);
+    const messages: ModelMessage[] = [];
+
+    if (input.system) {
+      const systemMessage: ModelMessage & { providerOptions?: unknown } = {
+        role: 'system',
+        content: input.system,
+      };
+      if (cachingActive && this.cacheSystemEnabled()) {
+        systemMessage.providerOptions = this.provider.cacheMarker(ttls.systemTtl, input.model);
+      }
+      messages.push(systemMessage);
+    }
+
+    if (input.messages) {
+      messages.push(
+        ...(cachingActive && this.cacheHistoryEnabled()
+          ? this.markLastHistoryMessage(input.messages, ttls.historyTtl, input.model)
+          : input.messages),
+      );
+    }
+
+    return {
+      messages,
+      tools: this.sortAndMarkTools(input.tools ?? {}, cachingActive, this.cacheToolsEnabled(), ttls.toolsTtl, input.model),
+    };
+  }
+
+  private cacheSystemEnabled(): boolean {
+    return this.options.cacheSystem ?? this.provider.promptCachingConfig().cacheSystem;
+  }
+
+  private cacheToolsEnabled(): boolean {
+    return this.options.cacheTools ?? this.provider.promptCachingConfig().cacheTools;
+  }
+
+  private cacheHistoryEnabled(): boolean {
+    return this.options.cacheLastHistory ?? this.provider.promptCachingConfig().cacheHistory;
+  }
+
+  private resolveTtls(model: LanguageModel | string): {
+    systemTtl: KloPromptCacheTtl;
+    toolsTtl: KloPromptCacheTtl;
+    historyTtl: KloPromptCacheTtl;
+  } {
+    const cfg = this.provider.promptCachingConfig();
+    if (cfg.vertexFallbackTo5m && this.provider.activeBackend() === 'vertex' && isAnthropicProtocolModel(model)) {
+      return { systemTtl: '5m', toolsTtl: '5m', historyTtl: '5m' };
+    }
+    return { systemTtl: cfg.systemTtl, toolsTtl: cfg.toolsTtl, historyTtl: cfg.historyTtl };
+  }
+
+  private wrapLeading(currentMessage: ModelMessage, leadingUserContext?: string): ModelMessage {
+    if (!leadingUserContext) {
+      return currentMessage;
+    }
+    const reminderPart = {
+      type: 'text' as const,
+      text: `<system-reminder>\n${leadingUserContext}\n</system-reminder>`,
+    };
+    if (typeof currentMessage.content === 'string') {
+      return {
+        ...currentMessage,
+        content: [reminderPart, { type: 'text' as const, text: currentMessage.content }],
+      } as ModelMessage;
+    }
+    if (Array.isArray(currentMessage.content)) {
+      return { ...currentMessage, content: [reminderPart, ...currentMessage.content] } as ModelMessage;
+    }
+    return currentMessage;
+  }
+
+  private markLastHistoryMessage(
+    history: ModelMessage[],
+    ttl: KloPromptCacheTtl,
+    model: LanguageModel | string,
+  ): ModelMessage[] {
+    if (history.length === 0) {
+      return history;
+    }
+    const out = [...history];
+    const last = out[out.length - 1];
+    const marker = this.provider.cacheMarker(ttl, model);
+    if (!marker) {
+      return history;
+    }
+    if (typeof last.content === 'string') {
+      out[out.length - 1] = {
+        ...last,
+        content: [{ type: 'text', text: last.content, providerOptions: marker }],
+      } as ModelMessage;
+      return out;
+    }
+    if (Array.isArray(last.content) && last.content.length > 0) {
+      const parts = [...last.content];
+      const lastPart = parts[parts.length - 1];
+      parts[parts.length - 1] = Object.assign({}, lastPart, { providerOptions: marker });
+      out[out.length - 1] = { ...last, content: parts } as ModelMessage;
+    }
+    return out;
+  }
+
+  private sortAndMarkTools(
+    tools: ToolMap,
+    cachingActive: boolean,
+    cacheTools: boolean,
+    ttl: KloPromptCacheTtl,
+    model: LanguageModel | string,
+  ): ToolMap {
+    const keys = Object.keys(tools).sort();
+    const sorted: Record<string, unknown> = {};
+    for (const key of keys) {
+      sorted[key] = tools[key as keyof typeof tools];
+    }
+    if (cachingActive && cacheTools && keys.length > 0) {
+      const lastKey = keys[keys.length - 1];
+      const marker = this.provider.cacheMarker(ttl, model);
+      if (marker) {
+        sorted[lastKey] = { ...(sorted[lastKey] as Record<string, unknown>), providerOptions: marker };
+      }
+    }
+    return sorted as ToolMap;
+  }
+}