Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Inference] Image content #205371

Merged
merged 4 commits into from
Jan 7, 2025
Merged

[Inference] Image content #205371

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
f8feecf
[Inference] Image content
dgieselaar Jan 2, 2025
ed4394a
Add tests
dgieselaar Jan 2, 2025
87f3f6d
TSDoc
dgieselaar Jan 2, 2025
b5552c8
Merge branch 'main' into inference-image-content
elasticmachine Jan 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions x-pack/platform/packages/shared/ai-infra/inference-common/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,9 @@ export {
ChatCompletionEventType,
ToolChoiceType,
type Message,
type MessageContentImage,
type MessageContentText,
type MessageContent,
type AssistantMessage,
type ToolMessage,
type UserMessage,
Expand Down
3 changes: 3 additions & 0 deletions x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,9 @@ export {
} from './events';
export {
MessageRole,
type MessageContent,
type MessageContentImage,
type MessageContentText,
type Message,
type AssistantMessage,
type UserMessage,
Expand Down
16 changes: 14 additions & 2 deletions x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/messages.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,26 @@ interface MessageBase<TRole extends MessageRole> {
role: TRole;
}

export interface MessageContentText {
type: 'text';
text: string;
}

export interface MessageContentImage {
type: 'image';
source: { data: string; mimeType: string };
}

export type MessageContent = string | Array<MessageContentText | MessageContentImage>;

/**
* Represents a message from the user.
*/
export type UserMessage = MessageBase<MessageRole.User> & {
/**
* The text content of the user message
* The text or image content of the user message
*/
content: string;
content: MessageContent;
};

/**
Expand Down
86 changes: 86 additions & 0 deletions ...ins/shared/inference/server/chat_complete/adapters/bedrock/bedrock_claude_adapter.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -256,6 +256,92 @@ describe('bedrockClaudeAdapter', () => {
expect(system).toEqual('Some system message');
});

it('correctly formats messages with content parts', () => {
bedrockClaudeAdapter.chatComplete({
executor: executorMock,
logger,
messages: [
{
role: MessageRole.User,
content: [
{
type: 'text',
text: 'question',
},
],
},
{
role: MessageRole.Assistant,
content: 'answer',
},
{
role: MessageRole.User,
content: [
{
type: 'image',
source: {
data: 'aaaaaa',
mimeType: 'image/png',
},
},
{
type: 'image',
source: {
data: 'bbbbbb',
mimeType: 'image/png',
},
},
],
},
],
});

expect(executorMock.invoke).toHaveBeenCalledTimes(1);

const { messages } = getCallParams();
expect(messages).toEqual([
{
rawContent: [
{
text: 'question',
type: 'text',
},
],
role: 'user',
},
{
rawContent: [
{
text: 'answer',
type: 'text',
},
],
role: 'assistant',
},
{
rawContent: [
{
type: 'image',
source: {
data: 'aaaaaa',
mediaType: 'image/png',
type: 'base64',
},
},
{
type: 'image',
source: {
data: 'bbbbbb',
mediaType: 'image/png',
type: 'base64',
},
},
],
role: 'user',
},
]);
});

it('correctly format tool choice', () => {
bedrockClaudeAdapter.chatComplete({
executor: executorMock,
Expand Down
21 changes: 19 additions & 2 deletions .../plugins/shared/inference/server/chat_complete/adapters/bedrock/bedrock_claude_adapter.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ import {
} from '@kbn/inference-common';
import { parseSerdeChunkMessage } from './serde_utils';
import { InferenceConnectorAdapter } from '../../types';
import type { BedRockMessage, BedrockToolChoice } from './types';
import type { BedRockImagePart, BedRockMessage, BedRockTextPart, BedrockToolChoice } from './types';
import {
BedrockChunkMember,
serdeEventstreamIntoObservable,
Expand Down Expand Up @@ -153,7 +153,24 @@ const messagesToBedrock = (messages: Message[]): BedRockMessage[] => {
case MessageRole.User:
return {
role: 'user' as const,
rawContent: [{ type: 'text' as const, text: message.content }],
rawContent: (typeof message.content === 'string'
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

NIT:

rawContent: typeof message.content === 'string' 
  ? [{ text: message.content, type: 'text' } satisfies BedRockTextPart]
  : message.content.map((contentPart) => {
    if (contentPart.type === 'text') {
      return { text: contentPart.text, type: 'text' } satisfies BedRockTextPart;
    } else {
      return {
        type: 'image',
        source: {
          data: contentPart.source.data,
          mediaType: contentPart.source.mimeType,
          type: 'base64',
        },
      } satisfies BedRockImagePart;
    }
}),

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

wdyt about castArray() instead?

? [message.content]
: message.content
).map((contentPart) => {
if (typeof contentPart === 'string') {
return { text: contentPart, type: 'text' } satisfies BedRockTextPart;
} else if (contentPart.type === 'text') {
return { text: contentPart.text, type: 'text' } satisfies BedRockTextPart;
}
return {
type: 'image',
source: {
data: contentPart.source.data,
mediaType: contentPart.source.mimeType,
type: 'base64',
},
} satisfies BedRockImagePart;
}),
};
case MessageRole.Assistant:
return {
Expand Down
39 changes: 31 additions & 8 deletions x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/bedrock/types.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,15 +17,38 @@ export interface BedRockMessage {
/**
* Bedrock message parts
*/
export interface BedRockTextPart {
type: 'text';
text: string;
}

export interface BedRockToolUsePart {
type: 'tool_use';
id: string;
name: string;
input: Record<string, unknown>;
}

export interface BedRockToolResultPart {
type: 'tool_result';
tool_use_id: string;
content: string;
}

export interface BedRockImagePart {
type: 'image';
source: {
type: 'base64';
mediaType: string;
data: string;
};
}

export type BedRockMessagePart =
| { type: 'text'; text: string }
| {
type: 'tool_use';
id: string;
name: string;
input: Record<string, unknown>;
}
| { type: 'tool_result'; tool_use_id: string; content: string };
| BedRockTextPart
| BedRockToolUsePart
| BedRockToolResultPart
| BedRockImagePart;

export type BedrockToolChoice = { type: 'auto' } | { type: 'any' } | { type: 'tool'; name: string };

Expand Down
80 changes: 80 additions & 0 deletions ...form/plugins/shared/inference/server/chat_complete/adapters/gemini/gemini_adapter.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -239,6 +239,86 @@ describe('geminiAdapter', () => {
]);
});

it('correctly formats content parts', () => {
geminiAdapter.chatComplete({
executor: executorMock,
logger,
messages: [
{
role: MessageRole.User,
content: [
{
type: 'text',
text: 'question',
},
],
},
{
role: MessageRole.Assistant,
content: 'answer',
},
{
role: MessageRole.User,
content: [
{
type: 'image',
source: {
data: 'aaaaaa',
mimeType: 'image/png',
},
},
{
type: 'image',
source: {
data: 'bbbbbb',
mimeType: 'image/png',
},
},
],
},
],
});

expect(executorMock.invoke).toHaveBeenCalledTimes(1);

const { messages } = getCallParams();
expect(messages).toEqual([
{
parts: [
{
text: 'question',
},
],
role: 'user',
},
{
parts: [
{
text: 'answer',
},
],
role: 'assistant',
},
{
parts: [
{
inlineData: {
data: 'aaaaaa',
mimeType: 'image/png',
},
},
{
inlineData: {
data: 'bbbbbb',
mimeType: 'image/png',
},
},
],
role: 'user',
},
]);
});

it('groups messages from the same user', () => {
geminiAdapter.chatComplete({
logger,
Expand Down
20 changes: 15 additions & 5 deletions .../platform/plugins/shared/inference/server/chat_complete/adapters/gemini/gemini_adapter.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,11 +196,21 @@ function messageToGeminiMapper() {
case MessageRole.User:
const userMessage: GeminiMessage = {
role: 'user',
parts: [
{
text: message.content,
},
],
parts: (typeof message.content === 'string' ? [message.content] : message.content).map(
(contentPart) => {
if (typeof contentPart === 'string') {
return { text: contentPart } satisfies Gemini.TextPart;
} else if (contentPart.type === 'text') {
return { text: contentPart.text } satisfies Gemini.TextPart;
}
return {
inlineData: {
data: contentPart.source.data,
mimeType: contentPart.source.mimeType,
},
} satisfies Gemini.InlineDataPart;
}
),
Comment on lines +199 to +213
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same NIT as for claude, but that's probably a personal preference at this point.

};
return userMessage;

Expand Down
Loading