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}`));

                 requestBody.tools[requestBody.tools.length - 1]['cache_control'] = { type: 'ephemeral' };
+            }
+        }
 -        if (enableSystemPromptCache) {
++
 +        if (cachingAtDepth !== -1) {
 +            cachingAtDepthForClaude(convertedPrompt.messages, cachingAtDepth);
 +        }
++
 +        if (enableSystemPromptCache || cachingAtDepth !== -1) {
             additionalHeaders['anthropic-beta'] = 'prompt-caching-2024-07-31';
+        }
++
         console.log('Claude request:', requestBody);
         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;
 +        }
 +    }
 +}