From 699f0f3fe73bad07ffb3143f051e8ff8a006a766 Mon Sep 17 00:00:00 2001 From: Roo Code Date: Wed, 16 Jul 2025 19:05:41 +0000 Subject: [PATCH] feat: add prompt caching support for LiteLLM - Add litellmUsePromptCache field to provider settings schema - Update LiteLLM handler to add cache_control metadata when enabled - Add UI checkbox for enabling/disabling prompt caching - Support cache-related token usage tracking (cache reads/writes) - Add translation keys for new UI elements Fixes #5791 --- packages/types/src/provider-settings.ts | 1 + src/api/providers/lite-llm.ts | 56 ++++++++++++++++++- .../components/settings/providers/LiteLLM.tsx | 16 +++++- webview-ui/src/i18n/locales/en/settings.json | 2 + 4 files changed, 71 insertions(+), 4 deletions(-) diff --git a/packages/types/src/provider-settings.ts b/packages/types/src/provider-settings.ts index 3b53627295..f1561188be 100644 --- a/packages/types/src/provider-settings.ts +++ b/packages/types/src/provider-settings.ts @@ -218,6 +218,7 @@ const litellmSchema = baseProviderSettingsSchema.extend({ litellmBaseUrl: z.string().optional(), litellmApiKey: z.string().optional(), litellmModelId: z.string().optional(), + litellmUsePromptCache: z.boolean().optional(), }) const defaultSchema = z.object({ diff --git a/src/api/providers/lite-llm.ts b/src/api/providers/lite-llm.ts index e8cd58b12c..ecc69beb04 100644 --- a/src/api/providers/lite-llm.ts +++ b/src/api/providers/lite-llm.ts @@ -39,11 +39,16 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa ): ApiStream { const { id: modelId, info } = await this.fetchModel() - const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [ + const baseOpenAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [ { role: "system", content: systemPrompt }, ...convertToOpenAiMessages(messages), ] + // Apply prompt caching if enabled + const openAiMessages = this.options.litellmUsePromptCache + ? this.addCacheControlToMessages(baseOpenAiMessages) + : baseOpenAiMessages + // Required by some providers; others default to max tokens allowed let maxTokens: number | undefined = info.maxTokens ?? undefined @@ -80,12 +85,21 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa } if (lastUsage) { + // Extract cache-related information if available + const cacheWriteTokens = + lastUsage.cache_creation_input_tokens || lastUsage.prompt_cache_miss_tokens || 0 + const cacheReadTokens = + lastUsage.cache_read_input_tokens || + lastUsage.prompt_cache_hit_tokens || + lastUsage.prompt_tokens_details?.cached_tokens || + 0 + const usageData: ApiStreamUsageChunk = { type: "usage", inputTokens: lastUsage.prompt_tokens || 0, outputTokens: lastUsage.completion_tokens || 0, - cacheWriteTokens: lastUsage.cache_creation_input_tokens || 0, - cacheReadTokens: lastUsage.prompt_tokens_details?.cached_tokens || 0, + cacheWriteTokens: cacheWriteTokens > 0 ? cacheWriteTokens : undefined, + cacheReadTokens: cacheReadTokens > 0 ? cacheReadTokens : undefined, } usageData.totalCost = calculateApiCostOpenAI( @@ -130,9 +144,45 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa throw error } } + + /** + * Add cache control metadata to messages for prompt caching + * Based on Cline's implementation: adds cache_control to system message and last two user messages + */ + private addCacheControlToMessages( + messages: OpenAI.Chat.ChatCompletionMessageParam[], + ): OpenAI.Chat.ChatCompletionMessageParam[] { + const cacheControl = { cache_control: { type: "ephemeral" } } + + // Find user message indices + const userMsgIndices = messages.reduce( + (acc, msg, index) => (msg.role === "user" ? [...acc, index] : acc), + [] as number[], + ) + + const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1 + const secondLastUserMsgIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1 + + return messages.map((message, index) => { + // Add cache control to system message (first message) + if (index === 0 && message.role === "system") { + return { ...message, ...cacheControl } + } + + // Add cache control to last two user messages + if (index === lastUserMsgIndex || index === secondLastUserMsgIndex) { + return { ...message, ...cacheControl } + } + + return message + }) + } } // LiteLLM usage may include an extra field for Anthropic use cases. interface LiteLLMUsage extends OpenAI.CompletionUsage { cache_creation_input_tokens?: number + prompt_cache_miss_tokens?: number + cache_read_input_tokens?: number + prompt_cache_hit_tokens?: number } diff --git a/webview-ui/src/components/settings/providers/LiteLLM.tsx b/webview-ui/src/components/settings/providers/LiteLLM.tsx index a2467b3c0b..3b7e66ddc1 100644 --- a/webview-ui/src/components/settings/providers/LiteLLM.tsx +++ b/webview-ui/src/components/settings/providers/LiteLLM.tsx @@ -1,5 +1,5 @@ import { useCallback, useState, useEffect, useRef } from "react" -import { VSCodeTextField } from "@vscode/webview-ui-toolkit/react" +import { VSCodeTextField, VSCodeCheckbox } from "@vscode/webview-ui-toolkit/react" import { type ProviderSettings, type OrganizationAllowList, litellmDefaultModelId } from "@roo-code/types" @@ -111,6 +111,20 @@ export const LiteLLM = ({ {t("settings:providers.apiKeyStorageNotice")} +
+ { + const isChecked = e.target.checked === true + setApiConfigurationField("litellmUsePromptCache", isChecked) + }}> + {t("settings:providers.litellmUsePromptCache")} + +
+ {t("settings:providers.litellmUsePromptCacheDescription")} +
+
+