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

1
packages/ai/CHANGELOG.md
View File

@@ -4,6 +4,7 @@
### Added
- Added prompt caching for Mistral requests using the pi session ID as `prompt_cache_key`, including cached-token usage and cost accounting ([#5854](https://github.com/earendil-works/pi/issues/5854)).
- Added the OpenRouter Fusion alias as `openrouter/fusion` ([#5866](https://github.com/earendil-works/pi/pull/5866) by [@dannote](https://github.com/dannote)).
## [0.79.7] - 2026-06-18

5
packages/ai/package.json
View File

@@ -69,9 +69,10 @@
"dependencies": {
"@anthropic-ai/sdk": "0.91.1",
"@aws-sdk/client-bedrock-runtime": "3.1048.0",
"@smithy/node-http-handler": "4.7.3",
"@google/genai": "1.52.0",
"@mistralai/mistralai": "2.2.1",
"@mistralai/mistralai": "2.2.6",
"@opentelemetry/api": "1.9.0",
"@smithy/node-http-handler": "4.7.3",
"http-proxy-agent": "7.0.2",
"https-proxy-agent": "7.0.6",
"openai": "6.26.0",

2
packages/ai/scripts/generate-models.ts
View File

@@ -948,7 +948,7 @@ async function loadModelsDevData(): Promise<Model<any>[]> {
cost: {
input: m.cost?.input || 0,
output: m.cost?.output || 0,
cacheRead: m.cost?.cache_read || 0,
cacheRead: m.cost?.cache_read ?? (m.cost?.input ? roundCost(m.cost.input * 0.1) : 0),
cacheWrite: m.cost?.cache_write || 0,
},
contextWindow: m.limit?.context || 4096,

58
packages/ai/src/models.generated.ts
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,

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

18
packages/ai/test/mistral-reasoning-mode.test.ts
View File

@@ -6,6 +6,7 @@ import type { Context, Model, SimpleStreamOptions } from "../src/types.ts";
interface MistralPayload {
promptMode?: "reasoning";
reasoningEffort?: "none" | "high";
promptCacheKey?: string;
}
function makeContext(): Context {
@@ -77,4 +78,21 @@ describe("Mistral reasoning mode selection", () => {
expect(payload.reasoningEffort).toBeUndefined();
expect(payload.promptMode).toBeUndefined();
});
it("uses the session id as prompt cache key", async () => {
const payload = await capturePayload(getModel("mistral", "mistral-large-latest"), {
sessionId: "session-123",
});
expect(payload.promptCacheKey).toBe("session-123");
});
it("omits prompt cache key when cache retention is disabled", async () => {
const payload = await capturePayload(getModel("mistral", "mistral-large-latest"), {
sessionId: "session-123",
cacheRetention: "none",
});
expect(payload.promptCacheKey).toBeUndefined();
});
});