-
-
Notifications
You must be signed in to change notification settings - Fork 16.6k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #1034 from FlowiseAI/feature/BedrockEmbeddings
Feature/AWSBedrock embeddings
- Loading branch information
Showing
5 changed files
with
167 additions
and
33 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
152 changes: 152 additions & 0 deletions
152
packages/components/nodes/embeddings/AWSBedrockEmbedding/AWSBedrockEmbedding.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 } |
Binary file added
BIN
+61.4 KB
packages/components/nodes/embeddings/AWSBedrockEmbedding/awsBedrock.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters