🐛 fix: Anthropic Prompt Caching Edge Case (#3690)

This commit is contained in:
Danny Avila 2024-08-18 19:03:03 -04:00 committed by GitHub
parent 87d95a9d82
commit 8ca1e4f94f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 89 additions and 26 deletions

View file

@ -9,31 +9,31 @@ function addCacheControl(messages) {
}
const updatedMessages = [...messages];
let userMessagesFound = 0;
let userMessagesModified = 0;
for (let i = updatedMessages.length - 1; i >= 0 && userMessagesFound < 2; i--) {
if (updatedMessages[i].role === 'user') {
if (typeof updatedMessages[i].content === 'string') {
updatedMessages[i] = {
...updatedMessages[i],
content: [
{
type: 'text',
text: updatedMessages[i].content,
cache_control: { type: 'ephemeral' },
},
],
};
} else if (Array.isArray(updatedMessages[i].content)) {
updatedMessages[i] = {
...updatedMessages[i],
content: updatedMessages[i].content.map((item) => ({
...item,
cache_control: { type: 'ephemeral' },
})),
};
for (let i = updatedMessages.length - 1; i >= 0 && userMessagesModified < 2; i--) {
const message = updatedMessages[i];
if (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--) {
if (message.content[j].type === 'text') {
message.content[j].cache_control = { type: 'ephemeral' };
userMessagesModified++;
break;
}
}
userMessagesFound++;
}
}

View file

@ -116,6 +116,7 @@ describe('addCacheControl', () => {
content: [
{ type: 'text', text: 'Hello' },
{ type: 'image', url: 'http://example.com/image.jpg' },
{ type: 'text', text: 'This is an image' },
],
},
{ role: 'assistant', content: 'Hi there' },
@ -124,8 +125,9 @@ describe('addCacheControl', () => {
const result = addCacheControl(messages);
expect(result[0].content[0].cache_control).toEqual({ type: 'ephemeral' });
expect(result[0].content[1].cache_control).toEqual({ type: 'ephemeral' });
expect(result[0].content[0]).not.toHaveProperty('cache_control');
expect(result[0].content[1]).not.toHaveProperty('cache_control');
expect(result[0].content[2].cache_control).toEqual({ type: 'ephemeral' });
expect(result[2].content[0]).toEqual({
type: 'text',
text: 'How are you?',
@ -143,7 +145,6 @@ describe('addCacheControl', () => {
];
const result = addCacheControl(messages);
console.dir(result, { depth: null });
expect(result[0].content).toEqual('Hello');
expect(result[2].content[0]).toEqual({
@ -161,4 +162,66 @@ describe('addCacheControl', () => {
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 = [
{
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].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 = [
{
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' },
});
});
});