mirror of
https://github.com/farcasclaudiu/Flowise.git
synced 2026-06-28 17:01:00 +03:00
AWS Bedrock embeddings call retry extension (#2917)
* AWS Bedrock embeddings call retry extension * Removed unnecessary commented line --------- Co-authored-by: Rok Pajk Kosec <rok.pajkkosec@gmail.com>
This commit is contained in:
@@ -83,6 +83,24 @@ class AWSBedrockEmbedding_Embeddings implements INode {
|
|||||||
}
|
}
|
||||||
],
|
],
|
||||||
optional: true
|
optional: true
|
||||||
|
},
|
||||||
|
{
|
||||||
|
label: 'Batch Size',
|
||||||
|
name: 'batchSize',
|
||||||
|
description: 'Documents batch size to send to AWS API for Titan model embeddings. Used to avoid throttling.',
|
||||||
|
type: 'number',
|
||||||
|
optional: true,
|
||||||
|
default: 50,
|
||||||
|
additionalParams: true
|
||||||
|
},
|
||||||
|
{
|
||||||
|
label: 'Max AWS API retries',
|
||||||
|
name: 'maxRetries',
|
||||||
|
description: 'This will limit the nubmer of AWS API for Titan model embeddings call retries. Used to avoid throttling.',
|
||||||
|
type: 'number',
|
||||||
|
optional: true,
|
||||||
|
default: 5,
|
||||||
|
additionalParams: true
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
@@ -144,7 +162,9 @@ class AWSBedrockEmbedding_Embeddings implements INode {
|
|||||||
if (iModel.startsWith('cohere')) {
|
if (iModel.startsWith('cohere')) {
|
||||||
return await embedTextCohere(documents, client, iModel, inputType)
|
return await embedTextCohere(documents, client, iModel, inputType)
|
||||||
} else {
|
} else {
|
||||||
return Promise.all(documents.map((document) => embedTextTitan(document, client, iModel)))
|
const batchSize = nodeData.inputs?.batchSize as number
|
||||||
|
const maxRetries = nodeData.inputs?.maxRetries as number
|
||||||
|
return processInBatches(documents, batchSize, maxRetries, (document) => embedTextTitan(document, client, iModel))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return model
|
return model
|
||||||
@@ -195,4 +215,38 @@ const embedTextCohere = async (texts: string[], client: BedrockRuntimeClient, mo
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const processInBatches = async (
|
||||||
|
documents: string[],
|
||||||
|
batchSize: number,
|
||||||
|
maxRetries: number,
|
||||||
|
processFunc: (document: string) => Promise<number[]>
|
||||||
|
): Promise<number[][]> => {
|
||||||
|
let sleepTime = 0
|
||||||
|
let retryCounter = 0
|
||||||
|
let result: number[][] = []
|
||||||
|
for (let i = 0; i < documents.length; i += batchSize) {
|
||||||
|
let chunk = documents.slice(i, i + batchSize)
|
||||||
|
try {
|
||||||
|
let chunkResult = await Promise.all(chunk.map(processFunc))
|
||||||
|
result.push(...chunkResult)
|
||||||
|
retryCounter = 0
|
||||||
|
} catch (e) {
|
||||||
|
if (retryCounter < maxRetries && e.name.includes('ThrottlingException')) {
|
||||||
|
retryCounter = retryCounter + 1
|
||||||
|
i = i - batchSize
|
||||||
|
sleepTime = sleepTime + 100
|
||||||
|
} else {
|
||||||
|
// Split to distinguish between throttling retry error and other errors in trance
|
||||||
|
if (e.name.includes('ThrottlingException')) {
|
||||||
|
throw new Error('AWS Bedrock retry limit reached: ' + e)
|
||||||
|
} else {
|
||||||
|
throw new Error(e)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
await new Promise((resolve) => setTimeout(resolve, sleepTime))
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
module.exports = { nodeClass: AWSBedrockEmbedding_Embeddings }
|
module.exports = { nodeClass: AWSBedrockEmbedding_Embeddings }
|
||||||
|
|||||||
Reference in New Issue
Block a user