feat(ai): enable Mistral prompt caching

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

View File

@@ -5761,7 +5761,7 @@ export const MODELS = {
cost: {
input: 0.3,
output: 0.9,
cacheRead: 0,
cacheRead: 0.03,
cacheWrite: 0,
},
contextWindow: 256000,
@@ -5778,7 +5778,7 @@ export const MODELS = {
cost: {
input: 0.4,
output: 2,
cacheRead: 0,
cacheRead: 0.04,
cacheWrite: 0,
},
contextWindow: 262144,
@@ -5795,7 +5795,7 @@ export const MODELS = {
cost: {
input: 0.4,
output: 2,
cacheRead: 0,
cacheRead: 0.04,
cacheWrite: 0,
},
contextWindow: 262144,
@@ -5812,7 +5812,7 @@ export const MODELS = {
cost: {
input: 0.4,
output: 2,
cacheRead: 0,
cacheRead: 0.04,
cacheWrite: 0,
},
contextWindow: 128000,
@@ -5829,7 +5829,7 @@ export const MODELS = {
cost: {
input: 0.4,
output: 2,
cacheRead: 0,
cacheRead: 0.04,
cacheWrite: 0,
},
contextWindow: 262144,
@@ -5846,7 +5846,7 @@ export const MODELS = {
cost: {
input: 0.1,
output: 0.3,
cacheRead: 0,
cacheRead: 0.01,
cacheWrite: 0,
},
contextWindow: 128000,
@@ -5863,7 +5863,7 @@ export const MODELS = {
cost: {
input: 0.1,
output: 0.3,
cacheRead: 0,
cacheRead: 0.01,
cacheWrite: 0,
},
contextWindow: 128000,
@@ -5897,7 +5897,7 @@ export const MODELS = {
cost: {
input: 2,
output: 5,
cacheRead: 0,
cacheRead: 0.2,
cacheWrite: 0,
},
contextWindow: 128000,
@@ -5914,7 +5914,7 @@ export const MODELS = {
cost: {
input: 0.5,
output: 1.5,
cacheRead: 0,
cacheRead: 0.05,
cacheWrite: 0,
},
contextWindow: 128000,
@@ -5931,7 +5931,7 @@ export const MODELS = {
cost: {
input: 0.04,
output: 0.04,
cacheRead: 0,
cacheRead: 0.004,
cacheWrite: 0,
},
contextWindow: 128000,
@@ -5948,7 +5948,7 @@ export const MODELS = {
cost: {
input: 0.1,
output: 0.1,
cacheRead: 0,
cacheRead: 0.01,
cacheWrite: 0,
},
contextWindow: 128000,
@@ -5965,7 +5965,7 @@ export const MODELS = {
cost: {
input: 2,
output: 6,
cacheRead: 0,
cacheRead: 0.2,
cacheWrite: 0,
},
contextWindow: 131072,
@@ -5982,7 +5982,7 @@ export const MODELS = {
cost: {
input: 0.5,
output: 1.5,
cacheRead: 0,
cacheRead: 0.05,
cacheWrite: 0,
},
contextWindow: 262144,
@@ -5999,7 +5999,7 @@ export const MODELS = {
cost: {
input: 0.5,
output: 1.5,
cacheRead: 0,
cacheRead: 0.05,
cacheWrite: 0,
},
contextWindow: 262144,
@@ -6016,7 +6016,7 @@ export const MODELS = {
cost: {
input: 0.4,
output: 2,
cacheRead: 0,
cacheRead: 0.04,
cacheWrite: 0,
},
contextWindow: 131072,
@@ -6033,7 +6033,7 @@ export const MODELS = {
cost: {
input: 0.4,
output: 2,
cacheRead: 0,
cacheRead: 0.04,
cacheWrite: 0,
},
contextWindow: 262144,
@@ -6050,7 +6050,7 @@ export const MODELS = {
cost: {
input: 1.5,
output: 7.5,
cacheRead: 0,
cacheRead: 0.15,
cacheWrite: 0,
},
contextWindow: 262144,
@@ -6067,7 +6067,7 @@ export const MODELS = {
cost: {
input: 1.5,
output: 7.5,
cacheRead: 0,
cacheRead: 0.15,
cacheWrite: 0,
},
contextWindow: 262144,
@@ -6084,7 +6084,7 @@ export const MODELS = {
cost: {
input: 0.4,
output: 2,
cacheRead: 0,
cacheRead: 0.04,
cacheWrite: 0,
},
contextWindow: 262144,
@@ -6101,7 +6101,7 @@ export const MODELS = {
cost: {
input: 0.15,
output: 0.15,
cacheRead: 0,
cacheRead: 0.015,
cacheWrite: 0,
},
contextWindow: 128000,
@@ -6118,7 +6118,7 @@ export const MODELS = {
cost: {
input: 0.1,
output: 0.3,
cacheRead: 0,
cacheRead: 0.01,
cacheWrite: 0,
},
contextWindow: 128000,
@@ -6135,7 +6135,7 @@ export const MODELS = {
cost: {
input: 0.15,
output: 0.6,
cacheRead: 0,
cacheRead: 0.015,
cacheWrite: 0,
},
contextWindow: 256000,
@@ -6152,7 +6152,7 @@ export const MODELS = {
cost: {
input: 0.15,
output: 0.6,
cacheRead: 0,
cacheRead: 0.015,
cacheWrite: 0,
},
contextWindow: 256000,
@@ -6169,7 +6169,7 @@ export const MODELS = {
cost: {
input: 0.25,
output: 0.25,
cacheRead: 0,
cacheRead: 0.025,
cacheWrite: 0,
},
contextWindow: 8000,
@@ -6186,7 +6186,7 @@ export const MODELS = {
cost: {
input: 0.15,
output: 0.15,
cacheRead: 0,
cacheRead: 0.015,
cacheWrite: 0,
},
contextWindow: 128000,
@@ -6203,7 +6203,7 @@ export const MODELS = {
cost: {
input: 2,
output: 6,
cacheRead: 0,
cacheRead: 0.2,
cacheWrite: 0,
},
contextWindow: 64000,
@@ -6220,7 +6220,7 @@ export const MODELS = {
cost: {
input: 0.7,
output: 0.7,
cacheRead: 0,
cacheRead: 0.07,
cacheWrite: 0,
},
contextWindow: 32000,
@@ -6237,7 +6237,7 @@ export const MODELS = {
cost: {
input: 0.15,
output: 0.15,
cacheRead: 0,
cacheRead: 0.015,
cacheWrite: 0,
},
contextWindow: 128000,
@@ -6254,7 +6254,7 @@ export const MODELS = {
cost: {
input: 2,
output: 6,
cacheRead: 0,
cacheRead: 0.2,
cacheWrite: 0,
},
contextWindow: 128000,

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