diff --git a/packages/core/src/conversations/turns/context.test.ts b/packages/core/src/conversations/turns/context.test.ts index 3ded2a3e..e5da38ac 100644 --- a/packages/core/src/conversations/turns/context.test.ts +++ b/packages/core/src/conversations/turns/context.test.ts @@ -21,12 +21,17 @@ describe("modelPrompt", () => { "You are Host Agent (@host-agent), an agent in pod Release of workspace Suga", ); expect(prompt.system).toContain("Check facts carefully."); - // The history, then one trailing instruction with everything that varies - // between turns, so the prefix the provider caches never changes shape. + // The history, marked where it ends, then one trailing instruction with + // everything that varies between turns, so the prefix the provider caches + // never changes shape. expect(prompt.messages.slice(0, -1)).toEqual([ { role: "user", content: "Sam (@sam): Check the release" }, { role: "assistant", content: "I am checking it." }, - { role: "user", content: "Reviewer (@reviewer) (agent): Migration is present." }, + { + role: "user", + content: "Reviewer (@reviewer) (agent): Migration is present.", + cacheBreakpoint: true, + }, ]); const instruction = prompt.messages.at(-1); expect(instruction?.role).toBe("user"); diff --git a/packages/core/src/conversations/turns/context.ts b/packages/core/src/conversations/turns/context.ts index ca217622..31d672d9 100644 --- a/packages/core/src/conversations/turns/context.ts +++ b/packages/core/src/conversations/turns/context.ts @@ -36,7 +36,9 @@ export interface TurnEnvironment { * thread now, why this agent has the floor, the turn's environment) goes in one * trailing instruction after the history, where a change costs nothing. A * compacted thread starts its history with the Compaction agent's summary, which - * changes only when the thread is compacted again. + * changes only when the thread is compacted again. The history's last message + * is marked as the end of that shared prefix, for providers that cache only + * what a request marks. */ export function modelPrompt( context: TurnContext, @@ -65,13 +67,19 @@ export function modelPrompt( .filter(Boolean) .join("\n\n"); + const history = [ + ...(context.compaction ? [compactionSummary(context.compaction)] : []), + ...context.messages + .filter((message) => message.status === "complete") + .flatMap((message) => promptMessagesFor(message, context.agent.id, resultFiles)), + ]; + const lastInHistory = history.at(-1); return { system, messages: [ - ...(context.compaction ? [compactionSummary(context.compaction)] : []), - ...context.messages - .filter((message) => message.status === "complete") - .flatMap((message) => promptMessagesFor(message, context.agent.id, resultFiles)), + ...(lastInHistory + ? history.with(-1, { ...lastInHistory, cacheBreakpoint: true as const }) + : history), { role: "user" as const, content: turnInstruction(context, environment) }, ], }; diff --git a/packages/core/src/conversations/turns/repository.ts b/packages/core/src/conversations/turns/repository.ts index 032e6a9e..eca93374 100644 --- a/packages/core/src/conversations/turns/repository.ts +++ b/packages/core/src/conversations/turns/repository.ts @@ -810,7 +810,11 @@ export const TurnCheckpoint = Schema.Struct({ model: Schema.String, system: Schema.String, messages: Schema.Array( - Schema.Struct({ role: Schema.Literals(["user", "assistant"]), content: Schema.String }), + Schema.Struct({ + role: Schema.Literals(["user", "assistant"]), + content: Schema.String, + cacheBreakpoint: Schema.optional(Schema.Literal(true)), + }), ), }), reply: ReplyDraft, diff --git a/packages/core/src/providers/model-providers/dialects/anthropic.ts b/packages/core/src/providers/model-providers/dialects/anthropic.ts index 6e867f3b..ea468707 100644 --- a/packages/core/src/providers/model-providers/dialects/anthropic.ts +++ b/packages/core/src/providers/model-providers/dialects/anthropic.ts @@ -9,13 +9,15 @@ import { /** * Anthropic keys with a header of its own and pages its listing, twenty models * at a time unless asked for more. The listing has no capabilities, so the - * registry supplies them. + * registry supplies them. Its models cache only the prompt prefixes a + * request marks. */ export const anthropic: ProviderDialect = { name: "Anthropic", discoveryRoot: (baseUrl) => (baseUrl.endsWith("/v1") ? baseUrl : `${baseUrl}/v1`), listingUrl: (root) => `${root}/models?limit=1000`, authorization: (apiKey) => ({ "x-api-key": apiKey ?? "", "anthropic-version": "2023-06-01" }), + cacheBreakpoint: () => ({ anthropic: { cacheControl: { type: "ephemeral" } } }), listing: dataListing, model: Schema.Struct({ id: Schema.String, diff --git a/packages/core/src/providers/model-providers/dialects/dialect.ts b/packages/core/src/providers/model-providers/dialects/dialect.ts index 08308472..3cfc624e 100644 --- a/packages/core/src/providers/model-providers/dialects/dialect.ts +++ b/packages/core/src/providers/model-providers/dialects/dialect.ts @@ -5,6 +5,7 @@ import { type ProviderPresetId, providerModelCapabilitySchema, } from "@sugabots/contracts"; +import type { ModelMessage } from "ai"; import { Data, Effect, Schema, SchemaGetter } from "effect"; import type { EgressHttpClient } from "../../network/egress.ts"; @@ -19,8 +20,9 @@ import type { EgressHttpClient } from "../../network/egress.ts"; * to a provider never asks which one it is. * * Discovery, how a provider lists its models and what it says about them, is - * what a dialect covers so far. A header a provider wants on every request, - * or an option its completions take, belongs here too as the need arrives. + * most of what a dialect covers; the rest is the options its requests take. A + * header a provider wants on every request belongs here too as the need + * arrives. */ export interface DiscoveredModel { @@ -41,6 +43,9 @@ export interface ProviderIdentity { apiFormat: ProviderApiFormat; } +/** Options for one provider's SDK, keyed by the provider they are for. */ +export type ProviderOptions = NonNullable; + export interface ProviderDialect { /** Named in tests and in what a person is told. */ readonly name: string; @@ -69,6 +74,13 @@ export interface ProviderDialect { model: DiscoveredModel, http: EgressHttpClient, ): Effect.Effect; + /** + * The options that mark the end of a prompt prefix later requests repeat, + * for a model that caches only the prefixes a request marks. Absent, or + * undefined for a model, where the provider caches prefixes by itself or + * not at all. + */ + cacheBreakpoint?(modelId: string): ProviderOptions | undefined; } export class ModelInspectionFailed extends Data.TaggedError("ModelInspectionFailed")<{ diff --git a/packages/core/src/providers/model-providers/dialects/openrouter.ts b/packages/core/src/providers/model-providers/dialects/openrouter.ts index 7fd7ec64..f10bf985 100644 --- a/packages/core/src/providers/model-providers/dialects/openrouter.ts +++ b/packages/core/src/providers/model-providers/dialects/openrouter.ts @@ -13,15 +13,24 @@ const names = Schema.Array(Schema.String).pipe( Schema.catchDecoding(() => Effect.succeedSome([])), ); +/** How OpenRouter's ids for Anthropic's models start, as in `anthropic/claude-sonnet-4.5`. */ +const ANTHROPIC_MODELS = "anthropic/"; + /** * OpenRouter's listing says what each model takes in, gives out, and accepts - * as parameters, so nothing is left to the registry. + * as parameters, so nothing is left to the registry. Its Anthropic models, + * like Anthropic's own, cache only the prompt prefixes a request marks; + * the rest cache by themselves or not at all. */ export const openrouter: ProviderDialect = { name: "OpenRouter", // The public /models endpoint accepts invalid keys; /models/user requires authentication. listingUrl: (root) => `${root}/models/user`, authorization: openaiCompatible.authorization, + cacheBreakpoint: (modelId) => + modelId.startsWith(ANTHROPIC_MODELS) + ? { openrouter: { cacheControl: { type: "ephemeral" } } } + : undefined, listing: dataListing, model: Schema.Struct({ id: Schema.String, diff --git a/packages/core/src/providers/models/models.recording.test.ts b/packages/core/src/providers/models/models.recording.test.ts index c84ac8a4..4486a659 100644 --- a/packages/core/src/providers/models/models.recording.test.ts +++ b/packages/core/src/providers/models/models.recording.test.ts @@ -4,6 +4,7 @@ import { describe, expect, it } from "vitest"; import type { ModelRequests } from "../../accounting/model-requests.ts"; import { effectRunner } from "../../database/database.ts"; import { noDatabase } from "../../database/testing.ts"; +import type { ProviderIdentity } from "../model-providers/dialects/index.ts"; import { Models } from "./models.ts"; /** The fake provider never queries, so nothing here reaches the database. */ @@ -56,8 +57,17 @@ const reply = () => type Recorded = ModelRequests.Started & Partial> & { outcome: ModelRequests.Outcome }; -/** A model over an OpenAI provider that answers with `responses`, one per request. */ -function modelAnswering(responses: Array<(recorded: readonly Recorded[]) => Response>) { +const OPENAI = { + preset: "openai", + baseUrl: "https://api.openai.com/v1", + apiFormat: "openai", +} satisfies ProviderIdentity; + +/** A model over `provider`, OpenAI's unless given, that answers with `responses`, one per request. */ +function modelAnswering( + responses: Array<(recorded: readonly Recorded[]) => Response>, + provider: ProviderIdentity = OPENAI, +) { const recorded: Recorded[] = []; /** Each request's body, as the provider was sent it. */ const sent: Array> = []; @@ -68,9 +78,7 @@ function modelAnswering(responses: Array<(recorded: readonly Recorded[]) => Resp resolve: () => Effect.succeed({ providerId: PROVIDER_ID, - preset: "openai", - baseUrl: "https://api.openai.com/v1", - apiFormat: "openai", + ...provider, apiKey: "secret", headers: {}, configurationUpdatedAt: new Date(0), @@ -208,3 +216,75 @@ describe("a response at its step limit", () => { ]); }); }); + +describe("prompt caching", () => { + const marked = { type: "ephemeral" }; + const turn: Models.StreamRequest = { + ...input, + messages: [ + { role: "user", content: "Earlier." }, + { role: "assistant", content: "Noted.", cacheBreakpoint: true }, + { role: "user", content: "[Turn] Look it up." }, + ], + }; + + it("marks the system text, the history's end and each step's newest message for Anthropic", async () => { + const { model, sent } = modelAnswering([() => new Response("Unavailable", { status: 503 })], { + preset: "anthropic", + baseUrl: "https://api.anthropic.com", + apiFormat: "anthropic", + }); + + await answer(model, { ...turn, model: "claude-test" }); + + expect(sent[0]).toMatchObject({ + system: [{ text: input.system, cache_control: marked }], + messages: [ + { role: "user", content: [{ text: "Earlier." }] }, + { role: "assistant", content: [{ text: "Noted.", cache_control: marked }] }, + { role: "user", content: [{ text: "[Turn] Look it up.", cache_control: marked }] }, + ], + }); + }); + + it("moves the newest message's mark on each step, keeping the history's", async () => { + const { model, sent } = modelAnswering([toolCall, reply], { + preset: "openrouter", + baseUrl: "https://openrouter.ai/api/v1", + apiFormat: "openai", + }); + + await answer(model, { ...turn, model: "anthropic/claude-test" }); + + /** Which of a request's messages are marked, on the message or on one of its parts. */ + const markedMessages = (body: Record = {}) => + (body.messages as Array<{ content: unknown; cache_control?: unknown }>).map( + (message) => + message.cache_control !== undefined || + (Array.isArray(message.content) && + message.content.some((part) => part.cache_control !== undefined)), + ); + // System text, history, turn. + expect(markedMessages(sent[0])).toEqual([true, false, true, true]); + // The turn's mark has moved to the tool result. + expect(markedMessages(sent[1])).toEqual([true, false, true, false, false, true]); + }); + + it.each([ + ["OpenAI", OPENAI, "gpt-test"], + [ + "OpenRouter's other models", + { preset: "openrouter", baseUrl: "https://openrouter.ai/api/v1", apiFormat: "openai" }, + "openai/gpt-test", + ], + ] satisfies Array<[string, ProviderIdentity, string]>)( + "leaves requests to %s unmarked, since they cache by themselves", + async (_name, provider, modelId) => { + const { model, sent } = modelAnswering([toolCall, reply], provider); + + await answer(model, { ...turn, model: modelId }); + + expect(JSON.stringify(sent)).not.toContain("cache_control"); + }, + ); +}); diff --git a/packages/core/src/providers/models/models.ts b/packages/core/src/providers/models/models.ts index 0c09c00f..7381ad84 100644 --- a/packages/core/src/providers/models/models.ts +++ b/packages/core/src/providers/models/models.ts @@ -27,7 +27,8 @@ import { ModelRequests } from "../../accounting/model-requests.ts"; import { streamLedger } from "../../accounting/stream-ledger.ts"; import type { Database } from "../../database/database.ts"; import { type UserFacing, UserMessage } from "../../user-message.ts"; -import { type ModelRegistry, modelsDev } from "../model-providers/dialects/index.ts"; +import type { ProviderOptions } from "../model-providers/dialects/dialect.ts"; +import { dialectFor, type ModelRegistry, modelsDev } from "../model-providers/dialects/index.ts"; import { ModelProviderRepository } from "../model-providers/model-provider-repository.ts"; import { withSignInAccess } from "../model-providers/sign-in/sign-in.ts"; import { Egress, type EgressHttpClients } from "../network/egress.ts"; @@ -43,6 +44,13 @@ export interface Prompt { export interface PromptMessage { role: "user" | "assistant"; content: string; + /** + * Whether later requests repeat the prompt up to and including this + * message, as a turn's history is repeated by the thread's next turn. A + * provider that caches only the prefixes a request marks is told to cache + * it; one that caches by itself needs no telling. + */ + cacheBreakpoint?: true; } /** A response streamed as it arrives, with tools the model may call along the way. */ @@ -326,6 +334,8 @@ export function make({ modelProviders, httpClients, requests, registry }: Option connection, }); const codex = connection.preset === "chatgpt"; + const breakpoint = dialectFor(connection).cacheBreakpoint?.(input.model); + const markNewest = breakpoint ? newestMessageBreakpoint(breakpoint) : undefined; // The SDK does not throw a provider's error into the text stream: it // reports it here and ends the stream, and whatever is asked of the // result afterwards fails with "No output generated". Keeping the @@ -350,8 +360,17 @@ export function make({ modelProviders, httpClients, requests, registry }: Option }, }, } - : { system: input.system }), - messages: [...input.messages, ...(input.continuationMessages ?? [])], + : { + system: breakpoint + ? withBreakpoint({ role: "system", content: input.system }, breakpoint) + : input.system, + }), + messages: [ + ...input.messages.map(({ cacheBreakpoint, ...message }) => + cacheBreakpoint && breakpoint ? withBreakpoint(message, breakpoint) : message, + ), + ...(input.continuationMessages ?? []), + ], tools: input.tools, toolApproval: input.toolApproval, abortSignal: signal, @@ -379,9 +398,13 @@ export function make({ modelProviders, httpClients, requests, registry }: Option // with nothing said, so the last round is told to answer. The note // rides on the newest tool result, which no provider has cached yet: // changing the tools or system text instead would cost the cache - // every earlier call built. - prepareStep: ({ stepNumber, messages }) => - stepNumber === input.maxSteps - 1 ? { messages: withLastStepNote(messages) } : undefined, + // every earlier call built. Every other step marks its newest message + // for the next step to read up to from the cache; the last has no + // next step to read it. + prepareStep: ({ stepNumber, messages }) => { + if (stepNumber === input.maxSteps - 1) return { messages: withLastStepNote(messages) }; + return markNewest ? { messages: markNewest(messages) } : undefined; + }, stopWhen: stepCountIs(input.maxSteps), maxRetries: 0, }); @@ -489,6 +512,36 @@ function noted(output: ToolResultPart["output"]): ToolResultPart["output"] { } } +/** + * Marks each step's newest message with `breakpoint`, taking it off the + * message the step before marked. The SDK sends each step the messages the + * step before was sent, so a mark left behind would count against the four + * Anthropic allows a request. + */ +function newestMessageBreakpoint(breakpoint: ProviderOptions) { + let marked: { index: number; unmarked: ModelMessage } | undefined; + return (messages: ModelMessage[]): ModelMessage[] => { + const restored = marked ? messages.with(marked.index, marked.unmarked) : messages; + const index = restored.length - 1; + const newest = restored[index]; + if (!newest) return restored; + marked = { index, unmarked: newest }; + return restored.with(index, withBreakpoint(newest, breakpoint)); + }; +} + +/** `message` with `breakpoint`'s options added to whatever options it has for the same provider. */ +function withBreakpoint( + message: Message, + breakpoint: ProviderOptions, +): Message { + const providerOptions = { ...message.providerOptions }; + for (const [provider, options] of Object.entries(breakpoint)) { + providerOptions[provider] = { ...providerOptions[provider], ...options }; + } + return { ...message, providerOptions }; +} + function languageModel( connection: ModelProviderRepository.ProviderEndpoint, modelId: string,