fix(ai): layer prompt cache breakpoints (#38725)
This commit is contained in:
parent
9b640cf97d
commit
993f046dd9
8 changed files with 258 additions and 43 deletions
|
|
@ -2,32 +2,31 @@
|
|||
// the policy designates. Runs once at compile time, before the per-protocol
|
||||
// body builder, so the existing inline-hint lowering path handles the rest.
|
||||
//
|
||||
// The default `"auto"` shape places one breakpoint at the last tool definition,
|
||||
// one at the last system part, and one at the latest user message. This
|
||||
// matches what production agent harnesses (LangChain's caching middleware,
|
||||
// kern-ai's 10x cost-reduction playbook) converge on for tool-use loops: the
|
||||
// latest user message stays put while a single turn explodes into many
|
||||
// assistant/tool round-trips, so caching at that boundary lets every
|
||||
// intra-turn API call hit the prefix.
|
||||
// The default `"auto"` shape places breakpoints at the last tool definition,
|
||||
// the first and last distinct system parts, and the conversation tail. This
|
||||
// exposes reusable tool, base-agent, project, and session prefixes while
|
||||
// advancing the tail after each tool result keeps the previous cache entry
|
||||
// within Anthropic's 20-block lookback during long agent turns.
|
||||
//
|
||||
// Manual `cache: CacheHint` placements on individual parts are preserved —
|
||||
// this function only fills gaps the caller left empty.
|
||||
// Manual `cache: CacheHint` placements on individual parts are preserved and
|
||||
// count against the four-breakpoint budget; auto only fills remaining slots.
|
||||
import { CacheHint, type CachePolicy, type CachePolicyObject } from "./schema/options"
|
||||
import { LLMRequest, Message, ToolDefinition, type ContentPart } from "./schema/messages"
|
||||
|
||||
const AUTO: CachePolicyObject = {
|
||||
tools: true,
|
||||
system: true,
|
||||
messages: "latest-user-message",
|
||||
messages: { tail: 1 },
|
||||
}
|
||||
|
||||
const NONE: CachePolicyObject = {}
|
||||
const BREAKPOINT_CAP = 4
|
||||
|
||||
// Resolution rules:
|
||||
// - undefined → "auto" — caching is on by default. The math favors it:
|
||||
// Anthropic 5m-cache write is 1.25x base, read is 0.1x,
|
||||
// so a single reuse within 5 minutes already wins.
|
||||
// - "auto" → tools + system + latest user msg.
|
||||
// - "auto" → tools + first/last system + final message boundary.
|
||||
// - "none" → no auto placement; manual `CacheHint`s still flow.
|
||||
// - object form → exactly what the caller asked for.
|
||||
const resolve = (policy: CachePolicy | undefined): CachePolicyObject => {
|
||||
|
|
@ -44,18 +43,32 @@ const RESPECTS_INLINE_HINTS = new Set(["anthropic-messages", "bedrock-converse"]
|
|||
const makeHint = (ttlSeconds: number | undefined): CacheHint =>
|
||||
ttlSeconds !== undefined ? new CacheHint({ type: "ephemeral", ttlSeconds }) : new CacheHint({ type: "ephemeral" })
|
||||
|
||||
const markLastTool = (tools: ReadonlyArray<ToolDefinition>, hint: CacheHint): ReadonlyArray<ToolDefinition> => {
|
||||
interface Budget {
|
||||
remaining: number
|
||||
}
|
||||
|
||||
const markLastTool = (
|
||||
tools: ReadonlyArray<ToolDefinition>,
|
||||
hint: CacheHint,
|
||||
budget: Budget,
|
||||
): ReadonlyArray<ToolDefinition> => {
|
||||
if (tools.length === 0) return tools
|
||||
const last = tools.length - 1
|
||||
if (tools[last]!.cache) return tools
|
||||
if (tools[last]!.cache || budget.remaining === 0) return tools
|
||||
budget.remaining -= 1
|
||||
return tools.map((tool, i) => (i === last ? new ToolDefinition({ ...tool, cache: hint }) : tool))
|
||||
}
|
||||
|
||||
const markLastSystem = (system: LLMRequest["system"], hint: CacheHint): LLMRequest["system"] => {
|
||||
const markSystemBoundaries = (system: LLMRequest["system"], hint: CacheHint, budget: Budget): LLMRequest["system"] => {
|
||||
if (system.length === 0) return system
|
||||
const last = system.length - 1
|
||||
if (system[last]!.cache) return system
|
||||
return system.map((part, i) => (i === last ? { ...part, cache: hint } : part))
|
||||
let changed = false
|
||||
const next = system.map((part, index) => {
|
||||
if ((index !== 0 && index !== system.length - 1) || part.cache || budget.remaining === 0) return part
|
||||
budget.remaining -= 1
|
||||
changed = true
|
||||
return { ...part, cache: hint }
|
||||
})
|
||||
return changed ? next : system
|
||||
}
|
||||
|
||||
const lastIndexOfRole = (messages: ReadonlyArray<Message>, role: Message["role"]): number =>
|
||||
|
|
@ -64,14 +77,20 @@ const lastIndexOfRole = (messages: ReadonlyArray<Message>, role: Message["role"]
|
|||
// Mark the last text part of `messages[index]`. If no text part exists, mark
|
||||
// the last content part regardless of type — that's the breakpoint position
|
||||
// in tool-result-only messages too.
|
||||
const markMessageAt = (messages: ReadonlyArray<Message>, index: number, hint: CacheHint): ReadonlyArray<Message> => {
|
||||
const markMessageAt = (
|
||||
messages: ReadonlyArray<Message>,
|
||||
index: number,
|
||||
hint: CacheHint,
|
||||
budget: Budget,
|
||||
): ReadonlyArray<Message> => {
|
||||
if (index < 0 || index >= messages.length) return messages
|
||||
const target = messages[index]!
|
||||
if (target.content.length === 0) return messages
|
||||
const lastTextIndex = target.content.findLastIndex((part) => part.type === "text")
|
||||
const markAt = lastTextIndex >= 0 ? lastTextIndex : target.content.length - 1
|
||||
const existing = target.content[markAt]!
|
||||
if ("cache" in existing && existing.cache) return messages
|
||||
if (("cache" in existing && existing.cache) || budget.remaining === 0) return messages
|
||||
budget.remaining -= 1
|
||||
const nextContent = target.content.map((part, i) => (i === markAt ? ({ ...part, cache: hint } as ContentPart) : part))
|
||||
const next = new Message({ ...target, content: nextContent })
|
||||
// Single pass over `messages`, substituting the one updated entry. Long
|
||||
|
|
@ -86,25 +105,42 @@ const markMessages = (
|
|||
messages: ReadonlyArray<Message>,
|
||||
strategy: NonNullable<CachePolicyObject["messages"]>,
|
||||
hint: CacheHint,
|
||||
budget: Budget,
|
||||
): ReadonlyArray<Message> => {
|
||||
if (messages.length === 0) return messages
|
||||
if (strategy === "latest-user-message") return markMessageAt(messages, lastIndexOfRole(messages, "user"), hint)
|
||||
if (strategy === "latest-assistant") return markMessageAt(messages, lastIndexOfRole(messages, "assistant"), hint)
|
||||
if (strategy === "latest-user-message")
|
||||
return markMessageAt(messages, lastIndexOfRole(messages, "user"), hint, budget)
|
||||
if (strategy === "latest-assistant")
|
||||
return markMessageAt(messages, lastIndexOfRole(messages, "assistant"), hint, budget)
|
||||
const start = Math.max(0, messages.length - strategy.tail)
|
||||
let next = messages
|
||||
for (let i = start; i < messages.length; i++) next = markMessageAt(next, i, hint)
|
||||
for (let i = start; i < messages.length; i++) next = markMessageAt(next, i, hint, budget)
|
||||
return next
|
||||
}
|
||||
|
||||
const countHints = (request: LLMRequest) =>
|
||||
request.tools.reduce((count, tool) => count + (tool.cache === undefined ? 0 : 1), 0) +
|
||||
request.system.reduce((count, part) => count + (part.cache === undefined ? 0 : 1), 0) +
|
||||
request.messages.reduce(
|
||||
(count, message) =>
|
||||
count +
|
||||
message.content.reduce(
|
||||
(contentCount, part) => contentCount + ("cache" in part && part.cache !== undefined ? 1 : 0),
|
||||
0,
|
||||
),
|
||||
0,
|
||||
)
|
||||
|
||||
export const applyCachePolicy = (request: LLMRequest): LLMRequest => {
|
||||
if (!RESPECTS_INLINE_HINTS.has(request.model.route.id)) return request
|
||||
const policy = resolve(request.cache)
|
||||
if (!policy.tools && !policy.system && !policy.messages) return request
|
||||
|
||||
const hint = makeHint(policy.ttlSeconds)
|
||||
const tools = policy.tools ? markLastTool(request.tools, hint) : request.tools
|
||||
const system = policy.system ? markLastSystem(request.system, hint) : request.system
|
||||
const messages = policy.messages ? markMessages(request.messages, policy.messages, hint) : request.messages
|
||||
const budget = { remaining: Math.max(0, BREAKPOINT_CAP - countHints(request)) }
|
||||
const tools = policy.tools ? markLastTool(request.tools, hint, budget) : request.tools
|
||||
const system = policy.system ? markSystemBoundaries(request.system, hint, budget) : request.system
|
||||
const messages = policy.messages ? markMessages(request.messages, policy.messages, hint, budget) : request.messages
|
||||
|
||||
if (tools === request.tools && system === request.system && messages === request.messages) return request
|
||||
return LLMRequest.update(request, { tools, system, messages })
|
||||
|
|
|
|||
|
|
@ -251,11 +251,11 @@ export class CacheHint extends Schema.Class<CacheHint>("LLM.CacheHint")({
|
|||
// Auto-placement policy for prompt caching. The protocol-neutral lowering step
|
||||
// reads this and injects `CacheHint`s at the configured boundaries; the
|
||||
// per-protocol body builders then translate those hints into wire markers as
|
||||
// usual. `"auto"` is the recommended default for agent loops — it places one
|
||||
// breakpoint at the last tool definition, one at the last system part, and one
|
||||
// at the latest user message. The combination of provider invalidation
|
||||
// hierarchy (tools → system → messages) and Anthropic/Bedrock's 20-block
|
||||
// lookback means three trailing breakpoints reliably cover the static prefix.
|
||||
// usual. `"auto"` is the recommended default for agent loops — it places
|
||||
// breakpoints at the last tool definition, the first and last distinct system
|
||||
// parts, and the conversation tail. The rolling message breakpoint keeps a
|
||||
// prior cache entry within Anthropic/Bedrock's 20-block lookback during long
|
||||
// tool loops.
|
||||
//
|
||||
// Pass `"none"` to opt out entirely (the legacy behavior). Pass the granular
|
||||
// object form to override individual choices.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue