Extract adding cache control to a helper

Author: jahooma

backend/src/main-prompt.ts

@@ -36,7 +36,7 @@ import {
 } from './tools'
 import { trimMessagesToFitTokenLimit } from './util/messages'
 import { checkTerminalCommand } from './check-terminal-command'
-import { toContentString } from 'common/util/messages'
+import { withCacheControl, toContentString } from 'common/util/messages'
 
 export const mainPrompt = async (
   ws: WebSocket,
@@ -782,20 +782,7 @@ function getMessagesSubset(messages: Message[], otherTokens: number) {
   // Cache up to the last message!
   const lastMessage = messagesSubset[messagesSubset.length - 1]
   if (lastMessage) {
-    if (typeof lastMessage.content === 'string') {
-      // Transform to a content array
-      ;(lastMessage as any).content = [
-        {
-          type: 'text',
-          text: lastMessage.content,
-          cache_control: { type: 'ephemeral' as const },
-        },
-      ]
-    } else {
-      lastMessage.content[lastMessage.content.length - 1].cache_control = {
-        type: 'ephemeral' as const,
-      }
-    }
+    messagesSubset[messagesSubset.length - 1] = withCacheControl(lastMessage)
   } else {
     logger.debug(
       {

common/src/util/messages.ts

@@ -47,3 +47,27 @@ export function toContentString(msg: Message): string {
   if (typeof content === 'string') return content
   return content.map((item) => (item as any)?.text ?? '').join('\n')
 }
+
+export function withCacheControl(msg: Message): Message {
+  if (typeof msg.content === 'string') {
+    return {
+      ...msg,
+      content: [
+        {
+          type: 'text',
+          text: msg.content,
+          cache_control: { type: 'ephemeral' as const },
+        },
+      ],
+    }
+  } else {
+    return {
+      ...msg,
+      content: msg.content.map((item, i) =>
+        i === msg.content.length - 1
+          ? { ...item, cache_control: { type: 'ephemeral' as const } }
+          : item
+      ),
+    }
+  }
+}