cache-policy.ts 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  1. // Apply an `LLMRequest.cache` policy by injecting `CacheHint`s onto the parts
  2. // the policy designates. Runs once at compile time, before the per-protocol
  3. // body builder, so the existing inline-hint lowering path handles the rest.
  4. //
  5. // The default `"auto"` shape places one breakpoint at the last tool definition,
  6. // one at the last system part, and one at the latest user message. This
  7. // matches what production agent harnesses (LangChain's caching middleware,
  8. // kern-ai's 10x cost-reduction playbook) converge on for tool-use loops: the
  9. // latest user message stays put while a single turn explodes into many
  10. // assistant/tool round-trips, so caching at that boundary lets every
  11. // intra-turn API call hit the prefix.
  12. //
  13. // Manual `cache: CacheHint` placements on individual parts are preserved —
  14. // this function only fills gaps the caller left empty.
  15. import { CacheHint, type CachePolicy, type CachePolicyObject } from "./schema/options"
  16. import { LLMRequest, Message, ToolDefinition, type ContentPart } from "./schema/messages"
  17. const AUTO: CachePolicyObject = {
  18. tools: true,
  19. system: true,
  20. messages: "latest-user-message",
  21. }
  22. const NONE: CachePolicyObject = {}
  23. // Resolution rules:
  24. // - undefined → "auto" — caching is on by default. The math favors it:
  25. // Anthropic 5m-cache write is 1.25x base, read is 0.1x,
  26. // so a single reuse within 5 minutes already wins.
  27. // - "auto" → tools + system + latest user msg.
  28. // - "none" → no auto placement; manual `CacheHint`s still flow.
  29. // - object form → exactly what the caller asked for.
  30. const resolve = (policy: CachePolicy | undefined): CachePolicyObject => {
  31. if (policy === undefined || policy === "auto") return AUTO
  32. if (policy === "none") return NONE
  33. return policy
  34. }
  35. // Protocols whose wire format ignores inline cache markers (OpenAI's implicit
  36. // prefix caching, Gemini's implicit + out-of-band CachedContent). Skip the
  37. // whole policy pass for these — emitting hints would be harmless but pointless.
  38. const RESPECTS_INLINE_HINTS = new Set(["anthropic-messages", "bedrock-converse"])
  39. const makeHint = (ttlSeconds: number | undefined): CacheHint =>
  40. ttlSeconds !== undefined ? new CacheHint({ type: "ephemeral", ttlSeconds }) : new CacheHint({ type: "ephemeral" })
  41. const markLastTool = (tools: ReadonlyArray<ToolDefinition>, hint: CacheHint): ReadonlyArray<ToolDefinition> => {
  42. if (tools.length === 0) return tools
  43. const last = tools.length - 1
  44. if (tools[last]!.cache) return tools
  45. return tools.map((tool, i) => (i === last ? new ToolDefinition({ ...tool, cache: hint }) : tool))
  46. }
  47. const markLastSystem = (system: LLMRequest["system"], hint: CacheHint): LLMRequest["system"] => {
  48. if (system.length === 0) return system
  49. const last = system.length - 1
  50. if (system[last]!.cache) return system
  51. return system.map((part, i) => (i === last ? { ...part, cache: hint } : part))
  52. }
  53. const lastIndexOfRole = (messages: ReadonlyArray<Message>, role: Message["role"]): number =>
  54. messages.findLastIndex((m) => m.role === role)
  55. // Mark the last text part of `messages[index]`. If no text part exists, mark
  56. // the last content part regardless of type — that's the breakpoint position
  57. // in tool-result-only messages too.
  58. const markMessageAt = (messages: ReadonlyArray<Message>, index: number, hint: CacheHint): ReadonlyArray<Message> => {
  59. if (index < 0 || index >= messages.length) return messages
  60. const target = messages[index]!
  61. if (target.content.length === 0) return messages
  62. const lastTextIndex = target.content.findLastIndex((part) => part.type === "text")
  63. const markAt = lastTextIndex >= 0 ? lastTextIndex : target.content.length - 1
  64. const existing = target.content[markAt]!
  65. if ("cache" in existing && existing.cache) return messages
  66. const nextContent = target.content.map((part, i) => (i === markAt ? ({ ...part, cache: hint } as ContentPart) : part))
  67. const next = new Message({ ...target, content: nextContent })
  68. // Single pass over `messages`, substituting the one updated entry. Long
  69. // conversations call this on every request, so avoid `.map()` here — its
  70. // closure dispatch and identity copies show up in profiling.
  71. const result = messages.slice()
  72. result[index] = next
  73. return result
  74. }
  75. const markMessages = (
  76. messages: ReadonlyArray<Message>,
  77. strategy: NonNullable<CachePolicyObject["messages"]>,
  78. hint: CacheHint,
  79. ): ReadonlyArray<Message> => {
  80. if (messages.length === 0) return messages
  81. if (strategy === "latest-user-message") return markMessageAt(messages, lastIndexOfRole(messages, "user"), hint)
  82. if (strategy === "latest-assistant") return markMessageAt(messages, lastIndexOfRole(messages, "assistant"), hint)
  83. const start = Math.max(0, messages.length - strategy.tail)
  84. let next = messages
  85. for (let i = start; i < messages.length; i++) next = markMessageAt(next, i, hint)
  86. return next
  87. }
  88. export const applyCachePolicy = (request: LLMRequest): LLMRequest => {
  89. if (!RESPECTS_INLINE_HINTS.has(request.model.route.id)) return request
  90. const policy = resolve(request.cache)
  91. if (!policy.tools && !policy.system && !policy.messages) return request
  92. const hint = makeHint(policy.ttlSeconds)
  93. const tools = policy.tools ? markLastTool(request.tools, hint) : request.tools
  94. const system = policy.system ? markLastSystem(request.system, hint) : request.system
  95. const messages = policy.messages ? markMessages(request.messages, policy.messages, hint) : request.messages
  96. if (tools === request.tools && system === request.system && messages === request.messages) return request
  97. return LLMRequest.update(request, { tools, system, messages })
  98. }