Anthropic prompt caching

Ironclad · Sep 3, 2024 · b331b4c · b331b4c
1 parent bf543a5
commit b331b4c
Show file tree

Hide file tree

Showing 5 changed files with 182 additions and 71 deletions.
diff --git a/packages/app/src/components/editors/DefaultNodeEditor.tsx b/packages/app/src/components/editors/DefaultNodeEditor.tsx
@@ -72,6 +72,10 @@ export const defaultEditorContainerStyles = css`
       min-width: 75px;
     }
 
+    label:nth-child(2) {
+      min-width: 32px;
+    }
+
     &.use-input-toggle label:first-child {
       min-width: unset;
     }

diff --git a/packages/core/src/model/DataValue.ts b/packages/core/src/model/DataValue.ts
@@ -14,11 +14,17 @@ export type BoolDataValue = DataValueDef<'boolean', boolean>;
 export type SystemChatMessage = {
   type: 'system';
   message: ChatMessageMessagePart | ChatMessageMessagePart[];
+
+  /** If true, this message marks a breakpoint when used with prompt caching (as of right now, Anthropic-only). */
+  isCacheBreakpoint?: boolean;
 };
 
 export type UserChatMessage = {
   type: 'user';
   message: ChatMessageMessagePart | ChatMessageMessagePart[];
+
+  /** If true, this message marks a breakpoint when used with prompt caching (as of right now, Anthropic-only). */
+  isCacheBreakpoint?: boolean;
 };
 
 export type AssistantChatMessageFunctionCall = {
@@ -41,12 +47,18 @@ export type AssistantChatMessage = {
   function_call: AssistantChatMessageFunctionCall | undefined;
 
   function_calls: AssistantChatMessageFunctionCall[] | undefined;
+
+  /** If true, this message marks a breakpoint when used with prompt caching (as of right now, Anthropic-only). */
+  isCacheBreakpoint?: boolean;
 };
 
 export type FunctionResponseChatMessage = {
   type: 'function';
   message: ChatMessageMessagePart | ChatMessageMessagePart[];
   name: string;
+
+  /** If true, this message marks a breakpoint when used with prompt caching (as of right now, Anthropic-only). */
+  isCacheBreakpoint?: boolean;
 };
 
 export type ChatMessage = SystemChatMessage | UserChatMessage | AssistantChatMessage | FunctionResponseChatMessage;
@@ -389,6 +401,7 @@ export const scalarDefaults: { [P in ScalarDataType]: Extract<ScalarDataValue, {
   'chat-message': {
     type: 'user',
     message: '',
+    isCacheBreakpoint: undefined,
   },
   'control-flow-excluded': undefined,
   date: new Date().toISOString(),

diff --git a/packages/core/src/model/nodes/PromptNode.ts b/packages/core/src/model/nodes/PromptNode.ts
@@ -37,6 +37,9 @@ export type PromptNodeData = {
   useNameInput?: boolean;
   enableFunctionCall?: boolean;
   computeTokenCount?: boolean;
+
+  isCacheBreakpoint?: boolean;
+  useIsCacheBreakpointInput?: boolean;
 };
 
 export class PromptNodeImpl extends NodeImpl<PromptNode> {
@@ -160,6 +163,14 @@ export class PromptNodeImpl extends NodeImpl<PromptNode> {
         label: 'Compute Token Count',
         dataKey: 'computeTokenCount',
       },
+      {
+        type: 'toggle',
+        label: 'Is Cache Breakpoint',
+        dataKey: 'isCacheBreakpoint',
+        helperMessage:
+          'For Anthropic, marks this message as a cache breakpoint - this message and every message before it will be cached using Prompt Caching.',
+        useInputToggleDataKey: 'useIsCacheBreakpointInput',
+      },
       {
         type: 'code',
         label: 'Prompt Text',
@@ -208,6 +219,7 @@ export class PromptNodeImpl extends NodeImpl<PromptNode> {
     const outputValue = interpolate(this.chartNode.data.promptText, inputMap);
 
     const type = getInputOrData(this.data, inputs, 'type', 'string');
+    const isCacheBreakpoint = getInputOrData(this.data, inputs, 'isCacheBreakpoint', 'boolean');
 
     if (['assistant', 'system', 'user', 'function'].includes(type) === false) {
       throw new Error(`Invalid type: ${type}`);
@@ -219,13 +231,15 @@ export class PromptNodeImpl extends NodeImpl<PromptNode> {
         (type): ChatMessage => ({
           type,
           message: outputValue,
+          isCacheBreakpoint,
         }),
       )
       .with(
         'user',
         (type): ChatMessage => ({
           type,
           message: outputValue,
+          isCacheBreakpoint,
         }),
       )
       .with('assistant', (type): ChatMessage => {
@@ -248,6 +262,7 @@ export class PromptNodeImpl extends NodeImpl<PromptNode> {
           message: outputValue,
           function_call: functionCall as AssistantChatMessageFunctionCall,
           function_calls: functionCall ? [functionCall as AssistantChatMessageFunctionCall] : undefined,
+          isCacheBreakpoint,
         };
       })
       .with(
@@ -256,6 +271,7 @@ export class PromptNodeImpl extends NodeImpl<PromptNode> {
           type,
           message: outputValue,
           name: getInputOrData(this.data, inputs, 'name', 'string'),
+          isCacheBreakpoint,
         }),
       )
       .otherwise(() => {

diff --git a/packages/core/src/plugins/anthropic/anthropic.ts b/packages/core/src/plugins/anthropic/anthropic.ts
@@ -86,11 +86,12 @@ export const anthropicModelOptions = Object.entries(anthropicModels).map(([id, {
 export type Claude3ChatMessage = {
   role: 'user' | 'assistant';
   content: string | Claude3ChatMessageContentPart[];
-}
+};
 
 export type Claude3ChatMessageTextContentPart = {
   type: 'text';
   text: string;
+  cache_control: CacheControl;
 };
 
 export type Claude3ChatMessageImageContentPart = {
@@ -100,22 +101,25 @@ export type Claude3ChatMessageImageContentPart = {
     media_type: string;
     data: string;
   };
+  cache_control: CacheControl;
 };
 
 export type Claude3ChatMessageToolResultContentPart = {
   type: 'tool_result';
   tool_use_id: string;
-  content: string | { type: 'text'; text: string; }[];
+  content: string | { type: 'text'; text: string }[];
+  cache_control: CacheControl;
 };
 
 export type Claude3ChatMessageToolUseContentPart = {
   type: 'tool_use';
   id: string;
   name: string;
   input: object;
-}
+  cache_control: CacheControl;
+};
 
-export type Claude3ChatMessageContentPart = 
+export type Claude3ChatMessageContentPart =
   | Claude3ChatMessageTextContentPart
   | Claude3ChatMessageImageContentPart
   | Claude3ChatMessageToolResultContentPart
@@ -125,7 +129,7 @@ export type ChatMessageOptions = {
   apiKey: string;
   model: AnthropicModels;
   messages: Claude3ChatMessage[];
-  system?: string;
+  system?: SystemPrompt;
   max_tokens: number;
   stop_sequences?: string[];
   temperature?: number;
@@ -138,6 +142,7 @@ export type ChatMessageOptions = {
     description: string;
     input_schema: object;
   }[];
+  beta?: string;
 };
 
 export type ChatCompletionOptions = {
@@ -159,62 +164,83 @@ export type ChatCompletionChunk = {
   model: string;
 };
 
-export type ChatMessageChunk = {
-  type: 'message_start';
-  message: {
-    id: string;
-    type: string;
-    role: string;
-    content: {
-      type: 'text';
-      text: string;
-    }[];
-    model: AnthropicModels;
-    stop_reason: string | null;
-    stop_sequence: string | null;
-    usage: {
-      input_tokens: number;
-      output_tokens: number;
-    };
-  };
-} | {
-  type: 'content_block_start';
-  index: number;
-  content_block: {
-    type: 'text';
-    text: string;
-  };
-} | {
-  type: 'ping';
-} | {
-  type: 'content_block_delta';
-  index: number;
-  delta: {
-    type: 'text_delta';
-    text: string;
-  }
-} | {
-  type: 'message_delta';
-  delta: {
-    stop_reason: string | null;
-    stop_sequence: string | null;
-    usage: {
-      output_tokens: number;
-    }
-  }
-} | {
-  type: 'message_stop';
+export type CacheControl = null | {
+  type: 'ephemeral';
 };
 
+export type SystemPrompt = string | SystemPromptMessage[];
+
+export type SystemPromptMessage = {
+  cache_control: CacheControl;
+  type: 'text';
+  text: string;
+};
+
+export type ChatMessageChunk =
+  | {
+      type: 'message_start';
+      message: {
+        id: string;
+        type: string;
+        role: string;
+        content: {
+          type: 'text';
+          text: string;
+        }[];
+        model: AnthropicModels;
+        stop_reason: string | null;
+        stop_sequence: string | null;
+        usage: {
+          input_tokens: number;
+          output_tokens: number;
+        };
+      };
+    }
+  | {
+      type: 'content_block_start';
+      index: number;
+      content_block: {
+        type: 'text';
+        text: string;
+      };
+    }
+  | {
+      type: 'ping';
+    }
+  | {
+      type: 'content_block_delta';
+      index: number;
+      delta: {
+        type: 'text_delta';
+        text: string;
+      };
+    }
+  | {
+      type: 'message_delta';
+      delta: {
+        stop_reason: string | null;
+        stop_sequence: string | null;
+        usage: {
+          output_tokens: number;
+        };
+      };
+    }
+  | {
+      type: 'message_stop';
+    };
+
 export type ChatMessageResponse = {
   id: string;
-  content: ({
-    text: string;
-  } | {
-    id: string;
-    name: string;
-    input: object;
-  })[];
+  content: (
+    | {
+        text: string;
+      }
+    | {
+        id: string;
+        name: string;
+        input: object;
+      }
+  )[];
   model: string;
   stop_reason: 'end_turn';
   stop_sequence: string;
@@ -230,7 +256,7 @@ export async function* streamChatCompletions({
   ...rest
 }: ChatCompletionOptions): AsyncGenerator<ChatCompletionChunk> {
   const defaultSignal = new AbortController().signal;
-  const response = await fetchEventSource('https://api.anthropic.com/v1/complete', {
+  const response = await fetchEventSource('https://api.anthropic.com/v1/completions', {
     method: 'POST',
     headers: {
       'Content-Type': 'application/json',
@@ -278,6 +304,7 @@ export async function callMessageApi({
   apiKey,
   signal,
   tools,
+  beta,
   ...rest
 }: ChatMessageOptions): Promise<ChatMessageResponse> {
   const defaultSignal = new AbortController().signal;
@@ -287,7 +314,7 @@ export async function callMessageApi({
       'Content-Type': 'application/json',
       'x-api-key': apiKey,
       'anthropic-version': '2023-06-01',
-      'anthropic-beta': tools ? 'tools-2024-04-04' : 'messages-2023-12-15',
+      ...(beta ? { 'anthropic-beta': beta } : {}),
     },
     body: JSON.stringify({
       ...rest,
@@ -306,17 +333,19 @@ export async function callMessageApi({
 export async function* streamMessageApi({
   apiKey,
   signal,
+  beta,
   ...rest
 }: ChatMessageOptions): AsyncGenerator<ChatMessageChunk> {
   // Use the Messages API for Claude 3 models
   const defaultSignal = new AbortController().signal;
+  console.dir({ rest }, { depth: null });
   const response = await fetchEventSource('https://api.anthropic.com/v1/messages', {
     method: 'POST',
     headers: {
       'Content-Type': 'application/json',
       'x-api-key': apiKey,
       'anthropic-version': '2023-06-01',
-      'anthropic-beta': 'messages-2023-12-15',
+      ...(beta ? { 'anthropic-beta': beta } : {}),
     },
     body: JSON.stringify({
       ...rest,