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.

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

feat (provider/bedrock): add file content part support #3383

Merged
merged 8 commits into from
Oct 28, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
5 changes: 5 additions & 0 deletions .changeset/curvy-needles-punch.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
'@ai-sdk/amazon-bedrock': patch
---

feat (provider/bedrock): add file content part support
33 changes: 31 additions & 2 deletions content/providers/01-ai-sdk-providers/08-amazon-bedrock.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -157,8 +157,6 @@ const model = bedrock('anthropic.claude-3-sonnet-20240229-v1:0', {

Documentation for additional settings based on the selected model can be found within the [Amazon Bedrock Inference Parameter Documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html).

### Example

You can use Amazon Bedrock language models to generate text with the `generateText` function:

```ts
Expand All @@ -174,6 +172,37 @@ const { text } = await generateText({
Amazon Bedrock language models can also be used in the `streamText` function
(see [AI SDK Core](/docs/ai-sdk-core)).

### File Inputs

<Note type="warning">
Amazon Bedrock supports file inputs on in combination with specific models,
e.g. `anthropic.claude-3-haiku-20240307-v1:0`.
</Note>

The Amazon Bedrock provider supports file inputs, e.g. PDF files.

```ts
import { bedrock } from '@ai-sdk/amazon-bedrock';
import { generateText } from 'ai';

const result = await generateText({
model: bedrock('anthropic.claude-3-haiku-20240307-v1:0'),
messages: [
{
role: 'user',
content: [
{ type: 'text', text: 'Describe the pdf in detail.' },
{
type: 'file',
data: fs.readFileSync('./data/ai.pdf'),
mimeType: 'application/pdf',
},
],
},
],
});
```

### Guardrails

You can use the `bedrock` provider metadata to utilize [Amazon Bedrock Guardrails](https://aws.amazon.com/bedrock/guardrails/):
Expand Down
29 changes: 29 additions & 0 deletions examples/ai-core/src/stream-text/amazon-bedrock-pdf.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
import { bedrock } from '@ai-sdk/amazon-bedrock';
import { streamText } from 'ai';
import 'dotenv/config';
import fs from 'node:fs';

async function main() {
const result = await streamText({
model: bedrock('anthropic.claude-3-haiku-20240307-v1:0'),
messages: [
{
role: 'user',
content: [
{ type: 'text', text: 'Describe the pdf in detail.' },
{
type: 'file',
data: fs.readFileSync('./data/ai.pdf'),
mimeType: 'application/pdf',
},
],
},
],
});

for await (const textPart of result.textStream) {
process.stdout.write(textPart);
}
}

main().catch(console.error);
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,9 @@ describe('system messages', () => {
});

describe('user messages', () => {
it('should convert messages with image and text parts to multiple parts', async () => {
it('should convert messages with file, image, and text parts to multiple parts', async () => {
const fileData = new Uint8Array([0, 1, 2, 3]);

const { messages } = convertToBedrockChatMessages([
{
role: 'user',
Expand All @@ -32,6 +34,11 @@ describe('user messages', () => {
image: new Uint8Array([0, 1, 2, 3]),
mimeType: 'image/png',
},
{
type: 'file',
data: Buffer.from(fileData).toString('base64'),
mimeType: 'application/pdf',
},
],
},
]);
Expand All @@ -47,6 +54,13 @@ describe('user messages', () => {
source: { bytes: new Uint8Array([0, 1, 2, 3]) },
},
},
{
document: {
format: 'pdf',
name: expect.any(String),
source: { bytes: Buffer.from(fileData) },
},
},
],
},
]);
Expand Down
27 changes: 26 additions & 1 deletion packages/amazon-bedrock/src/convert-to-bedrock-chat-messages.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,16 @@ import {
LanguageModelV1Prompt,
UnsupportedFunctionalityError,
} from '@ai-sdk/provider';
import { ImageFormat } from '@aws-sdk/client-bedrock-runtime';
import { createIdGenerator } from '@ai-sdk/provider-utils';
import { DocumentFormat, ImageFormat } from '@aws-sdk/client-bedrock-runtime';
import {
BedrockAssistantMessage,
BedrockMessagesPrompt,
BedrockUserMessage,
} from './bedrock-chat-prompt';

const generateFileId = createIdGenerator({ prefix: 'file', size: 16 });

export function convertToBedrockChatMessages(
prompt: LanguageModelV1Prompt,
): BedrockMessagesPrompt {
Expand Down Expand Up @@ -70,6 +73,28 @@ export function convertToBedrockChatMessages(
},
});

break;
}
case 'file': {
if (part.data instanceof URL) {
// The AI SDK automatically downloads files for user file parts with URLs
throw new UnsupportedFunctionalityError({
functionality: 'File URLs in user messages',
});
}

bedrockContent.push({
document: {
format: part.mimeType?.split(
'/',
)?.[1] as DocumentFormat,
name: generateFileId(),
source: {
bytes: Buffer.from(part.data, 'base64'),
},
},
});

break;
}
}
Expand Down
Loading