Merge pull request #3085 from honey-tree/claude-caching-at-depth Claude caching at depth

54db4983f4663d77db79ec1246888a5791bdb619

Cohee <18619528+Cohee1207@users.noreply.github.com>

Signed

Parents: dde9f93d77e58f06d8907c8dd5006e8495a5303b 3b03f07ebfb66f5a61597f5f022cea055e37d31b

default/config.yaml+8 -0 Viewed

   # (e.g {{random}} macro or lorebooks not as in-chat injections).
   # Otherwise, you'll just waste money on cache misses.
   enableSystemPromptCache: false
+  # Enables caching of the message history at depth (if supported).
+  # https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching
+  # -- IMPORTANT! --
+  # Use with caution. Behavior may be unpredictable and no guarantees can or will be made.
+  # Set to an integer to specify the desired depth. 0 (which does NOT include the prefill)
+  # should be ideal for most use cases.
+  # Any value other than a non-negative integer will be ignored and caching at depth will not be enabled.
+  cachingAtDepth: -1
 # -- SERVER PLUGIN CONFIGURATION --
 enableServerPlugins: false

src/endpoints/backends/chat-completions.js+19 -1 Viewed

     convertMistralMessages,
     convertAI21Messages,
     mergeMessages,
+    cachingAtDepthForOpenRouterClaude,
+    cachingAtDepthForClaude,
 } from '../../prompt-converters.js';
 import { readSecret, SECRET_KEYS } from '../secrets.js';

     const apiKey = request.body.reverse_proxy ? request.body.proxy_password : readSecret(request.user.directories, SECRET_KEYS.CLAUDE);
     const divider = '-'.repeat(process.stdout.columns);
     const enableSystemPromptCache = getConfigValue('claude.enableSystemPromptCache', false) && request.body.model.startsWith('claude-3');
+    let cachingAtDepth = getConfigValue('claude.cachingAtDepth', -1);
+    // Disabled if not an integer or negative, or if the model doesn't support it
+    if (!Number.isInteger(cachingAtDepth) || cachingAtDepth < 0 || !request.body.model.startsWith('claude-3')) {
+        cachingAtDepth = -1;
+    }
     if (!apiKey) {
         console.log(color.red(`Claude API key is missing.\n${divider}`));

@@ -138,9 +145,15 @@ async function sendClaudeRequest(request, response) {
138	requestBody.tools[requestBody.tools.length - 1]['cache_control'] = { type: 'ephemeral' };	145	requestBody.tools[requestBody.tools.length - 1]['cache_control'] = { type: 'ephemeral' };
139	}	146	}
140	}	147	}
141	if (enableSystemPromptCache) {	148
		149	if (cachingAtDepth !== -1) {
		150	cachingAtDepthForClaude(convertedPrompt.messages, cachingAtDepth);
		151	}
		152
		153	if (enableSystemPromptCache \|\| cachingAtDepth !== -1) {
142	additionalHeaders['anthropic-beta'] = 'prompt-caching-2024-07-31';	154	additionalHeaders['anthropic-beta'] = 'prompt-caching-2024-07-31';
143	}	155	}
		156
144	console.log('Claude request:', requestBody);	157	console.log('Claude request:', requestBody);
145		158
146	const generateResponse = await fetch(apiUrl + '/messages', {	159	const generateResponse = await fetch(apiUrl + '/messages', {

         if (request.body.use_fallback) {
             bodyParams['route'] = 'fallback';
+        }
+        let cachingAtDepth = getConfigValue('claude.cachingAtDepth', -1);
+        if (Number.isInteger(cachingAtDepth) && cachingAtDepth >= 0 && request.body.model.startsWith('anthropic/claude-3')) {
+            cachingAtDepthForOpenRouterClaude(request.body.messages, cachingAtDepth);
+        }
     } else if (request.body.chat_completion_source === CHAT_COMPLETION_SOURCES.CUSTOM) {
         apiUrl = request.body.custom_url;
         apiKey = readSecret(request.user.directories, SECRET_KEYS.CUSTOM);

src/prompt-converters.js+74 -0 Viewed

     });
     return messageStrings.join('\n') + '\nassistant:';
+}
+export function cachingAtDepthForClaude(messages, cachingAtDepth) {
+    let passedThePrefill = false;
+    let depth = 0;
+    let previousRoleName = '';
+    for (let i = messages.length - 1; i >= 0; i--) {
+        if (!passedThePrefill && messages[i].role === 'assistant') {
+            continue;
+        }
+        passedThePrefill = true;
+        if (messages[i].role !== previousRoleName) {
+            if (depth === cachingAtDepth || depth === cachingAtDepth + 2) {
+                const content = messages[i].content;
+                content[content.length - 1].cache_control = { type: 'ephemeral' };
+            }
+            if (depth === cachingAtDepth + 2) {
+                break;
+            }
+            depth += 1;
+            previousRoleName = messages[i].role;
+        }
+    }
+}
+/**
+ * Append cache_control headers to an OpenRouter request at depth. Directly modifies the
+ * messages array.
+ * @param {object[]} messages Array of messages
+ * @param {number} cachingAtDepth Depth at which caching is supposed to occur
+ */
+export function cachingAtDepthForOpenRouterClaude(messages, cachingAtDepth) {
+    //caching the prefill is a terrible idea in general
+    let passedThePrefill = false;
+    //depth here is the number of message role switches
+    let depth = 0;
+    let previousRoleName = '';
+    for (let i = messages.length - 1; i >= 0; i--) {
+        if (!passedThePrefill && messages[i].role === 'assistant') {
+            continue;
+        }
+        passedThePrefill = true;
+        if (messages[i].role !== previousRoleName) {
+            if (depth === cachingAtDepth || depth === cachingAtDepth + 2) {
+                const content = messages[i].content;
+                if (typeof content === 'string') {
+                    messages[i].content = [{
+                        type: 'text',
+                        text: content,
+                        cache_control: { type: 'ephemeral' },
+                    }];
+                } else {
+                    const contentPartCount = content.length;
+                    content[contentPartCount - 1].cache_control = {
+                        type: 'ephemeral',
+                    };
+                }
+            }
+            if (depth === cachingAtDepth + 2) {
+                break;
+            }
+            depth += 1;
+            previousRoleName = messages[i].role;
+        }
+    }
+}