shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ai): enable Mistral prompt caching

Browse Source
This commit is contained in:
Armin Ronacher
2026-06-18 23:47:25 +02:00
parent 8025fdd01f
commit 651d10d90b
8 changed files with 152 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
packages/ai/src/providers/mistral.ts
View File

@@ -226,7 +226,7 @@ function buildRequestOptions(model: Model<"mistral-conversations">, options?: Mi
// Mistral infrastructure uses `x-affinity` for KV-cache reuse (prefix caching).
// Respect explicit caller-provided header values.
if (options?.sessionId && !headers["x-affinity"]) {
if (shouldUsePromptCaching(options) && !headers["x-affinity"]) {
headers["x-affinity"] = options.sessionId;
}
@@ -255,6 +255,7 @@ function buildChatPayload(
if (options?.toolChoice) payload.toolChoice = mapToolChoice(options.toolChoice);
if (options?.promptMode) payload.promptMode = options.promptMode;
if (options?.reasoningEffort) payload.reasoningEffort = options.reasoningEffort;
if (shouldUsePromptCaching(options)) payload.promptCacheKey = options.sessionId;
if (context.systemPrompt) {
payload.messages.unshift({
@@ -266,6 +267,31 @@ function buildChatPayload(
return payload;
}
function shouldUsePromptCaching(options?: MistralOptions): options is MistralOptions & { sessionId: string } {
return options?.cacheRetention !== "none" && !!options?.sessionId;
}
function getMistralCachedPromptTokens(usage: unknown, promptTokens: number): number {
const rawUsage = usage as {
promptTokensDetails?: { cachedTokens?: unknown } | null;
prompt_tokens_details?: { cached_tokens?: unknown } | null;
promptTokenDetails?: { cachedTokens?: unknown } | null;
prompt_token_details?: { cached_tokens?: unknown } | null;
numCachedTokens?: unknown;
num_cached_tokens?: unknown;
};
const rawCachedTokens =
rawUsage.promptTokensDetails?.cachedTokens ??
rawUsage.prompt_tokens_details?.cached_tokens ??
rawUsage.promptTokenDetails?.cachedTokens ??
rawUsage.prompt_token_details?.cached_tokens ??
rawUsage.numCachedTokens ??
rawUsage.num_cached_tokens ??
0;
const cachedTokens = typeof rawCachedTokens === "number" && Number.isFinite(rawCachedTokens) ? rawCachedTokens : 0;
return Math.min(promptTokens, Math.max(0, cachedTokens));
}
async function consumeChatStream(
model: Model<"mistral-conversations">,
output: AssistantMessage,
@@ -305,11 +331,16 @@ async function consumeChatStream(
output.responseId ||= chunk.id;
if (chunk.usage) {
output.usage.input = chunk.usage.promptTokens || 0;
const promptTokens = chunk.usage.promptTokens || 0;
const cachedPromptTokens = getMistralCachedPromptTokens(chunk.usage, promptTokens);
output.usage.input = Math.max(0, promptTokens - cachedPromptTokens);
output.usage.output = chunk.usage.completionTokens || 0;
output.usage.cacheRead = 0;
output.usage.cacheRead = cachedPromptTokens;
output.usage.cacheWrite = 0;
output.usage.totalTokens = chunk.usage.totalTokens || output.usage.input + output.usage.output;
output.usage.totalTokens =
chunk.usage.totalTokens ||
output.usage.input + output.usage.output + output.usage.cacheRead + output.usage.cacheWrite;
calculateCost(model, output.usage);
}