illuma-agents 1.0.2 → 1.0.3

package/src/messages/cache.test.ts

@@ -0,0 +1,461 @@
+import type Anthropic from '@anthropic-ai/sdk';
+import type { AnthropicMessages } from '@/types/messages';
+import { addCacheControl, addBedrockCacheControl } from './cache';
+import { MessageContentComplex } from '@langchain/core/messages';
+import { ContentTypes } from '@/common/enum';
+
+describe('addCacheControl', () => {
+  test('should add cache control to the last two user messages with array content', () => {
+    const messages: AnthropicMessages = [
+      { role: 'user', content: [{ type: 'text', text: 'Hello' }] },
+      { role: 'assistant', content: [{ type: 'text', text: 'Hi there' }] },
+      { role: 'user', content: [{ type: 'text', text: 'How are you?' }] },
+      {
+        role: 'assistant',
+        content: [{ type: 'text', text: 'I\'m doing well, thanks!' }],
+      },
+      { role: 'user', content: [{ type: 'text', text: 'Great!' }] },
+    ];
+
+    const result = addCacheControl(messages);
+
+    expect(result[0].content[0]).not.toHaveProperty('cache_control');
+    expect(
+      (result[2].content[0] as Anthropic.TextBlockParam).cache_control
+    ).toEqual({ type: 'ephemeral' });
+    expect(
+      (result[4].content[0] as Anthropic.TextBlockParam).cache_control
+    ).toEqual({ type: 'ephemeral' });
+  });
+
+  test('should add cache control to the last two user messages with string content', () => {
+    const messages: AnthropicMessages = [
+      { role: 'user', content: 'Hello' },
+      { role: 'assistant', content: 'Hi there' },
+      { role: 'user', content: 'How are you?' },
+      { role: 'assistant', content: 'I\'m doing well, thanks!' },
+      { role: 'user', content: 'Great!' },
+    ];
+
+    const result = addCacheControl(messages);
+
+    expect(result[0].content).toBe('Hello');
+    expect(result[2].content[0]).toEqual({
+      type: 'text',
+      text: 'How are you?',
+      cache_control: { type: 'ephemeral' },
+    });
+    expect(result[4].content[0]).toEqual({
+      type: 'text',
+      text: 'Great!',
+      cache_control: { type: 'ephemeral' },
+    });
+  });
+
+  test('should handle mixed string and array content', () => {
+    const messages: AnthropicMessages = [
+      { role: 'user', content: 'Hello' },
+      { role: 'assistant', content: 'Hi there' },
+      { role: 'user', content: [{ type: 'text', text: 'How are you?' }] },
+    ];
+
+    const result = addCacheControl(messages);
+
+    expect(result[0].content[0]).toEqual({
+      type: 'text',
+      text: 'Hello',
+      cache_control: { type: 'ephemeral' },
+    });
+    expect(
+      (result[2].content[0] as Anthropic.TextBlockParam).cache_control
+    ).toEqual({ type: 'ephemeral' });
+  });
+
+  test('should handle less than two user messages', () => {
+    const messages: AnthropicMessages = [
+      { role: 'user', content: 'Hello' },
+      { role: 'assistant', content: 'Hi there' },
+    ];
+
+    const result = addCacheControl(messages);
+
+    expect(result[0].content[0]).toEqual({
+      type: 'text',
+      text: 'Hello',
+      cache_control: { type: 'ephemeral' },
+    });
+    expect(result[1].content).toBe('Hi there');
+  });
+
+  test('should return original array if no user messages', () => {
+    const messages: AnthropicMessages = [
+      { role: 'assistant', content: 'Hi there' },
+      { role: 'assistant', content: 'How can I help?' },
+    ];
+
+    const result = addCacheControl(messages);
+
+    expect(result).toEqual(messages);
+  });
+
+  test('should handle empty array', () => {
+    const messages: AnthropicMessages = [];
+    const result = addCacheControl(messages);
+    expect(result).toEqual([]);
+  });
+
+  test('should handle non-array input', () => {
+    const messages = 'not an array';
+    /** @ts-expect-error - This is a test */
+    const result = addCacheControl(messages);
+    expect(result).toBe('not an array');
+  });
+
+  test('should not modify assistant messages', () => {
+    const messages: AnthropicMessages = [
+      { role: 'user', content: 'Hello' },
+      { role: 'assistant', content: 'Hi there' },
+      { role: 'user', content: 'How are you?' },
+    ];
+
+    const result = addCacheControl(messages);
+
+    expect(result[1].content).toBe('Hi there');
+  });
+
+  test('should handle multiple content items in user messages', () => {
+    const messages: AnthropicMessages = [
+      {
+        role: 'user',
+        content: [
+          { type: 'text', text: 'Hello' },
+          {
+            type: 'image',
+            source: { type: 'url', url: 'http://example.com/image.jpg' },
+          },
+          { type: 'text', text: 'This is an image' },
+        ],
+      },
+      { role: 'assistant', content: 'Hi there' },
+      { role: 'user', content: 'How are you?' },
+    ];
+
+    const result = addCacheControl(messages);
+
+    expect(result[0].content[0]).not.toHaveProperty('cache_control');
+    expect(result[0].content[1]).not.toHaveProperty('cache_control');
+    expect(
+      (result[0].content[2] as Anthropic.TextBlockParam).cache_control
+    ).toEqual({ type: 'ephemeral' });
+    expect(result[2].content[0]).toEqual({
+      type: 'text',
+      text: 'How are you?',
+      cache_control: { type: 'ephemeral' },
+    });
+  });
+
+  test('should handle an array with mixed content types', () => {
+    const messages: AnthropicMessages = [
+      { role: 'user', content: 'Hello' },
+      { role: 'assistant', content: 'Hi there' },
+      { role: 'user', content: [{ type: 'text', text: 'How are you?' }] },
+      { role: 'assistant', content: 'I\'m doing well, thanks!' },
+      { role: 'user', content: 'Great!' },
+    ];
+
+    const result = addCacheControl(messages);
+
+    expect(result[0].content).toEqual('Hello');
+    expect(result[2].content[0]).toEqual({
+      type: 'text',
+      text: 'How are you?',
+      cache_control: { type: 'ephemeral' },
+    });
+    expect(result[4].content).toEqual([
+      {
+        type: 'text',
+        text: 'Great!',
+        cache_control: { type: 'ephemeral' },
+      },
+    ]);
+    expect(result[1].content).toBe('Hi there');
+    expect(result[3].content).toBe('I\'m doing well, thanks!');
+  });
+
+  test('should handle edge case with multiple content types', () => {
+    const messages: AnthropicMessages = [
+      {
+        role: 'user',
+        content: [
+          {
+            type: 'image',
+            source: {
+              type: 'base64',
+              media_type: 'image/png',
+              data: 'some_base64_string',
+            },
+          },
+          {
+            type: 'image',
+            source: {
+              type: 'base64',
+              media_type: 'image/png',
+              data: 'another_base64_string',
+            },
+          },
+          { type: 'text', text: 'what do all these images have in common' },
+        ],
+      },
+      { role: 'assistant', content: 'I see multiple images.' },
+      { role: 'user', content: 'Correct!' },
+    ];
+
+    const result = addCacheControl(messages);
+
+    expect(result[0].content[0]).not.toHaveProperty('cache_control');
+    expect(result[0].content[1]).not.toHaveProperty('cache_control');
+    expect(
+      (result[0].content[2] as Anthropic.ImageBlockParam).cache_control
+    ).toEqual({ type: 'ephemeral' });
+    expect(result[2].content[0]).toEqual({
+      type: 'text',
+      text: 'Correct!',
+      cache_control: { type: 'ephemeral' },
+    });
+  });
+
+  test('should handle user message with no text block', () => {
+    const messages: AnthropicMessages = [
+      {
+        role: 'user',
+        content: [
+          {
+            type: 'image',
+            source: {
+              type: 'base64',
+              media_type: 'image/png',
+              data: 'some_base64_string',
+            },
+          },
+          {
+            type: 'image',
+            source: {
+              type: 'base64',
+              media_type: 'image/png',
+              data: 'another_base64_string',
+            },
+          },
+        ],
+      },
+      { role: 'assistant', content: 'I see two images.' },
+      { role: 'user', content: 'Correct!' },
+    ];
+
+    const result = addCacheControl(messages);
+
+    expect(result[0].content[0]).not.toHaveProperty('cache_control');
+    expect(result[0].content[1]).not.toHaveProperty('cache_control');
+    expect(result[2].content[0]).toEqual({
+      type: 'text',
+      text: 'Correct!',
+      cache_control: { type: 'ephemeral' },
+    });
+  });
+});
+
+type TestMsg = {
+  role?: 'user' | 'assistant' | 'system';
+  content?: string | MessageContentComplex[];
+};
+
+describe('addBedrockCacheControl (Bedrock cache checkpoints)', () => {
+  it('returns input when not enough messages', () => {
+    const empty: TestMsg[] = [];
+    expect(addBedrockCacheControl(empty)).toEqual(empty);
+    const single: TestMsg[] = [{ role: 'user', content: 'only' }];
+    expect(addBedrockCacheControl(single)).toEqual(single);
+  });
+
+  it('wraps string content and appends separate cachePoint block', () => {
+    const messages: TestMsg[] = [
+      { role: 'user', content: 'Hello' },
+      { role: 'assistant', content: [{ type: ContentTypes.TEXT, text: 'Hi' }] },
+    ];
+    const result = addBedrockCacheControl(messages);
+    const last = result[1].content as MessageContentComplex[];
+    expect(Array.isArray(last)).toBe(true);
+    expect(last[0]).toEqual({ type: ContentTypes.TEXT, text: 'Hi' });
+    expect(last[1]).toEqual({ cachePoint: { type: 'default' } });
+  });
+
+  it('inserts cachePoint after the last text when multiple text blocks exist', () => {
+    const messages: TestMsg[] = [
+      {
+        role: 'user',
+        content: [
+          { type: ContentTypes.TEXT, text: 'Intro' },
+          { type: ContentTypes.TEXT, text: 'Details' },
+          {
+            type: ContentTypes.IMAGE_FILE,
+            image_file: { file_id: 'file_123' },
+          },
+        ],
+      },
+      {
+        role: 'assistant',
+        content: [
+          { type: ContentTypes.TEXT, text: 'Reply A' },
+          { type: ContentTypes.TEXT, text: 'Reply B' },
+        ],
+      },
+    ];
+
+    const result = addBedrockCacheControl(messages);
+
+    const first = result[0].content as MessageContentComplex[];
+    const second = result[1].content as MessageContentComplex[];
+
+    expect(first[0]).toEqual({ type: ContentTypes.TEXT, text: 'Intro' });
+    expect(first[1]).toEqual({ type: ContentTypes.TEXT, text: 'Details' });
+    expect(first[2]).toEqual({ cachePoint: { type: 'default' } });
+
+    const img = first[3] as MessageContentComplex;
+    expect(img.type).toBe(ContentTypes.IMAGE_FILE);
+    if (img.type === ContentTypes.IMAGE_FILE) {
+      expect('image_file' in img).toBe(true);
+    }
+
+    expect(second[0]).toEqual({ type: ContentTypes.TEXT, text: 'Reply A' });
+    expect(second[1]).toEqual({ type: ContentTypes.TEXT, text: 'Reply B' });
+    expect(second[2]).toEqual({ cachePoint: { type: 'default' } });
+  });
+
+  it('skips adding cachePoint when content is an empty array', () => {
+    const messages: TestMsg[] = [
+      { role: 'user', content: [] },
+      { role: 'assistant', content: [] },
+      { role: 'user', content: 'ignored because only last two are modified' },
+    ];
+
+    const result = addBedrockCacheControl(messages);
+
+    const first = result[0].content as MessageContentComplex[];
+    const second = result[1].content as MessageContentComplex[];
+
+    expect(Array.isArray(first)).toBe(true);
+    expect(first.length).toBe(0);
+
+    expect(Array.isArray(second)).toBe(true);
+    expect(second.length).toBe(0);
+    expect(second[0]).not.toEqual({ cachePoint: { type: 'default' } });
+  });
+
+  it('skips adding cachePoint when content is an empty string', () => {
+    const messages: TestMsg[] = [
+      { role: 'user', content: '' },
+      { role: 'assistant', content: '' },
+      { role: 'user', content: 'ignored because only last two are modified' },
+    ];
+
+    const result = addBedrockCacheControl(messages);
+
+    expect(result[0].content).toBe('');
+    expect(result[1].content).toBe('');
+  });
+
+  /** (I don't think this will ever occur in actual use, but its the only branch left uncovered so I'm covering it */
+  it('skips messages with non-string, non-array content and still modifies the previous to reach two edits', () => {
+    const messages: TestMsg[] = [
+      {
+        role: 'user',
+        content: [{ type: ContentTypes.TEXT, text: 'Will be modified' }],
+      },
+      { role: 'assistant', content: undefined },
+      {
+        role: 'user',
+        content: [{ type: ContentTypes.TEXT, text: 'Also modified' }],
+      },
+    ];
+
+    const result = addBedrockCacheControl(messages);
+
+    const last = result[2].content as MessageContentComplex[];
+    expect(last[0]).toEqual({ type: ContentTypes.TEXT, text: 'Also modified' });
+    expect(last[1]).toEqual({ cachePoint: { type: 'default' } });
+
+    expect(result[1].content).toBeUndefined();
+
+    const first = result[0].content as MessageContentComplex[];
+    expect(first[0]).toEqual({
+      type: ContentTypes.TEXT,
+      text: 'Will be modified',
+    });
+    expect(first[1]).toEqual({ cachePoint: { type: 'default' } });
+  });
+
+  it('works with the example from the langchain pr', () => {
+    const messages: TestMsg[] = [
+      {
+        role: 'system',
+        content: [
+          { type: ContentTypes.TEXT, text: 'You\'re an advanced AI assistant.' },
+        ],
+      },
+      {
+        role: 'user',
+        content: [
+          { type: ContentTypes.TEXT, text: 'What is the capital of France?' },
+        ],
+      },
+    ];
+
+    const result = addBedrockCacheControl(messages);
+
+    let system = result[0].content as MessageContentComplex[];
+    let user = result[1].content as MessageContentComplex[];
+
+    expect(system[0]).toEqual({
+      type: ContentTypes.TEXT,
+      text: 'You\'re an advanced AI assistant.',
+    });
+    expect(system[1]).toEqual({ cachePoint: { type: 'default' } });
+    expect(user[0]).toEqual({
+      type: ContentTypes.TEXT,
+      text: 'What is the capital of France?',
+    });
+    expect(user[1]).toEqual({ cachePoint: { type: 'default' } });
+
+    result.push({
+      role: 'assistant',
+      content: [
+        {
+          type: ContentTypes.TEXT,
+          text: 'Sure! The capital of France is Paris.',
+        },
+      ],
+    });
+
+    const result2 = addBedrockCacheControl(result);
+
+    system = result2[0].content as MessageContentComplex[];
+    user = result2[1].content as MessageContentComplex[];
+    const assistant = result2[2].content as MessageContentComplex[];
+
+    expect(system[0]).toEqual({
+      type: ContentTypes.TEXT,
+      text: 'You\'re an advanced AI assistant.',
+    });
+    expect(system[1]).toEqual({ cachePoint: { type: 'default' } });
+    expect(user[0]).toEqual({
+      type: ContentTypes.TEXT,
+      text: 'What is the capital of France?',
+    });
+    expect(user[1]).toEqual({ cachePoint: { type: 'default' } });
+
+    expect(assistant[0]).toEqual({
+      type: ContentTypes.TEXT,
+      text: 'Sure! The capital of France is Paris.',
+    });
+    expect(assistant[1]).toEqual({ cachePoint: { type: 'default' } });
+  });
+});

