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

Feature/AWSBedrock embeddings #1034

Merged
merged 1 commit into from
Oct 11, 2023
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
23 changes: 6 additions & 17 deletions packages/components/nodes/chatmodels/AWSBedrock/AWSChatBedrock.ts
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ class AWSChatBedrock_ChatModels implements INode {
this.type = 'AWSChatBedrock'
this.icon = 'awsBedrock.png'
this.category = 'Chat Models'
this.description = 'Wrapper around AWS Bedrock large language models'
this.description = 'Wrapper around AWS Bedrock large language models that use the Chat endpoint'
this.baseClasses = [this.type, ...getBaseClasses(ChatBedrock)]
this.credential = {
label: 'AWS Credential',
Expand Down Expand Up @@ -88,27 +88,18 @@ class AWSChatBedrock_ChatModels implements INode {
{ label: 'us-west-1', name: 'us-west-1' },
{ label: 'us-west-2', name: 'us-west-2' }
],
default: 'us-east-1',
optional: false
default: 'us-east-1'
},
{
label: 'Model Name',
name: 'model',
type: 'options',
options: [
{ label: 'amazon.titan-tg1-large', name: 'amazon.titan-tg1-large' },
{ label: 'amazon.titan-e1t-medium', name: 'amazon.titan-e1t-medium' },
{ label: 'stability.stable-diffusion-xl', name: 'stability.stable-diffusion-xl' },
{ label: 'ai21.j2-grande-instruct', name: 'ai21.j2-grande-instruct' },
{ label: 'ai21.j2-jumbo-instruct', name: 'ai21.j2-jumbo-instruct' },
{ label: 'ai21.j2-mid', name: 'ai21.j2-mid' },
{ label: 'ai21.j2-ultra', name: 'ai21.j2-ultra' },
{ label: 'anthropic.claude-instant-v1', name: 'anthropic.claude-instant-v1' },
{ label: 'anthropic.claude-v1', name: 'anthropic.claude-v1' },
{ label: 'anthropic.claude-v2', name: 'anthropic.claude-v2' }
],
default: 'anthropic.claude-v2',
optional: false
default: 'anthropic.claude-v2'
},
{
label: 'Temperature',
Expand All @@ -117,18 +108,16 @@ class AWSChatBedrock_ChatModels implements INode {
step: 0.1,
description: 'Temperature parameter may not apply to certain model. Please check available model parameters',
optional: true,
default: 0.7,
additionalParams: false
default: 0.7
},
{
label: 'Max Tokens to Sample',
name: 'max_tokens_to_sample',
type: 'number',
step: 10,
description: 'Max Tokens parameter may not apply to certain model. Please check available model parameters',
optional: false,
default: 200,
additionalParams: false
optional: true,
default: 200
}
]
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,152 @@
import { ICommonObject, INode, INodeData, INodeParams } from '../../../src/Interface'
import { getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils'
import { BedrockEmbeddings, BedrockEmbeddingsParams } from 'langchain/embeddings/bedrock'
import { BedrockRuntimeClient, InvokeModelCommand } from '@aws-sdk/client-bedrock-runtime'

class AWSBedrockEmbedding_Embeddings implements INode {
label: string
name: string
version: number
type: string
icon: string
category: string
description: string
baseClasses: string[]
credential: INodeParams
inputs: INodeParams[]

constructor() {
this.label = 'AWS Bedrock Embeddings'
this.name = 'AWSBedrockEmbeddings'
this.version = 1.0
this.type = 'AWSBedrockEmbeddings'
this.icon = 'awsBedrock.png'
this.category = 'Embeddings'
this.description = 'AWSBedrock embedding models to generate embeddings for a given text'
this.baseClasses = [this.type, ...getBaseClasses(BedrockEmbeddings)]
this.credential = {
label: 'AWS Credential',
name: 'credential',
type: 'credential',
credentialNames: ['awsApi'],
optional: true
}
this.inputs = [
{
label: 'Region',
name: 'region',
type: 'options',
options: [
{ label: 'af-south-1', name: 'af-south-1' },
{ label: 'ap-east-1', name: 'ap-east-1' },
{ label: 'ap-northeast-1', name: 'ap-northeast-1' },
{ label: 'ap-northeast-2', name: 'ap-northeast-2' },
{ label: 'ap-northeast-3', name: 'ap-northeast-3' },
{ label: 'ap-south-1', name: 'ap-south-1' },
{ label: 'ap-south-2', name: 'ap-south-2' },
{ label: 'ap-southeast-1', name: 'ap-southeast-1' },
{ label: 'ap-southeast-2', name: 'ap-southeast-2' },
{ label: 'ap-southeast-3', name: 'ap-southeast-3' },
{ label: 'ap-southeast-4', name: 'ap-southeast-4' },
{ label: 'ap-southeast-5', name: 'ap-southeast-5' },
{ label: 'ap-southeast-6', name: 'ap-southeast-6' },
{ label: 'ca-central-1', name: 'ca-central-1' },
{ label: 'ca-west-1', name: 'ca-west-1' },
{ label: 'cn-north-1', name: 'cn-north-1' },
{ label: 'cn-northwest-1', name: 'cn-northwest-1' },
{ label: 'eu-central-1', name: 'eu-central-1' },
{ label: 'eu-central-2', name: 'eu-central-2' },
{ label: 'eu-north-1', name: 'eu-north-1' },
{ label: 'eu-south-1', name: 'eu-south-1' },
{ label: 'eu-south-2', name: 'eu-south-2' },
{ label: 'eu-west-1', name: 'eu-west-1' },
{ label: 'eu-west-2', name: 'eu-west-2' },
{ label: 'eu-west-3', name: 'eu-west-3' },
{ label: 'il-central-1', name: 'il-central-1' },
{ label: 'me-central-1', name: 'me-central-1' },
{ label: 'me-south-1', name: 'me-south-1' },
{ label: 'sa-east-1', name: 'sa-east-1' },
{ label: 'us-east-1', name: 'us-east-1' },
{ label: 'us-east-2', name: 'us-east-2' },
{ label: 'us-gov-east-1', name: 'us-gov-east-1' },
{ label: 'us-gov-west-1', name: 'us-gov-west-1' },
{ label: 'us-west-1', name: 'us-west-1' },
{ label: 'us-west-2', name: 'us-west-2' }
],
default: 'us-east-1'
},
{
label: 'Model Name',
name: 'model',
type: 'options',
options: [{ label: 'amazon.titan-embed-text-v1', name: 'amazon.titan-embed-text-v1' }],
default: 'amazon.titan-embed-text-v1'
}
]
}

async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
const iRegion = nodeData.inputs?.region as string
const iModel = nodeData.inputs?.model as string

const obj: BedrockEmbeddingsParams = {
model: iModel,
region: iRegion
}

const credentialData = await getCredentialData(nodeData.credential ?? '', options)
if (credentialData && Object.keys(credentialData).length !== 0) {
const credentialApiKey = getCredentialParam('awsKey', credentialData, nodeData)
const credentialApiSecret = getCredentialParam('awsSecret', credentialData, nodeData)
const credentialApiSession = getCredentialParam('awsSession', credentialData, nodeData)

obj.credentials = {
accessKeyId: credentialApiKey,
secretAccessKey: credentialApiSecret,
sessionToken: credentialApiSession
}
}

const client = new BedrockRuntimeClient({
region: obj.region,
credentials: obj.credentials
})

const model = new BedrockEmbeddings(obj)

// Avoid Illegal Invocation
model.embedQuery = async (document: string): Promise<number[]> => {
return await embedText(document, client, iModel)
}

model.embedDocuments = async (documents: string[]): Promise<number[][]> => {
return Promise.all(documents.map((document) => embedText(document, client, iModel)))
}
return model
}
}

const embedText = async (text: string, client: BedrockRuntimeClient, model: string): Promise<number[]> => {
// replace newlines, which can negatively affect performance.
const cleanedText = text.replace(/\n/g, ' ')

const res = await client.send(
new InvokeModelCommand({
modelId: model,
body: JSON.stringify({
inputText: cleanedText
}),
contentType: 'application/json',
accept: 'application/json'
})
)

try {
const body = new TextDecoder().decode(res.body)
return JSON.parse(body).embedding
} catch (e) {
throw new Error('An invalid response was returned by Bedrock.')
}
}

module.exports = { nodeClass: AWSBedrockEmbedding_Embeddings }
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
22 changes: 7 additions & 15 deletions packages/components/nodes/llms/AWSBedrock/AWSBedrock.ts
Original file line number Diff line number Diff line change
Expand Up @@ -88,8 +88,7 @@ class AWSBedrock_LLMs implements INode {
{ label: 'us-west-1', name: 'us-west-1' },
{ label: 'us-west-2', name: 'us-west-2' }
],
default: 'us-east-1',
optional: false
default: 'us-east-1'
},
{
label: 'Model Name',
Expand All @@ -98,17 +97,12 @@ class AWSBedrock_LLMs implements INode {
options: [
{ label: 'amazon.titan-tg1-large', name: 'amazon.titan-tg1-large' },
{ label: 'amazon.titan-e1t-medium', name: 'amazon.titan-e1t-medium' },
{ label: 'stability.stable-diffusion-xl', name: 'stability.stable-diffusion-xl' },
{ label: 'cohere.command-text-v14', name: 'cohere.command-text-v14' },
{ label: 'ai21.j2-grande-instruct', name: 'ai21.j2-grande-instruct' },
{ label: 'ai21.j2-jumbo-instruct', name: 'ai21.j2-jumbo-instruct' },
{ label: 'ai21.j2-mid', name: 'ai21.j2-mid' },
{ label: 'ai21.j2-ultra', name: 'ai21.j2-ultra' },
{ label: 'anthropic.claude-instant-v1', name: 'anthropic.claude-instant-v1' },
{ label: 'anthropic.claude-v1', name: 'anthropic.claude-v1' },
{ label: 'anthropic.claude-v2', name: 'anthropic.claude-v2' }
],
default: 'anthropic.claude-v2',
optional: false
{ label: 'ai21.j2-ultra', name: 'ai21.j2-ultra' }
]
},
{
label: 'Temperature',
Expand All @@ -117,18 +111,16 @@ class AWSBedrock_LLMs implements INode {
step: 0.1,
description: 'Temperature parameter may not apply to certain model. Please check available model parameters',
optional: true,
default: 0.7,
additionalParams: false
default: 0.7
},
{
label: 'Max Tokens to Sample',
name: 'max_tokens_to_sample',
type: 'number',
step: 10,
description: 'Max Tokens parameter may not apply to certain model. Please check available model parameters',
optional: false,
default: 200,
additionalParams: false
optional: true,
default: 200
}
]
}
Expand Down
3 changes: 2 additions & 1 deletion packages/components/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
},
"license": "SEE LICENSE IN LICENSE.md",
"dependencies": {
"@aws-sdk/client-bedrock-runtime": "3.422.0",
"@aws-sdk/client-dynamodb": "^3.360.0",
"@dqbd/tiktoken": "^1.0.7",
"@getzep/zep-js": "^0.6.3",
Expand Down Expand Up @@ -45,7 +46,7 @@
"graphql": "^16.6.0",
"html-to-text": "^9.0.5",
"ioredis": "^5.3.2",
"langchain": "^0.0.157",
"langchain": "^0.0.165",
"langfuse-langchain": "^1.0.14-alpha.0",
"langsmith": "^0.0.32",
"linkifyjs": "^4.1.1",
Expand Down