package/src/messages/cache.ts

@@ -0,0 +1,151 @@
+import { BaseMessage, MessageContentComplex } from '@langchain/core/messages';
+import type { AnthropicMessage } from '@/types/messages';
+import type Anthropic from '@anthropic-ai/sdk';
+import { ContentTypes } from '@/common/enum';
+
+type MessageWithContent = {
+  content?: string | MessageContentComplex[];
+};
+
+/**
+ * Anthropic API: Adds cache control to the appropriate user messages in the payload.
+ * @param messages - The array of message objects.
+ * @returns - The updated array of message objects with cache control added.
+ */
+export function addCacheControl<T extends AnthropicMessage | BaseMessage>(
+  messages: T[]
+): T[] {
+  if (!Array.isArray(messages) || messages.length < 2) {
+    return messages;
+  }
+
+  const updatedMessages = [...messages];
+  let userMessagesModified = 0;
+
+  for (
+    let i = updatedMessages.length - 1;
+    i >= 0 && userMessagesModified < 2;
+    i--
+  ) {
+    const message = updatedMessages[i];
+    if ('getType' in message && message.getType() !== 'human') {
+      continue;
+    } else if ('role' in message && message.role !== 'user') {
+      continue;
+    }
+
+    if (typeof message.content === 'string') {
+      message.content = [
+        {
+          type: 'text',
+          text: message.content,
+          cache_control: { type: 'ephemeral' },
+        },
+      ];
+      userMessagesModified++;
+    } else if (Array.isArray(message.content)) {
+      for (let j = message.content.length - 1; j >= 0; j--) {
+        const contentPart = message.content[j];
+        if ('type' in contentPart && contentPart.type === 'text') {
+          (contentPart as Anthropic.TextBlockParam).cache_control = {
+            type: 'ephemeral',
+          };
+          userMessagesModified++;
+          break;
+        }
+      }
+    }
+  }
+
+  return updatedMessages;
+}
+
+/**
+ * Adds Bedrock Converse API cache points to the last two messages.
+ * Inserts `{ cachePoint: { type: 'default' } }` as a separate content block
+ * immediately after the last text block in each targeted message.
+ * @param messages - The array of message objects.
+ * @returns - The updated array of message objects with cache points added.
+ */
+export function addBedrockCacheControl<
+  T extends Partial<BaseMessage> & MessageWithContent,
+>(messages: T[]): T[] {
+  if (!Array.isArray(messages) || messages.length < 2) {
+    return messages;
+  }
+
+  const updatedMessages: T[] = messages.slice();
+  let messagesModified = 0;
+
+  for (
+    let i = updatedMessages.length - 1;
+    i >= 0 && messagesModified < 2;
+    i--
+  ) {
+    const message = updatedMessages[i];
+
+    if (
+      'getType' in message &&
+      typeof message.getType === 'function' &&
+      message.getType() === 'tool'
+    ) {
+      continue;
+    }
+
+    const content = message.content;
+
+    if (typeof content === 'string' && content === '') {
+      continue;
+    }
+
+    if (typeof content === 'string') {
+      message.content = [
+        { type: ContentTypes.TEXT, text: content },
+        { cachePoint: { type: 'default' } },
+      ] as MessageContentComplex[];
+      messagesModified++;
+      continue;
+    }
+
+    if (Array.isArray(content)) {
+      let hasCacheableContent = false;
+      for (const block of content) {
+        if (block.type === ContentTypes.TEXT) {
+          if (typeof block.text === 'string' && block.text !== '') {
+            hasCacheableContent = true;
+            break;
+          }
+        }
+      }
+
+      if (!hasCacheableContent) {
+        continue;
+      }
+
+      let inserted = false;
+      for (let j = content.length - 1; j >= 0; j--) {
+        const block = content[j] as MessageContentComplex;
+        const type = (block as { type?: string }).type;
+        if (type === ContentTypes.TEXT || type === 'text') {
+          const text = (block as { text?: string }).text;
+          if (text === '' || text === undefined) {
+            continue;
+          }
+          content.splice(j + 1, 0, {
+            cachePoint: { type: 'default' },
+          } as MessageContentComplex);
+          inserted = true;
+          break;
+        }
+      }
+      if (!inserted) {
+        content.push({
+          cachePoint: { type: 'default' },
+        } as MessageContentComplex);
+      }
+      messagesModified++;
+    }
+  }
+
+  return updatedMessages;
+}